Learning spatio-temporal patterns with Neural Cellular Automata

TL;DR

This work presents Neural Cellular Automata as a minimal, differentiable framework that learns local update rules to generate and evolve complex spatio-temporal patterns. By tying NCAs to PDE discretisations and validating on Gray-Scott reaction-diffusion dynamics and image morphing tasks, the authors demonstrate robust generalisation to unseen initial conditions, the ability to capture transient and stable structures, and the impact of symmetry constraints and hyperparameters on stability. The approach offers a data-driven yet interpretable pathway to model biological pattern formation and other emergent phenomena, with potential extensions to multiscale and time-dependent rules. Overall, the study broadens the scope of NCA as a practical, modular tool for learning dynamic patterns from imaging data and PDE trajectories.

Abstract

Neural Cellular Automata (NCA) are a powerful combination of machine learning and mechanistic modelling. We train NCA to learn complex dynamics from time series of images and PDE trajectories. Our method is designed to identify underlying local rules that govern large scale dynamic emergent behaviours. Previous work on NCA focuses on learning rules that give stationary emergent structures. We extend NCA to capture both transient and stable structures within the same system, as well as learning rules that capture the dynamics of Turing pattern formation in nonlinear Partial Differential Equations (PDEs). We demonstrate that NCA can generalise very well beyond their PDE training data, we show how to constrain NCA to respect given symmetries, and we explore the effects of associated hyperparameters on model performance and stability. Being able to learn arbitrary dynamics gives NCA great potential as a data driven modelling framework, especially for modelling biological pattern formation.

Figures (12)

• Figure 1: Schematic of an update step of the NCA. For each $C$ channel pixel in the ${\color{black} S}\times {\color{black} S}$ lattice $x^{(n)}$ at step $n$, a perception vector $z^{(n)}$ is constructed to encode local information via convolution with hard-coded kernels $K$. This perception vector is fed through a dense neural network $F_\theta$ with trainable weights $W_1$, $W_2$, and biases $v$. The nonlinear activation function $u(\cdot)$ is applied on the single hidden layer of the network. The output of this network yields the incremental update to that pixel, which is applied in parallel to all pixels with the stochastic mask $\sigma$ to determine the lattice state $x^{(n+1)}$ at step $n+1$.
• Figure 2: 1D phase space representation of NCA trajectories, predictions ${\color{black} \hat{x}}^{({\color{black} m} t,r)}$ and true states $y^{({\color{black} m} ,r)}$. Here ${\color{black} M}=3,R=2$. The first batch ($x^{(\cdot,1)}$) is trained with re-initialised intermediate states, whereas the second batch ($x^{(\cdot,2)}$) is trained with propagated intermediate states.
• Figure 3: Snapshots taken from the training data used for learning PDE dynamics. PDE is run for $N=1024$ steps with timestep 1 and $D_A=0.1,D_B=0.05,\alpha = 0.06230,\gamma = 0.06268$.
• Figure 4: A: loss as a function of time sampling $t$. Training loss shows the minimum loss during training epochs, averaged over 4 random initialisations, with standard deviation as error bars. Test loss shows how the best trained NCA (minimal training loss) performs on unseen initial conditions. B: Initial condition and true state (PDE simulation) at $n=2048$. C: Snapshots of NCA trajectories (at $n=2048$) based on unseen initial conditions, with varying time sampling $t$. Each NCA is trained for 4000 epochs, with a mini-batch size $B=64$.
• Figure 5: Snapshots of PDE and NCA trajectories from an unseen initial condition. NCA trained with $C=8$, Identity and Laplacian kernels, relu activation, trained on sampling $t=32$ for 4000 epochs with Euclidean loss.