Source-linked AI summary
Learning PDE Time-Stepping with Neural Cellular Automata
Esha Saha, Hao Wang
TL;DR
Repeated PDE simulation under varying conditions can make classical solvers expensive. This paper learns a local NCA update rule for repeated PDE time-stepping and benchmarks it against PDE-Net, PINN, and FNO. At long horizons, NCA outperformed the other baselines on a majority of benchmark PDEs, though it was not uniformly superior.
Problem
Classical PDE solvers are computationally expensive for repeated simulations across varying initial conditions, motivating learned surrogates.
Method
NCA learns a local, homogeneous update rule applied repeatedly at each grid cell to produce long PDE rollouts.
Results
At long time horizons, NCA outperformed PDE-Net, FNO, and PINN on a majority of the benchmark PDEs.
Takeaways & Limitations
NCA provides a simple local model for predicting PDE solutions over many time steps while maintaining good accuracy and stability in the reported benchmarks.
Takeaways & Limitations
NCA’s performance depends on local perception and design choices, and standard NCA does not explicitly include physical constraints such as conservation laws or boundary conditions.
Abstract
from arXiv · showhide
Classical numerical solvers for partial differential equations (PDEs) are computationally expensive to solve repeatedly across varying initial conditions, motivating the need for learned surrogates. In this paper, we propose a trainable Neural Cellular Automata (NCA) based surrogate model for learning long time PDE dynamics. Rather than mapping an entire initial field to a full trajectory in one shot, our proposed model learns a small, local, homogeneous update rule that is applied identically and repeatedly at every grid cell, mirroring the locality of differential operators. We benchmark this framework against three baselines: PDE - Net, a modified physics-informed neural network (PINN), and a Fourier Neural Operator (FNO), on five canonical PDEs (heat, advection, Burgers, Allen - Cahn, and Fisher - KPP), evaluated at temporal domain two times beyond the training temporal domain. The proposed model achieves the lowest long-horizon relative errors on the majority of the experiments.
1 Introduction
Classical PDE solvers are accurate and interpretable but costly for repeated simulations, motivating learned surrogates. The paper explores NCA as a local, homogeneous, repeatedly applied PDE time-stepping model and evaluates it against learned-solver baselines.
- Classical finite-difference, finite-element, and spectral solvers can be computationally expensive when equations must be solved repeatedly under varying conditions.
- NCA learns a small local update rule applied identically at every grid cell using only its immediate neighborhood.This construction mirrors the locality of differential operators and makes the model translation-invariant.
- Repeated application of the learned rule produces a full spatiotemporal rollout trained against trajectories generated by a classical numerical solver.
- The framework is tested on five canonical PDEs and benchmarked against PDE-Net, PINN, and FNO up to twice the training temporal domain.
2 Related Work
Related work spans classical and physics-informed solvers, operator learning, graph-based local simulators, PDE-Net, and general Neural Cellular Automata. This paper recasts local NCA updates as learned PDE time-stepping under shared finite-difference evaluation.
- Classical discretization methods remain standard for PDE simulation, while PINNs parameterize solution fields and penalize governing-equation residuals.PINNs may require retraining for new initial or boundary conditions and can struggle with stiff or multi-scale dynamics.
- Operator learning maps between function spaces, with FNO operating in the spectral domain and DeepONet using a branch-trunk architecture.
- Graph-network simulators learn local message-passing updates on mesh or particle graphs, while PDE-Net learns trainable convolutional filters with a pointwise nonlinearity.
- The paper recasts NCA’s local-update construction as a learned numerical PDE time-stepper with gated finite-difference perception and synchronous residual updates.Its evaluation compares NCA with PDE-Net, PINN, and FNO using a shared finite-difference solver across five canonical PDEs.
3 Methodology
The method represents PDE evolution with a multi-channel NCA state whose local, gated stencil perception and shared neural update are repeatedly applied in closed loop. It trains these synchronous residual steps against finite-difference trajectories and evaluates long-horizon rollouts on canonical PDEs.
- State representation: The NCA state stores the physical PDE variable in its first channel and latent states in the remaining C −1 channels for additional information.The latent channels can retain implicit derivatives or temporal context.
- NCA update: Each rollout step computes a local perception of the full state, passes it through a learned update network, and applies the resulting C-channel update synchronously.Every cell updates at every step without a stochastic firing mask.
- Perception: Perception uses four fixed 3×3 filters—identity, Sobel-x, Sobel-y, and discrete Laplacian—with trainable scalar gates controlling their contributions.The fixed filters are applied independently to each channel, and an ℓ1 penalty with λ=10^-4 encourages sparse stencil preferences.
- Neural update: A pointwise MLP implemented by 1×1 convolutions maps gated perception features to the state update using the same rule at every spatial location.The output layer is zero-initialized so training begins as an identity map.
- Training and evaluation: Training advances the NCA and finite-difference solver in lockstep for R=500 steps, uses the model’s own state without teacher forcing, and truncates BPTT at K=1.Evaluation removes detachment and rolls out to T=1000; errors are reported as mean ± std relative ℓ2 errors over 200 held-out initial conditions at horizons 100, 200, 600, 800, and 1000.
- Experimental setup: The benchmark covers five canonical PDEs on a 64×64 grid with periodic boundaries, using smooth random initial fields and finite-difference trajectories.The experimental setup includes the main finite-difference solver and benchmark training settings.
4 Results
NCA generally delivers the strongest long-horizon accuracy across the benchmark, though PDE-Net is better on Allen–Cahn and FNO narrowly leads on advection. Its advantage is most evident under repeated rollouts, while noise robustness varies by equation.
- Long-horizon benchmark: At T = 1000 on heat, NCA reaches relative ℓ2 error 0.561 versus 0.709 for PDE - Net, 0.824 for FNO, and 1.439 for PINN.This is approximately a 21% reduction relative to the next-best baseline, PDE - Net.
- Long-horizon benchmark: On Burgers, NCA error rises from 0.040 at T = 100 to 0.196 at T = 1000, while PDE - Net rises from 0.093 to 1.754 and FNO from 0.301 to 0.853.NCA remains below 0.20 throughout the evaluated interval.
- Long-horizon benchmark: NCA has the lowest Fisher - KPP error at every reported horizon, reaching 0.315 at T = 1000 versus 0.674 for PDE - Net, 0.820 for FNO, and 1.817 for PINN.The error starts at 0.012 at T = 100.
- Equation-specific results: PDE - Net is consistently more accurate than NCA on Allen - Cahn, with T = 1000 errors of 0.015 and 0.052, respectively.NCA nevertheless remains below PINN (0.126) and FNO (0.138) at T = 1000, showing that performance depends on the PDE configuration.
- Equation-specific results: FNO is strongest on advection at the longest horizon, with relative error 0.437 compared with 0.457 for NCA.Across heat, Burgers, and Fisher - KPP, NCA improves over FNO at T = 1000 by approximately 32%, 77%, and 62%, respectively.
- Long-horizon benchmark: NCA provides the strongest long-horizon performance on the majority of the benchmark PDEs.Relative error generally increases with temporal horizon, but NCA shows the most stable behavior on most PDEs.
- Robustness to noise: Under additive Gaussian training noise, NCA retains relatively low heat and Fisher - KPP errors but becomes substantially less stable on advection.For heat, the T=1000 error decreases from 0.561 to 0.246 at 5% noise; for Fisher - KPP, it decreases from 0.315 to 0.123.
- Learned operators: NCA’s perception gates favor the Laplacian on heat, Allen - Cahn, and Fisher - KPP, while Burgers also receives a relatively larger ∂x gate.This matches the diffusion-dominated structure of the first three equations and Burgers’ combination of diffusion with advective flux.
5 Discussion
NCA is trained as a local one-step update and evaluated through repeated composition, with long-horizon performance compared against PDE-Net, PINN, and FNO. Across the discussion and ablations, NCA is generally stronger at long rollouts, although results depend on the PDE, resolution, kernel, and training setup.
- Method: NCA learns a one-step map from immediate neighboring cells and evaluates it by composing the update over long rollouts.Training uses closed-loop states, while evaluation measures stability over T = 1000 steps.
- Baseline comparison: On heat, PDE-Net errors are 0.067 versus NCA’s 0.011 at T = 100, and 0.709 versus 0.561 at T = 1000.PDE-Net outperforms NCA on Allen-Cahn, while its Burgers error rises from 0.174 to 1.754 between T = 200 and T = 1000.
- Baseline comparison: PINN performance lies between PDE-Net and PINN on heat, Burgers, and Fisher-KPP, but is substantially behind NCA at T = 1000.The PINN maps an initial condition directly to a queried horizon, unlike the composed Markov updates used by NCA, PDE-Net, and FNO.
- Ablation and efficiency: At matched settings, NCA improves with resolution, while FNO remains less accurate and can require nearly 10× longer training than NCA.On heat at N = 64, training takes 23000 seconds for FNO versus 2300 seconds for NCA; PINN trains in about 380 seconds but performs poorly at longer horizons.
- Ablation and efficiency: At N = 64 and T = 1000 on heat, widening the perception kernel raises relative error from 0.179 for 3 × 3 to 0.825 for 5 × 5 and 2.231 for 7 × 7.The same degradation pattern holds on Burgers, favoring immediate-neighbor perception.
6 Conclusions
The paper concludes that a local NCA update can maintain accurate, stable PDE rollouts and outperform other learned solvers on most benchmark equations. It also identifies limits from local perception, repeated updates, design choices, physical constraints, and data-driven generalization.
- Conclusion: NCA outperformed PDE-Net, FNO, and PINN at long horizons for a majority of the benchmark PDEs.The conclusion links this behavior to the stability of the learned update under repeated application.
- Conclusion: Ablations showed better NCA performance on heat and Burgers across resolutions, with optimal results when learning from immediate neighbors.The conclusion emphasizes locality as a design condition rather than simply favoring larger perception fields.
- Limitations and future work: NCA’s local interactions limit each cell’s available information, while repeated application can accumulate errors when the learned update is insufficiently stable.Performance also depends on the perception kernel, update rule, and training rollout length.
- Limitations and future work: Standard NCA does not explicitly enforce conservation laws, boundary conditions, or known symmetries unless these are added to the model or training procedure.Accuracy on the considered PDEs therefore does not necessarily imply satisfaction of the underlying physical laws.
- Limitations and future work: Future work includes incorporating physical properties and developing adaptive or multi-scale perception mechanisms for local and longer-range interactions.The proposed directions avoid assuming that simply increasing kernel size improves performance.