Learning PDE Time-Stepping with Neural Cellular Automata

2026-08-31Machine Learning

Machine Learning
AI summary

The authors developed a new type of neural network based on Cellular Automata to predict the behavior of systems described by partial differential equations (PDEs) over long times. Instead of predicting everything at once, their model learns a simple rule that updates each small part of the system repeatedly, similar to how PDEs work locally. They tested their method on five common PDEs and found it generally more accurate for making predictions beyond the time ranges it was trained on compared to other neural network approaches. This suggests their model can effectively mimic complex PDE dynamics over extended periods.

Partial Differential EquationsNeural Cellular AutomataSurrogate ModelPhysics-Informed Neural NetworksFourier Neural OperatorHeat EquationAdvection EquationBurgers' EquationAllen-Cahn EquationFisher-KPP Equation
Authors
Esha Saha, Hao Wang
Abstract
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.