# Learning spatio-temporal patterns with Neural Cellular Automata

Alex D. Richardson111School of Physics and Astronomy, University of Edinburgh,
Edinburgh EH9 3FD, UK 222To whom correspondence should be addressed. E-mail:
<EMAIL_ADDRESS>Tibor Antal 333School of Mathematics and
Maxwell Institute for Mathematical Sciences, University of Edinburgh,
Edinburgh EH9 3FD, UK Richard A. Blythe ††footnotemark:  Linus J. Schumacher
444Institute of regeneration and repair, University of Edinburgh, Edinburgh
EH16 4UU, UK

###### 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.

###### Contents

  1. 1 Introduction
  2. 2 Model and methods
    1. 2.1 Model details and parameters
    2. 2.2 Training techniques
      1. 2.2.1 Loss functions
  3. 3 Results
    1. 3.1 Gray-Scott reaction diffusion equations
      1. 3.1.1 Effect of training hyperparameters
    2. 3.2 Image morphing
      1. 3.2.1 Effect of model hyperparameters
      2. 3.2.2 Stability analysis
  4. 4 Discussion
  5. A Gradient Calculation
  6. B Videos

##  1 Introduction

Many complex natural phenomena—such as organ growth, the structure of
materials or the patterns of neural activity in our brains—are emergent [1].
These are typically characterised by many simple interacting components that
collectively exhibit behaviour that is far richer than that of the individual
parts, and cannot easily be predicted from them. Emergence is especially
prevalent in complex systems of biological nature across a wide range of
scales - from gene expression dictating cell fates, interacting cells forming
structures during morphogenesis, synaptic connections in the brain, or the
interactions of organisms in ecology.

Cellular Automata (CA) provide simple models of spatio-temporal emergent
behaviour, where a discrete lattice of ‘cells’ are equipped with an internal
state and a rule that updates each cell state depending on itself and its
local neighbours. The classic _Game of Life_ [2] is a famous example, where
cell states and the update rule utilise simple Boolean logic, but the emergent
complexity has fascinated and inspired much research [3, 4]. CA are a natural
modelling framework of a wide range of biological processes such as: skin
patterning [5, 6], limb polydactyly [7], chimerism [8], cancer [9] and
landscape ecology [10]. In these cases the CA rules are constructed with
expert knowledge of likely mechanisms, however in general the space of
possible CA rules is vast, and there is a non-uniqueness by which several
rules can result in qualitatively similar emergent behaviours. As such the
inverse problem of inferring mechanistic interactions (CA rules) that might
generate a given observed emergent behaviour is much more challenging than the
forward problem. Establishing the emergent consequences of a known set of
mechanistic interactions between components is conceptually straightforward —
one sets them up in a computational model or in vivo and then observes the
collective behaviour that emerges. In the case of CA, once a rule is defined,
any initial condition can trivially be propagated forward to obtain the
emergent behaviour.

In this work, we establish and extend the utility of Neural Cellular Automata
(NCA, [11])—a special case of CA where each cell state is a real vector, and
the update rule is determined by a neural network. The update rule is
parameterised by the neural network by minimising a cost function that
measures how similar the trajectory generated by iteratively applying the
update rule is to training data. Since the gradient of the loss function can
be evaluated straightforwardly, gradient-based optimisation allows for
efficiently learning (non-unique) local update rules to match desired global
behaviour, in the usual manner of machine learning [12]. This allows us to
tackle the inverse problem of inferring local mechanistic rules given observed
emergent behaviour.

We investigate the potential for NCA to be applied as a data-driven
alternative, where the underlying mechanisms are not assumed to be known. This
could include biological systems, which are inherently complex and may feature
interactions that are not directly measured by a given experimental procedure.
We do this by exploring the behaviour of NCAs on two idealised systems. First,
we train the NCA on Turing patterns generated by the solutions of certain
partial differential equations (PDEs). We show that the underlying dynamics
are well represented, for example, by testing on initial conditions that were
not part of the training set. Second, to understand the generality of the
method, we build on previously observed behaviour of NCA on the artificial
problem of morphing from one image to another [11]. Here any underlying
dynamics is _a priori_ unknown and likely to be highly complex, as opposed to
the PDE learning case where we know and understand the PDE. Nonetheless we
show that such dynamics are learnable and are robust to perturbations. In
addition, we achieve this with the minimal neural network complexity of NCA
[13], in line with the principles of Explainable Artificial Intelligence [14,
15].

In Section 2 we set out the structure of the NCA, and show that it can be
viewed as a machine learning approach to approximating the finite difference
discretisation of an unknown PDE. There is already a fairly strong link
between cellular automata and reaction diffusion equations [6, 5], in that
both are used to model similar systems, and a suitably chosen CA rule will
correspond to the finite difference approximation of any PDE. In Section 3 we
present our results. We first (Section 3.1) benchmark the NCA by assessing its
ability to capture the types of Turing patterns [16] that emerge from the
Gray-Scott [17] reaction-diffusion equations. These equations describe the
population of two chemical species with a nonlinear interaction between them,
and is capable of generating a variety of patterns. In Section 3.2, we show
that the same basic model and training techniques can also be applied to an
image morphing task. Thus we conclude that NCA are capable of constructing
microscopic dynamical rules that lead to a wide range of prescribed emergent
behaviour. In Section 3.2.2 we further explore constraining NCA to respect
basic symmetry requirements placed upon them, and investigate the robustness
of trained NCA to initial condition perturbations. We discuss implications for
further development of the method Section 4.

##  2 Model and methods

We now define the NCA model, and discuss the motivations behind design choices
and hyperparameters. We further discuss the training methods, as it turns out
that most training parameters can be kept constant between tasks. The main
exception to this is how frequently to sample in time: PDEs have a clear
notion of time, whereas image morphing does not. All the models and software
developed here have been implemented within the Tensorflow [18] framework
https://github.com/AlexDR1998/NCA which permits efficient GPU parallelisation.

###  2.1 Model details and parameters

Neural Cellular Automata (NCA) are a class of cellular automata defined on a
lattice of real vectors, with the local update rule encoded in a neural
network. As with all neural network models, there is freedom to choose the
network structure and how the input data are preprocessed for training and
testing. We refer to the set of choices that are not learned via the training
data as _hyperparameters_. Figure 1 shows a schematic of a single NCA update,
which maps the state of the system at time step $n$ to a corresponding state
at time $n+1$, where $n=0,1,\dots,$. This update comprises a sequence of
stages—depicted counterclockwise starting top left in the figure—which we now
describe in turn.

##### System state

We take the state of the system to be described through a vector of $C$ real
numbers at each point on an $M\times M$ lattice. For example, each of the $C$
values could represent the concentration of a chemical or biological species
at a given point in space. These observable channels are shown with coloured
shading in Figure 1. These can be augmented with _hidden channels_
(transparent in the figure), the state of which can influence the dynamics
within the _observable channels_. In a biological context, these hidden
channels could represent concentration profiles of proteins or chemicals that
are not measured in a particular experimental setting, but can be inferred by
the machine learning algorithm. The number of hidden channels is a
hyperparameter of the model. The number of hidden and observable channels sum
to $C$. Mathematically, we denote the state of the system at timestep $n$ as
the vector $x^{(n)}\in\mathcal{X}=\mathds{I}^{M\times M\times C}$, where
$\mathds{I}\in[a,b]$ is some interval of real numbers. The elements of this
vector are $x^{(n)}_{ijc}\in\mathds{I}$, which $i\in[1,M]$, $j\in[1,M]$ and
$c\in[1,C]$ denote the $x$-coordinate, $y$-coordinate and channel number,
respectively. We emphasise that during training (Section 2.2), only the
observable channels are compared to data.

Figure 1: Schematic of an update step of the NCA. For each $C$ channel pixel
in the $M\times M$ 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$.

##### Convolution kernels

The first stage of the update is to apply _convolution kernels_ to the spatial
data in each channel. These kernels are chosen to mimic differential
operators, such as gradients (Sobel filters) and Laplacians. We denote the set
of convolution kernels $g^{k}$, labelled with the index $k=1,\ldots,K$. Each
kernel $g^{k}$ is a square matrix, in our case $(3\times 3)$. This generates
the expanded _perception vector_ $z^{(n)}$ whose elements are given by:

$z^{(n)}_{ijck}=\sum_{\Delta i,\Delta j\in[-1,0,1]}g_{\Delta i,\Delta
j}^{k}x^{(n)}_{i+\Delta i,j+\Delta j,c}\equiv g\ast x^{(n)}$ (1)

Crucially the kernels are applied in parallel on all channels $C$
independently: that is, all kernels are applied _depthwise_. The idea of
decomposing an arbitrary convolution to separate depthwise and channel-mixing
convolutions [19] was inspired by the deep link between Convolutional Neural
Networks (CNNs) and Cellular Automata [11, 20]. In particular, this
facilitates representing the NCA in standard Tensorflow code. In principle,
kernels can be learnable rather than hard-coded; however this makes the
trained models less interpretable and so we do not pursue this approach here.
As the $3\times 3$ convolution kernels only encode information about the Moore
neighbourhood (i.e. adjacent and diagonal cells), we would never need any more
than 9 kernels, as any more would be linearly dependant on those already
present.

The purpose of applying the kernels is to make a clearer correspondence
between NCAs and numerical methods for solving PDEs. Essentially, they provide
the neural network with a basic set of differential operators to work with.
The set of kernels to include is an important hyperparameter. For example, if
one anticipates that the update rules should be invariant under a global
rotation of the system—i.e., that the dynamics are isotropic—one can justify
excluding Sobel kernels and just using identity, Laplacian, and local
averages. We already have translational invariance as a direct consequence of
the NCA construction, but isotropic symmetry can only be realistically
achieved by only using isotropic kernels [21] or data augmentation. We explore
this further in Section 3.2.2. The explicit forms of the convolution kernels
used in this work are

$\underbrace{\begin{bmatrix}0&0&0\\\ 0&1&0\\\
0&0&0\end{bmatrix}}_{\text{identity}},\quad\underbrace{\frac{1}{9}\begin{bmatrix}1&1&1\\\
1&1&1\\\
1&1&1\end{bmatrix}}_{\text{average}},\quad\underbrace{\frac{1}{8}\begin{bmatrix}1&2&1\\\
0&0&0\\\
-1&-2&-1\end{bmatrix}}_{\text{Sobel}_{x}},\quad\underbrace{\frac{1}{8}\begin{bmatrix}1&0&-1\\\
2&0&-2\\\
1&0&-1\end{bmatrix}}_{\text{Sobel}_{y}},\quad\underbrace{\frac{1}{4}\begin{bmatrix}1&2&1\\\
2&-12&2\\\ 1&2&1\end{bmatrix}}_{\text{Laplacian}}.$ (2)

##### Neural network

The perception vector $z^{(n)}$ is then applied to the input layer of a neural
network (see lower part of Figure 1). The values on the output layer,
$F(z^{(n)})$, form a vector of increments in the original state space
$\mathcal{X}$. In a deterministic update, one would simply add $F(z^{(n)})$ to
$x^{(n)}$ to obtain the updated state, $x^{(n+1)}$. Taking $F(z^{(n)})$ as an
increment, rather than the new state vector $x^{(n+1)}$, implies that the NCA
is a residual, rather than a naive, recurrent neural network (RCNN) [22].

In the present context, residual RCNNs have several benefits. Firstly a
residual RCNN is easier to compare to numerical discretisiations of PDEs,
aiding interpretability of our models. Secondly the residual RCNN minimises
the problem of vanishing or exploding gradients that the naive RCNN would
experience. In the naive approach, recurrent iterations of our neural network
would lead to exponentially large or small gradient updates, leading to a
failure to learn optimal model weights. A consequence of the vanishing
gradients problem is that information from previous timesteps is quickly lost,
so $x^{(m)}$ has little bearing on $x^{(n)}$ for $m\ll n$. In principle a
naive RCNN can learn long term dependencies, but in practice this is very
challenging. As such residual RCNNs are better suited to learning long term
behaviours, as $x^{(n+1)}$ depends linearly on $x^{(n)}$, so information
preservation is in some sense the default behaviour of our model. This
behaviour is especially clear during training, in that initialising the
residual RCNN to perform the identity mapping $x^{(n+1)}=x^{(n)}$ is
straightforward: one simply arranges for $F(z^{(n)})=0$ by setting the final
layer weights in the neural network to zero. Initialising the weights such
that $F(z^{(n)})=x^{(n)}$, which would be required in the naive case, is much
harder. This ‘do nothing’ update is a better starting point than one that
quickly forgets the initial state of the system. In the latter case, the NCA
may resort to learning how to construct the desired $x^{(n)}$ as a global
attractor, irrespective of initial conditions, for example ‘growing’ $x^{(n)}$
from fixed boundary conditions. Preserving the initial system state allows the
model to better learn dynamics particular to those initial conditions, whilst
still allowing for boundary driven behaviour to be learned.

It remains to specify the neural network structure, that is, the number and
size of any hidden layers, and how they are connected. Here we aim to keep the
architecture as simple as possible, as a minimal yet sufficiently functional
network architecture has several advantages. Training a small model is
computationally cheaper, and smaller models are far more interpretable [15].
Specifically, we use just one hidden layer, as shown in Figure 1. As noted
previously, we do not mix spatial data in the neural network, only between
channels and kernels. That is, the network shown in Figure 1 is replicated for
each pixel $i,j$ in $z^{(n)}$, and takes as input the $K\times C$ elements of
$z^{(n)}$ that correspond to a given pixel, and transforms to $C$ channel
values, consistent with the original state vector.

The hyperparameters associated with this neural network structure are the
choice of activation function and size of the hidden layer. We fix the hidden
layer size to $H=4C$ where $C$ is the number of channels. This way the network
size scales with the number of channels, so we just explore them together as
one hyperparameter. Denoting the hidden-layer activation function as $u$, we
can specify the mapping $F$ through the elements of the output vector as

$f^{(n)}_{ijc}=\sum_{h\in[1,H]}\Bigg{(}W_{1}^{ch}u\Big{(}\sum_{\begin{subarray}{c}c^{\prime}\in[1,C]\\\
k\in[1,K]\end{subarray}}W_{2}^{c^{\prime}kh}z^{(n)}_{ijc^{\prime}k}\Big{)}\Bigg{)}+v^{c}\equiv
F(z^{(n)})\;.$ (3)

The weights $W_{2}$ mix information between the channels and kernels
independently of the position $i,j$ to determine the activation of each hidden
node $h$. The weights $W_{1}$ then combine the hidden nodes to construct the
output value for each channel $c$, again independently of $i,j$. We emphasise
that the same set of weights and biases is applied at every pixel, consistent
with the separation of spatial and channel mixing between the two stages of
the process.

##### Stochastic mask

The final step is to increment a random subset of state vector elements
$x^{(n)}_{ijc}$ by applying a _mask_ $\sigma^{(n)}=(\sigma^{(n)}_{ijc})$,
where $\sigma^{(n)}_{ijc}$ are independent Bernoulli random variables with
parameter $1-p$. That is $\sigma^{(n)}_{ijc}=0$ with probability $p$, which is
related to the dropout rate in machine learning regularisation, and
$\sigma^{(n)}_{ijc}=1$ otherwise. The purpose of this mask is to break any
global synchronisation between cells.

Given the above, the update specified by the NCA is

$x^{(n+1)}=x^{(n)}+\sigma^{(n)}F(g\ast x^{(n)}))\equiv\Phi(x^{(n)},\theta)$
(4)

where we have introduced the mapping
$\Phi(\cdot,\theta):\mathcal{X}\rightarrow\mathcal{X}$ from one state vector
to the next, where $\theta$ encodes the network parameters ($W_{1},W_{2},v$).
In terms of individual elements, this corresponds to

$x^{(n+1)}_{ijc}=x^{(n)}_{ijc}+\sigma^{(n)}_{ijc}f^{(n)}_{ijc}$ (5)

where the elements $f^{(n)}_{ijc}$ are given by Eq. (3) above. See line 7 of
Algorithm 1. Hence, given $x^{(0)}$, the NCA provides recursively the sequence
$(x^{(n)})_{n=0,1,2,\dots,N}$.

Algorithm 1 Pseudocode description of a single NCA update step. Here $x$ is an
$M\times M$ lattice with $C$ channels. $g*x$ represent the convolutions
described in Eq.1 . $W_{1}$ and $W_{2}$ are the neural network weight
matrices, with $v$ being a vector of biases, all of which are encoded in
$\theta$. $u()$ is the activation function. Note that in practice the For
loops in line 4 and convolutions in line 3 are efficiently parallelised in
Tensorflow’s GPU implementation. Random() samples a uniform $[0,1)$
distribution.

1:function $\Phi$($z,\theta$)

2: $W_{1},W_{2},v\leftarrow\theta$

3: $z\leftarrow g*x$

4: for all $(i,j)\in[1,M]^{2}$ do

5: $dx\leftarrow W_{1}u(W_{2}\bullet z[i,j])+v$

6: if Random() $\geq p$ then

7: $x[i,j]\leftarrow x[i,j]+dx$

8: end if

9: end for

10: return $x$

11:end function

##### Batch parallelism

When implementing this model in tensorflow, we make use of _batch parallelism_
, where instead of training on one trajectory $(x^{(n)})_{n}$, we train
simultaneously on a set (batch) of trajectories $(x^{(n,r)})_{n,r}$, where
superscript $r=1,2,\dots,R$ denotes the batch number. In effect this just adds
an extra batch dimension to $x^{(n)}$, so $x^{(n)}_{ijc},z^{(n)}_{ijdk}$ and
$f^{(n)}_{ijc}$ become $x^{(n,r)}_{ijc},z^{(n,r)}_{ijdk}$ and
$f^{(n,r)}_{ijc}$ respectively. This is normally done to leverage low-level
speed-up, as training the network on batches involves matrix-matrix rather
than matrix-vector multiplications, which are well optimised on parallel
architectures (GPUs). However in the case of NCA, batch parallelism enables
far more diverse systems to be learned, for example learning several distinct
trajectories by the same rule, or improving stability through data
augmentation. It is also crucial for extending the existing NCA training
algorithm [11] to learning longer sequences of data, as discussed in section
2.2.

To summarise, the NCA can be described in terms of neural network language as
a residual (calculating increments to each pixel) Recurrent Convolutional
Neural Network (RCNN) with per-pixel dropout (stochastic mask). The
hyperparameters are the number of hidden channels, the set of convolution
kernels and activation functions on the hidden layer of the network. The
effect of varying the hyperparameters is explored in Section 3.

###  2.2 Training techniques

The neural network architecture described above is very minimal in comparison
to the state-of-the-art in the literature [23], featuring only a single hidden
layer applied in parallel. By contrast, the training process is fairly
complex. We set out key steps below, with corresponding pseudo-code set out as
Algorithm 2. NCA trajectories can be considered as paths in $\mathcal{X}$
(Figure 2), where the training process constrains the NCA parameters such that
the trajectories pass as close (defined by the loss function) as possible to
the observed data points in $\mathcal{X}$. Projecting these paths onto 1D
helps visualise the training process, especially when training to multiple
batches.

The technique of training NCA is based on backpropagation through time, a
typical method for training RNNs [24]. We have established the batch of NCA
trajectories $(x^{(n,r)})_{n=1,\dots,N}$, with $n$ and $r$ denoting time and
batch respectively. Originally [11], training the NCA consisted of one set of
initial states being mapped to one set of final states $x^{(0,r)}\rightarrow
x^{(t,r)}$. We extend this to learn the set of transitions
$x^{(n-1,r)}\rightarrow x^{(n,r)}$, where the batch parallelism allows us to
train the NCA on each transition simultaneously. This allows training NCA to
far more diverse and complex dynamics. For clarity we drop the explicit batch
index (i.e. set $r=1$) for now, the context where it matters is discussed
later, but even so batch parallelism is still exploited for learning the
different timesteps.

We have a time series of data $(y^{(\delta)})_{\delta=0,1,\dots,D}$, but only
for every $t$ NCA timesteps, that is at times $\delta t$ for
$\delta=0,1,\dots,D$, where $N=Dt$. Hence we need to compare $x^{(\delta t)}$
to $y^{(\delta)}$. We initialise $x^{(\delta t)}=y^{(\delta)}$ for
$\delta=0,\dots,D-1$, and propagate each state through
$\Phi^{t}(\cdot,\theta)=\Phi\circ\dots\circ\Phi(\cdot,\theta)$ ($t$ nested
function compositions). To compute the loss, we compare $x^{(\delta t)}$ to
$y^{(\delta)}$ for $\delta=1,\dots,D$, averaging over different $\delta$:
$\frac{1}{D}\sum_{\delta=1}^{D}\mathcal{L}(x^{(\delta t)},y^{(\delta)})$,
where the loss function $\mathcal{L}:\mathcal{X}^{2}\rightarrow\mathds{R}$ is
a meaningful measure of distance between any pair of $x^{(i)}$ and $y^{(j)}$.
Training is achieved by minimising the loss function, which requires partial
gradients with respect to the trainable parameters $\theta$ to be evaluated.
For the full gradient calculations, see appendix A.

Algorithm 2 Training an NCA $\Phi$ to $R$ data trajectories of length $D$.
Split into $B$ mini-batches to reduce memory usage. Here $x^{(n,r)}$ and
$y^{(n,r)}$ denote predicted state and data at step $n$ of trajectory $r$
respectively. $\hat{x}^{(n,r)}$ is a temporary variable to store new
intermediate states at each training iteration. Lines 20 and 21 perform
gradient normalisation followed by a parameter update handled by the Nadam
algorithm (or any other optimiser of choice). The nested For Loops on lines
9,11 and 12 are easily parallelised. Typical choices of mini-batching are such
that $10<(D\times R)//B<100$. $\Phi^{t}$ denotes $t$ iterations of $\Phi$. The
choice of $t$ is an important hyperparameter, and relates to the temporal
resolution of the data $x$. The model parameters $\theta$ are assumed to
either be initialised appropriately, or already partially trained.
RandomShuffle$(A,B)$ randomly shuffles $A$ and splits it into $B$ equal sized
chunks. $\mathcal{L}(A,B):\mathcal{X}\times\mathcal{X}\rightarrow\mathds{R}$
computes the loss between states $A$ and $B$.

1:function Train($\Phi,\theta,y,t$,B,EPOCHS)

2: $x\leftarrow y$

3: $D\leftarrow y.shape[0]$

4: $R\leftarrow y.shape[1]$

5: for $i\in$ EPOCHS do

6: Grad $\leftarrow\vec{0}$

7: DS $\leftarrow$ RandomShuffle([1,D],B)

8: RS $\leftarrow$ RandomShuffle([1,R],B)

9: for $b\in[1,B]$ do

10: Loss $\leftarrow 0$

11: for $\delta\in$ DS$[b]$ do

12: for $r\in$ RS$[b]$ do

13: $\hat{x}^{(\delta,r)}\leftarrow\Phi^{t}(x^{(\delta-1,r)},\theta)$

14: Loss $\leftarrow$ Loss $+\mathcal{L}(\hat{x}^{(\delta,r)},y^{(\delta,r)})$

15: end for

16: end for

17: Loss $\leftarrow\frac{1}{D\times R}$ Loss

18: Grad $\leftarrow$ Grad $+\frac{\partial\text{Loss}}{\partial\theta}$

19: end for

20: Grad $\leftarrow$ Norm(Grad)

21: Update($\theta$,Grad,$i$)

22: for $\delta\in[1,D]$ do

23: for $r\in[2,R]$ do

24: $x^{(\delta,r)}\leftarrow\hat{x}^{(\delta,r)}$

25: end for

26: end for

27: end for

28: return $\Phi$

29:end function

Figure 2: 1D phase space representation of NCA trajectories, predictions
$x^{(\delta t,r)}$ and true states $y^{(\delta,r)}$. Here $D=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.

There are additional practical considerations for optimising this training
method. After each iteration, we have the choice of keeping and further
propagating the values in $x^{(\delta t)}$ for $\delta=1\dots D$, or re-
initialising them: $x^{(\delta t)}=y^{(\delta)}$. Propagating them allows the
NCA to better learn long term dynamics (particularly of the hidden channels)
over many training iterations. However, we observe that re-initialising helps
speed up the training process in the earlier steps. As both approaches have
their advantages, we return to the batch parallelised case and do both. We
regard re-initialising the states as a form of data augmentation (Figure 2),
so in practice we only re-initialise one batch: $x^{(\delta
t,1)}=y^{(\delta,1)}$. This choice of only re-initialising one batch performs
well, but is arbitrary and could be further tuned for specific problems. The
NCA is initialised with random hidden layer weights ($W_{2}$), and zero final
layer weights ($W_{1}$) and bias ($v$).

When implementing the training procedure, as described in Algorithm 2, the
additional subtlety of mini-batching is required [25, 26]. Rather than
computing the gradient for transitions $x^{(\delta-1,r)}\rightarrow
x^{(\delta,r)}$ for all $\delta\in[1,D],r\in[1,R]$, we randomly split
$[1,D]\times[1,R]$ into $B$ _mini-batches_ , and separately compute the loss
gradient for each mini-batch. After iterating through each mini-batch, the
gradients are averaged and applied once to the model parameters. The need for
mini-batching is due to memory constraints: if a large enough number of
batches $R$ or timesteps $D$ is used, computing the gradient over the full set
of transitions is unfeasible. In Algorithm 2 the memory cost of calculating
the gradient (line 18) scales like $D\times R\times M^{4}\times C^{2}\times t$
(where $|\mathcal{X}|=M^{2}C$). By contrast the memory cost of storing and
adding to the calculated gradient over each mini-batch is fixed as
$\|\theta\|$ (i.e. does not scale with $B$), which is minimal given the small
size of the network. $M$ and $C$ are fixed by the spatial (and channel)
resolution of the data, but mini-batching reduces the memory burden of $D$ and
$R$. For the full calculation see appendix A. In the case of $B=1$, the mini-
batching reduces such that the For loops at lines $9,11$ and $12$ collapse
into simpler loops over $[1,D]$ and $[1,R]$. When training on image morphing
[11], the size of data does not require mini-batching as $D$ and $R$ are
small. To accurately capture PDE dynamics, much larger $D$ is used, and as
such mini-batching is required. We do not explore spatial mini-batching, where
random subsets of pixels are tracked for gradients, as this makes efficient
parallelisation more challenging, however it could be very useful to further
explore as it could enable training of NCA to higher resolution data.

####  2.2.1 Loss functions

We now turn to the question of how to define the distance between a NCA
trajectory and target data. This problem is split naturally into two parts:
first, how to find the difference between corresponding (time and batch
labelled) points in $\mathcal{X}$, and then how to combine all these
difference measures to the distance between two sets of points in
$\mathcal{X}$. For the latter choice, we adopt an arithmetic mean over the
time points and batches considered (as shown in lines 14 and 17 of Algorithm
2), although these could be weighted (e.g., the states at a specific time or
batch being most important).

We compared several loss functions for corresponding points in $\mathcal{X}$,
but we found that the standard euclidean norm
$\mathcal{L}(x,y)=\big{(}\sum_{i,j,c}(x_{ijc}-y_{ijc})^{2}\big{)}^{\frac{1}{2}}$
worked best in all contexts. Probability mass based losses (Hellinger [27] and
Bhattacharyya [28]) and distance between spatial Fourier transforms of points
in $\mathcal{X}$ all work well in the PDE modelling case, but perform poorly
on image morphing. Various approximations to the Wasserstein distance [29]
performed poorly in both contexts but still remain promising given their
success in texture synthesis [30, 31, 32]. As such, for the following results
we stick to the euclidean distance, although we recommend experimenting with
different loss functions depending on the system one is modelling. Any
differentiable function
$\mathcal{L}:\mathcal{X}\times\mathcal{X}\rightarrow\mathds{R}$ can be used,
if minimising it’s output constrains the inputs in a desirable way.

##  3 Results

We now demonstrate the applicability of the NCA to two contrasting use cases.
First, we consider training data comprising numerical solutions of coupled
nonlinear reaction-diffusion equations, with parameters that produce Turing
patterns. Equations in this class are widely used to model complex biological
systems such as in: developmental biology [33, 34]; ecology [35]; and skin
pattern morphogenesis [36]. We adopt a representative example that is capable
of generating a wide variety of spatial patterns, in the context of chemical
reactions [17]. We demonstrate in particular that the NCA can generalise
beyond the set of initial conditions in its training data. We then turn to an
artificial image morphing problem, inspired by [11], and show that the same
NCA is capable of constructing local rules to effect the desired dynamics. In
contrast to the reaction-diffusion system, such rules are not known _a priori_
, so it is not obvious that they exist. This problem also lends itself to
testing the robustness of the rules that result. When training to PDEs, the
focus is on exploring training hyperparameters (loss function, time sampling).
After determining suitable training hyperparameters, we then show that this
generalises to the image morphing task, where we explore model hyperparameters
(kernels, activation functions, number of hidden channels).

###  3.1 Gray-Scott reaction diffusion equations

The Gray-Scott [17] reaction diffusion equations read

$\displaystyle\partial_{t}A$
$\displaystyle=D_{A}(\partial_{xx}+\partial_{yy})A-AB^{2}+\alpha(1-A)$
$\displaystyle\partial_{t}B$
$\displaystyle=D_{B}(\partial_{xx}+\partial_{yy})B+AB^{2}-(\gamma+\alpha)B$

in which $D_{A},D_{B},\alpha$ and $\gamma$ are parameters. These describe two
species, $A$ and $B$, which diffuse in two-dimensional space with diffusion
constants $D_{A}$ and $D_{B}$, respectively. Species $A$ grows towards a
density of $A=1$ at a rate $\alpha$, whilst species $B$ dies out at rate
$\gamma+\alpha$. The species also undergo the reaction $A+2B\to 3B$.

With $D_{A}=0.1,D_{B}=0.05,\alpha=0.06230,\gamma=0.06268$ we obtain maze-like
patterning (Figure 3). We solve these PDEs with an Euler finite-difference
discretisation scheme, with time-step of 1, for $N=1024$ steps, and on an
$M\times M$ lattice of size $M=64$. $\alpha$ and $\gamma$ parameterise the
patterning type, whereas $D_{A}$ and $D_{B}$ re-scale the patterns, and must
be chosen in line with the timestep size to achieve numerical stability.

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$.

####  3.1.1 Effect of training hyperparameters

We begin with a basic NCA architecture that employs just $K=2$ kernels, the
identity and Laplacian, $p=0$ (purely deterministic), $C=8$ channels in total
(so 2 observable and 6 hidden channels) and a rectified linear unit (relu,
$u(z)=\frac{|z|+z}{2}$) activation function. We found that the Nadam optimiser
[37] consistently performed well. Nadam is a modification of the widely used
Adam optimiser [38], with the only difference being the use of Nesterov
momentum [37]. Optimisers based on Nesterov momentum perform well, both in
theoretical convergence and generalisability of trained deep neural networks
[37, 39]. Note that we also employ gradient normalisation [40] before passing
gradient information to the optimiser - this was found to significantly
improve training performance. This just leaves the time sampling ($t$) as the
main hyperparameter to optimise. Time sampling is subtle in the case of PDEs
as numerical integration of the PDE system necessarily involves a discrete
integration time step. We can sample the trajectories at coarser intervals,
increasing $t$ in Algorithm 2 such that each NCA update corresponds to a
timestep of the PDE solver. In other words, we only compare every
$t^{\text{th}}$ PDE and NCA step.

We found that while training loss increased with greater sampling intervals
$t$, tests on unseen initial conditions achieve comparable loss for most
sampling intervals, with modest improvements for greater $t$ (Figure 4). The
training loss is calculated for $N=1024$ steps (as in Figure 3), whereas the
test loss is calculated over $N=2048$ steps from an unseen initial condition.
This demonstrates generalisation both to unseen initial conditions, and longer
simulation times. Note that the unseen initial condition used for testing
features high frequency components not observed at all during training, and
that NCA trained with high sampling were sometimes numerically unstable to
these high frequency inputs (missing test loss points in Figure 4A, or
snapshots at $t=2$,$t=6$ in Figure 4B).

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:
snapshots of NCA trajectories (at $n=2048$) based on unseen initial
conditions, with varying sampling $t$. Each NCA is trained for 4000 epochs,
with a mini-batch size $B=64$.

Figure 5 shows various snapshots from true (PDE) trajectories alongside the
corresponding snapshots from an NCA trained with $t=32$. This extrapolates an
unseen initial condition far beyond the time observed during training
($n\in[0,1024]$), demonstrating that the NCA does learn the underlying rules
of the dynamics rather than overfitting to the training trajectories. When we
considered finer sampling $t$ (Figure 4), we observe more frequent numerical
instabilities, or complete failure to learn dynamics. Coarse time sampling
appears to both stabilise these numerical problems, and yield more
generalisable models. We posit this is due to coarse time sampling allowing
the NCA to be less constrained during training, in that intermediate states
may explore more possible states, increasing the chances of finding $\theta$
that gives the correct dynamics. Fine time sampling perhaps over-constrains
the NCA, leading to instabilities or training converging to sub-optimal local
minima of the loss landscape. Alternatively, the behaviour at fine time
sampling is consistent with overfitting, so coarser time sampling could be
considered a regularising technique here.

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.

In summary, we have found that the NCA architecture set out in Section 2 is
capable of learning update rules that reproduce the solution of a certain pair
of coupled nonlinear PDEs. Our main finding is that good rules can be learnt,
but coarse time sampling improves numerical stability, and learns rules that
generalise better to unseen initial conditions. In the next section we show
that the algorithm still performs well when there is no known underlying
integration timestep.

###  3.2 Image morphing

Figure 6: Image morphing task. Given a space invader initial condition, morph
through a microbe and remain stable at a rooster pattern.

We now test whether a training method that works well for PDEs also works on
the image morphing task. The task comprises an initial image, an intermediate
image to morph through, and then a final image that is intended to remain
stable (i.e., be an attractor of the dynamics) (Figure 6). This latter
requirement is incorporated into the training data by repeating the final
image twice. The images shown were downsampled to a resolution of $60\times
60$ to reduce computational cost. We further impose fixed boundary conditions,
that is, to insist that the state vector $x^{(n)}$ vanishes at the boundary
points. This is because the system is not periodic, as was the case for the
PDE problem.

Reliable training of the NCA requires a careful construction of the training
data. A side-effect of the fixed boundary conditions is that the NCA can learn
to grow an image from the boundary, rather than from the initial condition. We
would like the pattern formation to remain translationally invariant - the
image morphing should behave independently of the boundaries, and should only
depend on the input images. To enforce this, we embed the training data within
a larger lattice, thereby reducing the influence of the boundaries.
Translational invariance is further encouraged by training to several copies
of the image sequence, each randomly shifted in space, to prevent learning any
effective long range boundary interaction. It is possible to train NCA to
produce textures as _global_ attractors [30, 13] as textures are
translationally invariant. We believe it is impossible to have a fixed pattern
(i.e. not a translationaly invariant texture) as a global attractor — if
boundary effects are removed so the whole NCA system is translationaly
invariant. Avoiding the desired dynamics being a global attractor ensures that
the NCA has learned a rule that maps the input state through the sequence of
target states, rather than just generating the target states irrespective of
initial conditions. This is further verified by exploring stability under
perturbation in Section 3.2.2

We also find that to train the NCA to reach a stable state, in effect mapping
an image to itself, augmenting the training data with noise is necessary.
Without noise augmentation, the training process crashes as gradients diverge
(when training the final stable transition), but adding a very small amount of
random noise to the data fixes this. We can understand the effect of this
noise as introducing a basin of attraction around the desired final state, and
training to noisy images enhances the robustness of the NCA to noisy
perturbations.

####  3.2.1 Effect of model hyperparameters

The training hyperparameter $t$ that corresponds to the frequency of time
sampling cannot be assumed to translate from the PDE case to the image
morphing task. As there is no underlying physical mechanism connecting the
images, there is nothing to provide a basic unit of time. We are however
guided by the fact that the update rules are local, and therefore initially
take the number of timesteps to be 64, which is similar to the lattice size
and therefore gives sufficient time for information to propagate across it. As
we explore this point more below, we find in fact that fewer timesteps can
also be sufficient.

While a deterministic update ($p=0$) is appropriate in the PDE case, as this
was a feature of the training data, for the image morphing problem, we update
stochastically with $p=\frac{1}{2}$, as removing global synchronisation
between each cell acts like dropout regularisation, and can be biologically
motivated [11]. Note that a choice of $p=\frac{1}{2}$ effectively halves the
number of timesteps between images. We found that varying the update
probability $p$ had very little direct impact on model performance (except for
extreme values close to 1 or 0), instead the effective number of timesteps
between images was explored.

Figure 7: NCA trained on image morphing task with different kernels and
activations. 16 channels, 64 steps between images. A,B: training loss and
snapshots of NCA with relu activation and various kernels. C,D: training loss
and snapshots of NCA with Identity, Sobel and Laplacian kernels, for various
activation functions.

The system state has 4 observable channels—the red, green, blue and alpha
(transparency) components—and 12 hidden channels. Since the images are not
rotationally symmetric, and we have no prior underlying mechanisms to
constrain symmetries of update rules, we consider adding Sobel kernels to the
identity and Laplacian kernels that were used in the PDE case. The Laplacian
kernel detects spatial changes (and curvature) in patterns, whereas the Sobel
kernels also detect the orientation of any changes. We find that including the
symmetry breaking Sobel kernels improves the performance over just using
symmetric kernels (Figure 7A,B). This does however break the symmetry of the
NCA—such an NCA is unlikely to be stable under rotations or reflections of the
initial condition, as discussed in Section 3.2.2. We also find that the best
performing activation function is relu (Figure 7C,D). The linear case,
$u(z)=z$, can be thought of as an absence of an activation function.
Surprisingly, the overall shape of the final image is reasonably well
reproduced, although it clearly lacks the definition achieved with the other
activation functions. This justifies the additional complexity of nonlinear
activations.

Figure 8: NCA trained on image morphing task. Relu activation; Identity, Sobel
and Laplacian kernels. A,B: Training loss and snapshots of 16 channel NCAs
trained with different time sampling. C.D: Training loss and snapshots of NCAs
trained with time sampling of 32, and various numbers of channels.

Exploring how NCA behaviour scales with number of channels, we find that
unsurprisingly more hidden channels performs better (Figure 8A,B), capturing
more of the details in the rooster image. The number of channels functions as
a clear ‘model size’ parameter, and we find that model performance scales
nicely with this measure of model size. We also explore how NCA train for
different numbers of timesteps between images (Figure 8C,D). It is surprising
that with as few as 8 timesteps, the basic shape and colour of the rooster are
correct (although details are better at 16 or 32 steps), which highlights the
locality of the update rule. The image resolution is $60\times 60$, and with 8
timesteps between images only cells less than 16 pixels away can communicate
(from initial condition to reaching the stable rooster pattern). However with
the stochastic cell updates, the effective communication range from initial to
final condition here is halved to just 8 pixels. This emphasises that the
update rule is local, in that local structures of the initial condition morph
into local structures of the final state at similar locations.

####  3.2.2 Stability analysis

Figure 9: Local stability behaviour of two NCA. A:i 32 channels, 32 steps
between images. B: 16 channels, 64 steps between images. Top left heatmap in
each case shows how many pixels of the final image change (by more than 0.1 to
account for random fluctuations) when that pixel is perturbed in the initial
condition. The other images all show snapshots of the final state when the
initial condition is perturbed locally, for different perturbation locations.

With a trained NCA, there is an obvious question of stability—if an initial
condition is perturbed away from what the NCA is trained on, how does this
affect the behaviour? We consider three kinds of perturbations of the initial
condition: local perturbations, global perturbations, and symmetry
perturbations. Local perturbations change one pixel, and allow us to explore
how errors propagate through the spatial part of the NCA. Global perturbations
can show how resilient NCA are to noisy inputs. Symmetry perturbations, such
as rotations or reflections of initial conditions, allow us to explore how NCA
respect desirable symmetries.

Stability under local perturbations depends strongly on how many timesteps
between images the NCA is trained on (Figure 9). We find that NCAs with fewer
time-steps are more stable to local perturbations, or conversely that allowing
more NCA steps between training images gives more time for local perturbations
to travel. In both cases the perturbations remain mostly local. Using local
perturbations could help calibrate the number of timesteps to use when
modelling real systems, in that the NCA should have the same response to
perturbations as the underlying data being modelled.

Figure 10: Rightmost column shows extrapolation beyond training time,
demonstrating stability of the final state. Top row shows snapshots from
unperturbed trajectory. Middle row shows snapshots from minimal initial
perturbation that destroys the final state (minimising
$\kappa^{(96)}(x^{(0)},\tilde{x}^{(0)})$). Bottom row shows snapshots from
maximal initial perturbation that preserves the final state (maximising
$\kappa^{(96)}(x^{(0)},\tilde{x}^{(0)})$). NCA (32 channels; Identity, Sobel
and Laplacian kernels; time sampling $t=32$, relu activation) trained on image
morphing task.

To address the question of stability with respect to global perturbations, we
frame it as an optimisation problem. Let
$\kappa^{n}(x^{(0)},\tilde{x}^{(0)})=\|\tilde{x}^{(0)}\|-\|\Phi^{n}(x^{(0)}+\tilde{x}^{(0)})-\Phi^{n}(x^{(0)})\|$,
where $\tilde{x}^{(0)}$ is a perturbation of initial condition $x^{(0)}$. By
finding $\tilde{x}^{(0)}$ that maximises or minimises $\kappa^{n}$, we can
find a maximally perturbed initial condition $x^{(0)}+\tilde{x}^{(0)}$ that
leaves a future state $x^{(n)}$ unchanged, or a minimal perturbation that
destroys $x^{(n)}$. This allows us to explore the space of initial conditions
around which the NCA was trained, and can reveal which features of an initial
condition are important. For example, it may be only the edges and corners of
an image are learned from. As the whole NCA process is differentiable, we can
use gradient based optimisation on $\kappa^{n}$ to find $\tilde{x}^{(0)}$.
Finding minimal perturbations that destroy $x^{(n)}$ is similar to adversarial
attacks of classifier networks, where small changes to an image completely
destroy the behaviour of an image classifier. Figure 10 shows the behaviour of
a trained NCA (best performing model shown in figure 8A,B) starting from
examples of these adversarial initial conditions. It is possible to find
initial conditions that are visually similar to the true initial condition,
and yet they destroy the stable rooster pattern ($x^{(96)}$). We can also find
large perturbations of the initial condition that leave the target state
($n=96$) unperturbed, however the long term stability of the rooster pattern
is still damaged.

Figure 11: Behaviour of trained NCA on symmetrically perturbed inputs. Left
column shows inputs, middle two shows final state behaviour for NCA with
asymmetric kernels (identity, Sobel and Laplacian), rightmost two shows final
state behaviour for NCA with symmetric kernels (identity, Laplacian, average).
Augmented data examples show NCAs trained to trajectories rotated to random
angles.

Finally, we compare the behaviour of different NCA models on symmetrically
perturbed initial conditions (Figure 11). By rotating or flipping the input
image we obtain symmetrically perturbed inputs. One NCA is trained on _normal
data_ , that is, data that has only been translationally perturbed to minimise
boundary effects. The other is trained on _augmented data_ that also includes
the same training data after applying global rotations about random angles. We
also explore the effect of restricting the NCA to include only symmetric
kernels, rather than both symmetric and asymmetric kernels. We find that even
without any data augmentation, the symmetric kernel NCA already performs very
well, although it struggles with the off lattice 45 degree rotations. When
trained on rotationally augmented data, the asymmetric kernel NCA improves its
performance on rotated inputs, but is still outperformed by the symmetric
kernel NCA for on-lattice rotations (with or without data augmentation). Off
lattice rotations are the most challenging, and seem to be where data
augmentation is necessary even for symmetric kernel NCA. Overall, we find that
the symmetric kernel NCA better handles symmetric perturbations, whilst the
asymmetric kernel NCA performs best on the unperturbed data. As one might
expect, building symmetry into the model allows it to solve a broader range of
symmetrically-related problems, whereas leaving it out promotes specialisation
towards a single problem.

##  4 Discussion

We have demonstrated NCA as a framework for modelling emergent spatio-temporal
patterning. Many systems in biology are characterised by complex emergent
phenomena of locally interacting components, and finding interaction rules or
mechanisms that lead to specific emergent behaviours is a challenging inverse
problem. By making classic cellular automata differentiable, NCA present a new
approach to this class of problems. [11] demonstrated that a trained NCA can
generate complex structures (specifically emojis) that remain stable over time
and to perturbation. Here we have extended this approach to learn dynamics
from snapshots of a pattern at multiple timepoints (rather than just the end-
point), i.e., we show the ability to learn dynamic patterns, specifically
those arising from PDEs with Turing instabilities that been widely used to
study biological pattern formation.

Specifically, we showed in Section 3.1 that NCA can infer update rules
equivalent to those obtained by discretising and iterating a set of PDEs. NCA
have an inductive bias to learning underlying dynamics, rather than just
overfitting to data, due to their minimal parameters and hard coded local
kernels. We demonstrate this by presenting the trained NCA with initial
conditions that were not part of the training data, and finding that the
predicted trajectories are similar to those obtained directly from the PDEs -
the trained NCA generalise well. This suggests that an NCA trained on
experimental data could be used to predict the behaviour of that system under
conditions that have not been directly observed. We have also discussed NCA
hyperparameters in more detail than most previous work, which can provide
guidance for future exploration. For example, tuning the number of timesteps
between images can constrain how far local information is allowed to spread,
and the number of channels required to accurately capture a patterning
behaviour could function as a heuristic for the complexity of that pattern.

More generally, the findings of Section 3.2 confirm that NCA can be used as a
tool to construct local dynamical update rules whose emergent properties
satisfy certain constraints. These constraints include the end state, but
extend also to the stability of that configuration, invariance of the dynamics
under certain symmetry operations and the effect of boundary conditions. Given
that for any observed emergent patterning, many possible microscopic update
rules could exist, constraining how NCA respect symmetries or behave around
boundaries helps reduce the set of possible rules. As the training process for
NCA amounts to a differential optimisation procedure, microscopic rules that
yield desired emergent behaviour can be efficiently found, even when
constraints are imposed, as exemplified by the image morphing task.

Using NCA as an in silico way to study the behaviour of growing systems is a
recurring theme in the literature [11, 21], discussed more concretely in [41].
Further applications or extensions of NCA models have been explored in the
context of image processing and synthesis [42, 43]. Here NCA models have been
coupled to CLIP text embeddings [44], in line with recent machine learning
based text to image techniques. NCA are clearly capable of a diverse range of
behaviours, but all the previous literature just focuses on training one set
of transitions, rather than full dynamic trajectories. We believe that the
extension to learning dynamics of arbitrarily long sequences dramatically
increases the already wide range of systems and behaviours NCA can model. We
believe that being able to train to sequences of images will better enable NCA
to be applicable to modelling real biological systems, and it will probably
enable more interesting image synthesis techniques.

Compared to most current machine-learning research, our chosen neural network
architecture is economical in terms of the number of trainable parameters.
This not only makes training more computationally efficient, but also adheres
also to the aesthetic guidance of modelling traditions in physics and
mathematics that simpler models are preferred over more complex models when
they have comparable descriptive power. Here we have found that a single
hidden layer is sufficient to model a variety of systems and reproduce a wide
range of behaviours. The reason for this may lie in part due to hidden
channels in the state space, as these can encode complex environmental
information such as boundary conditions, as well as encoding memory for longer
term dynamics. This spatially distributed memory encoding in the hidden
channels could be likened to previous work on differentiable neural computers
[45]. We note that NCA also link back to older work on amorphous computing
[46], providing a connection between modern machine learning and theory of
spatially distributed computational frameworks.

There are however a few shortcomings of NCA, mainly the underlying assumption
that purely local interactions are sufficient. There will be systems with non-
local (or multiscale) interactions that cannot be elegantly explained with
purely local rules. We have also assumed that the update rules are constant
over time, even though many complex (physical or biological) systems are
highly time dependent [47]. Whilst it is possible that the hidden channels in
the NCA could encode nonlocality and time-dependence, it might be more natural
(and interpretable) to extend the representation of the dynamics in the neural
network to incorporate such dependencies explicitly, for example by increasing
the size/stride of convolution kernels, or by including an explicit time
parameter as a network input. A possible risk with including explicit time
dependence is that the NCA could over-fit to the timestep, rather than
learning the microscopic interactions that yield the emergent behaviour. To
tackle such questions, it might be desirable to augment the loss functions
with measures of model complexity as a means to converge on the most
parsimonious description, for example by sparsity regularisation. Generalising
NCA to better work on multiscale systems could also be worth exploring, for
example by coupling NCA-like models on lattices of different resolutions.

Similarities can be drawn between NCA and other data driven equation discovery
techniques like SINDy [48] or Extended Dynamic Mode Decomposition (EDMD) [49].
Both SINDy and EDMD have the main purpose of fitting dynamical systems to
data, both in the cases of ODEs and PDEs. SINDy enforces parsimonious
solutions through sparity regularisation, whereas Dynamic Mode Decomposition
is analogous to Singular Value Decomposition for time series data (and the
Extended DMD is a nonlinear generalisation). The key areas where NCA differ is
in background motivation, and the sorts of systems they’re suited to. NCA act
as a bridge between machine learning based image processing, and learning
simple models of complex systems; whereas SINDy and EDMD were developed in the
context of data driven engineering and fluid dynamics respectively. Cellular
automata models are more general than (local) PDEs. Although we restrict
ourselves to differential operator kernels, we don’t have to - learning
arbitrary kernels would provide far more general expressive power, whilst
still keeping a minimal (and importantly local) model. The class of models
that could be learned with arbitrary kernels includes local PDEs, but is far
more general.

A further development of these NCA models is to make them truly distributed
during learning. When computing the loss, information of the whole lattice is
needed, which places limits on the size of a lattice that can be handled with
available computation time and memory. If instead an NCA could be trained with
a purely local loss, such that model weights are updated for each pixel based
on its neighbours, more advanced training procedures could be exploited. In
essence, if the only global communication is the updates to the model weights,
rather than the full lattice state, NCA could be trained using online
training, or on much larger lattice sizes. An alternative approach to increase
the resolution that can be trained on would be to randomly sub-sample elements
[50] of the NCA lattice when computing losses.

Although we explored the stability of NCA under perturbations to the
trajectory (initial condition), we did not address stability under
perturbation of model parameters. This could naturally tie in to
interpretation of the trained model, for example, by assessing stability under
perturbation of network weights. Performing network pruning [51] could be
another powerful approach, yielding even more minimal models. The recently
popular field of explainable AI [15, 14] likely offers some other tools that
would enable this. A worthwhile further development would be to reverse-
engineer a concise analytic expression for the underlying PDE from the trained
NCA parameters, for example with symbolic regression techniques like Sparse
Identification of Nonlinear Dynamics (SINDy) [48]. Such an approach could be
tested with NCAs trained on known PDEs, but it would be interesting to then
apply this to systems where we don’t know the underlying mechanics, such as
image morphing or biological data.

## Acknowledgements

This work has made use of the resources provided by the Edinburgh Compute and
Data Facility (ECDF) (http://www.ecdf.ed.ac.uk/). For the purpose of open
access, the author has applied a Creative Commons Attribution (CC BY) licence
to any Author Accepted Manuscript version arising from this submission. Alex D
Richardson was supported by the EPSRC Centre for Doctoral Training in
Mathematical Modelling, Analysis and Computation (MAC-MIGS) funded by the UK
Engineering and Physical Sciences Research Council (grant EP/S023291/1),
Heriot-Watt University and the University of Edinburgh.

## References

  * [1] Somayeh Kalantari, Eslam Nazemi and Behrooz Masoumi  “Emergence phenomena in self-organizing systems: a systematic literature review of concepts, researches, and future prospects”  In _Journal of Organizational Computing and Electronic Commerce_ 30.3 Taylor & Francis, 2020, pp. 224–265  DOI: 10.1080/10919392.2020.1748977
  * [2] Martin Gardner  “Mathematical Games”  In _Scientific American_ 223.4 Springer ScienceBusiness Media LLC, 1970, pp. 120–123  DOI: 10.1038/scientificamerican1070-120
  * [3] Palash Sarkar  “A Brief History of Cellular Automata”  In _ACM Comput. Surv._ 32.1 New York, NY, USA: Association for Computing Machinery, 2000, pp. 80–107  DOI: 10.1145/349194.349202
  * [4] Yanhua Zhuang et al.  “A Bibliographic Review of Cellular Automaton Publications in the Last 50 Years”  In _Journal of Cellular Automata_ 12.6, 2017, pp. 475–492
  * [5] Liana Manukyan et al.  “A living mesoscopic cellular automaton made of skin scales”  In _Nature_ 544.7649 Springer ScienceBusiness Media LLC, 2017, pp. 173–179  DOI: 10.1038/nature22031
  * [6] Anamarija Fofonjka and Michel C. Milinkovitch  “Reaction-diffusion in a growing 3D domain of skin scales generates a discrete cellular automaton”  In _Nature Communications_ 12.1 Springer ScienceBusiness Media LLC, 2021  DOI: 10.1038/s41467-021-22525-1
  * [7] Axel Lange, Hans L. Nemeschkal and Gerd B. Müller  “A threshold model for polydactyly” Biological Challenges in Morphogenesis  In _Progress in Biophysics and Molecular Biology_ 137, 2018, pp. 1–11  DOI: https://doi.org/10.1016/j.pbiomolbio.2018.04.007
  * [8] J.. Bloomfield, J.. Sherratt, K.. Painter and G. Landini  “Cellular automata and integrodifferential equation models for cell renewal in mosaic tissues” OA status: green_published  In _Journal of The Royal Society Interface_ 7.52, 2010, pp. 1525–1535  DOI: 10.1098/rsif.2010.0071
  * [9] All Masoudi-Nejad et al.  “Cancer systems biology and modeling: Microscopic scale and multiscale approaches”  In _Seminars in Cancer Biology_ 30, 2015, pp. 60–69  DOI: 10.1016/j.semcancer.2014.03.003
  * [10] Elisabete A. Silva, Jack Ahern and Jack Wileden  “Strategies for landscape ecology: An application using cellular automata models”  In _Progress in Planning_ 70.4, 2008, pp. 133–177  DOI: 10.1016/j.progress.2008.05.002
  * [11] Alexander Mordvintsev, Ettore Randazzo, Eyvind Niklasson and Michael Levin  “Growing Neural Cellular Automata” https://distill.pub/2020/growing-ca  In _Distill_ , 2020  DOI: 10.23915/distill.00023
  * [12] Y. LeCun, Y. Bengio and G. Hinton  “Deep learning” LeCun, Yann Bengio, Yoshua Hinton, Geoffrey Jeong, Yongwook/N-7413-2016; Ma, Jialin/ABG-2965-2021 1476-4687  In _Nature_ 521.7553, 2015, pp. 436–444  DOI: 10.1038/nature14539
  * [13] Alexander Mordvintsev and Eyvind Niklasson  “microNCA: Texture Generation with Ultra-Compact Neural Cellular Automata”  arXiv, 2021  DOI: 10.48550/ARXIV.2111.13545
  * [14] Roberto Confalonieri, Ludovik Coba, Benedikt Wagner and Tarek R. Besold  “A historical perspective of explainable Artificial Intelligence”  In _Wiley Interdisciplinary Reviews-Data Mining and Knowledge Discovery_ 11.1, 2021  DOI: 10.1002/widm.1391
  * [15] Plamen P. Angelov et al.  “Explainable artificial intelligence: an analytical review”  In _Wiley Interdisciplinary Reviews-Data Mining and Knowledge Discovery_ 11.5, 2021  DOI: 10.1002/widm.1424
  * [16] “The chemical basis of morphogenesis”  In _Philosophical Transactions of the Royal Society of London. Series B, Biological Sciences_ 237.641 The Royal Society, 1952, pp. 37–72  DOI: 10.1098/rstb.1952.0012
  * [17] John E. Pearson  “Complex Patterns in a Simple System”  In _Science_ 261.5118, 1993, pp. 189–192  DOI: 10.1126/science.261.5118.189
  * [18] Martı́n Abadi al  “TensorFlow: Large-Scale Machine Learning on Heterogeneous Systems” Software available from tensorflow.org, 2015  URL: https://www.tensorflow.org/
  * [19] François Chollet  “Xception: Deep Learning with Depthwise Separable Convolutions”  arXiv, 2016  DOI: 10.48550/ARXIV.1610.02357
  * [20] William Gilpin  “Cellular automata as convolutional neural networks”  In _Physical Review E_ 100.3 American Physical Society (APS), 2019  DOI: 10.1103/physreve.100.032402
  * [21] Alexander Mordvintsev, Ettore Randazzo and Craig Fouts  “Growing Isotropic Neural Cellular Automata”  arXiv, 2022  DOI: 10.48550/ARXIV.2205.01681
  * [22] Md Zahangir Alom et al.  “Recurrent Residual Convolutional Neural Network based on U-Net (R2U-Net) for Medical Image Segmentation”, 2018  arXiv:1802.06955 [cs.CV]
  * [23] Shi Dong, Ping Wang and Khushnood Abbas  “A survey on deep learning and its applications”  In _Computer Science Review_ 40, 2021, pp. 100379  DOI: https://doi.org/10.1016/j.cosrev.2021.100379
  * [24] A.. Tsoi and A. Back  “Discrete time recurrent neural network architectures: A unifying review” Tsoi, AC Back, A Tsoi, Ah Chung/0000-0003-2904-7008  In _Neurocomputing_ 15.3-4, 1997, pp. 183–223  DOI: 10.1016/s0925-2312(97)00161-6
  * [25] Dominic Masters and Carlo Luschi  “Revisiting Small Batch Training for Deep Neural Networks”, 2018  arXiv:1804.07612 [cs.LG]
  * [26] Jason Brownlee  “A gentle introduction to mini-batch gradient descent and how to configure batch size”  In _MachineLearningMastery.com_ , 2019  URL: https://machinelearningmastery.com/gentle-introduction-mini-batch-gradient-descent-configure-batch-size/
  * [27] “Hellinger distance - Encyclopedia of Mathematics”  In _Hellinger distance - Encyclopedia of Mathematics_ URL: https://encyclopediaofmath.org/wiki/Hellinger_distance
  * [28] “Bhattacharyya distance - Encyclopedia of Mathematics”  In _Bhattacharyya distance - Encyclopedia of Mathematics_ URL: https://encyclopediaofmath.org/index.php?title=Bhattacharyya_distance
  * [29] Cédric Villani  “The Wasserstein distances”  In _Optimal Transport: Old and New_ Berlin, Heidelberg: Springer Berlin Heidelberg, 2009, pp. 93–111  DOI: 10.1007/978-3-540-71050-9“˙6
  * [30] Eyvind Niklasson, Alexander Mordvintsev, Ettore Randazzo and Michael Levin  “Self-Organising Textures” https://distill.pub/selforg/2021/textures  In _Distill_ , 2021  DOI: 10.23915/distill.00027.003
  * [31] Eric Heitz, Kenneth Vanhoey, Thomas Chambon and Laurent Belcour  “A Sliced Wasserstein Loss for Neural Texture Synthesis”  arXiv, 2020  DOI: 10.48550/ARXIV.2006.07229
  * [32] Antoine Houdard, Arthur Leclaire, Nicolas Papadakis and Julien Rabin  “Wasserstein Generative Models for Patch-based Texture Synthesis”  In _CoRR_ abs/2007.03408, 2020  arXiv: https://arxiv.org/abs/2007.03408
  * [33] Himanshu Kaul and Peter W. Zandstra  “GARMEN: GRN Agent-based Reaction-Diffusion Modeling Environment”, 2022  DOI: 10.25392/leicester.data.21375642.v1
  * [34] Sapna Chhabra et al.  “Dissecting the dynamics of signaling events in the BMP, WNT, and NODAL cascade during self-organized fate patterning in human gastruloids”  In _PLOS Biology_ 17.10 Public Library of Science (PLoS), 2019, pp. e3000498  DOI: 10.1371/journal.pbio.3000498
  * [35] Weiming Wang, Lei Zhang, Hailing Wang and Zhenqing Li  “Pattern formation of a predator–prey system with Ivlev-type functional response”  In _Ecological Modelling_ 221.2, 2010, pp. 131–140  DOI: https://doi.org/10.1016/j.ecolmodel.2009.09.011
  * [36] Shigeru Kondo, Masakatsu Watanabe and Seita Miyazawa  “Studies of Turing pattern formation in zebrafish skin”  In _Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences_ 379.2213 The Royal Society, 2021  DOI: 10.1098/rsta.2020.0274
  * [37] Timothy Dozat  “Incorporating Nesterov Momentum into Adam”, 2016
  * [38] Diederik P. Kingma and Jimmy Ba  “Adam: A Method for Stochastic Optimization”  arXiv, 2014  DOI: 10.48550/ARXIV.1412.6980
  * [39] Xingyu Xie et al.  “Adan: Adaptive Nesterov Momentum Algorithm for Faster Optimizing Deep Models”, 2023  arXiv:2208.06677 [cs.LG]
  * [40] “A.3 Normalized Gradient Descent — jermwatt.github.io” [Accessed 27-09-2023], https://jermwatt.github.io/machine_learning_refined/notes/3_First_order_methods/3_9_Normalized.html
  * [41] Sam Greydanus  “Studying Growth with Neural Cellular Automata”, {https://greydanus.github.io/2022/05/24/studying-growth/}
  * [42] Mainakdeb  “text-2-cellular-automata: Neural cellular automata + clip”  In _GitHub_ , https://github.com/Mainakdeb/text-2-cellular-automata
  * [43] Magnus Petersen  “Neural-Cellular-Automata-Image-Manipulation”  In _GitHub_ , https://github.com/MagnusPetersen/Neural-Cellular-Automata-Image-Manipulation
  * [44] Alec Radford et al.  “Learning Transferable Visual Models From Natural Language Supervision”, 2021  arXiv:2103.00020 [cs.CV]
  * [45] Alex Graves et al.  “Hybrid computing using a neural network with dynamic external memory”  In _Nature_ 538.7626 Springer ScienceBusiness Media LLC, 2016, pp. 471–476  DOI: 10.1038/nature20101
  * [46] W. Stark  “Amorphous computing: examples, mathematics and theory”  In _Natural Computing_ 12.3 Springer ScienceBusiness Media LLC, 2013, pp. 377–392  DOI: 10.1007/s11047-013-9370-0
  * [47] L Rensing, U Meyer-Grahle and P Ruoff  “Biological timing and the clock metaphor: Oscillatory and hourglass mechanisms”  In _Chronobiology International_ 18.3, 2001, pp. 329–369  DOI: 10.1081/CBI-100103961
  * [48] Steven L. Brunton, Joshua L. Proctor and J. Kutz  “Discovering governing equations from data by sparse identification of nonlinear dynamical systems”  In _Proceedings of the National Academy of Sciences_ 113.15 Proceedings of the National Academy of Sciences, 2016, pp. 3932–3937  DOI: 10.1073/pnas.1517384113
  * [49] Carl Folkestad et al.  “Extended Dynamic Mode Decomposition with Learned Koopman Eigenfunctions for Prediction and Control”, 2020  arXiv:1911.08751 [eess.SY]
  * [50] Antonio Matas-Gil and Robert G. Endres  “Unraveling biochemical spatial patterns: machine learning approaches to the inverse problem of Turing patterns”, 2023  arXiv:2309.06339 [physics.bio-ph]
  * [51] Davis Blalock, Jose Javier Gonzalez Ortiz, Jonathan Frankle and John Guttag  “What is the State of Neural Network Pruning?”, 2020  arXiv:2003.03033 [cs.LG]

##  Appendix A Gradient Calculation

To compute the gradients, we use the shorthand
$\varphi_{j}=x^{((\delta-1)t+j)}$, that is the NCA state propagated from time
$(\delta-1)t$ by $j$ time steps, which we can write step by step as

$\displaystyle\varphi_{0}$ $\displaystyle=x^{((\delta-1)t)}$
$\displaystyle\varphi_{1}$
$\displaystyle=\Phi(\varphi_{0},\theta)=x^{((\delta-1)t+1)}$
$\displaystyle\vdots$ $\displaystyle\varphi_{j}$
$\displaystyle=\Phi(\varphi_{j-1},\theta)=\Phi^{j}(\varphi_{0},\theta)$
$\displaystyle\vdots$ $\displaystyle\varphi_{t}$
$\displaystyle=\Phi(\varphi_{t-1},\theta)=\Phi^{t}(\varphi_{0},\theta)=x^{(\delta
t)}$

We have a distance $\mathcal{L}(x,y)$ between two states. We want to minimize
the total distance between the NCA prediction and the data averaged over all
$\delta$ in terms of the parameters $\theta$, so we look for

$\frac{d}{d\theta}\frac{1}{D}\sum_{\delta}\mathcal{L}(x^{(\delta
t)},y^{(\delta)})=\frac{1}{D}\sum_{\delta}\frac{\partial\mathcal{L}}{\partial
x^{(\delta t)}}\frac{dx^{(\delta t)}}{d\theta}$

since only the prediction $x^{(\delta t)}$ depends on the parameters $\theta$
and not the distance function or the data. The second term can be written as

$\frac{dx^{(\delta
t)}}{d\theta}=\frac{d\varphi_{t}}{d\theta}=\frac{d\Phi(\varphi_{t-1},\theta)}{d\theta}=\frac{\partial\Phi(\varphi_{t-1},\theta)}{\partial\varphi_{t-1}}\frac{d\varphi_{t-1}}{d\theta}+\frac{\partial\Phi(\varphi_{t-1},\theta)}{\partial\theta}$

We decreased the subscript $t$, which we can do in general

$\frac{d\varphi_{j}}{d\theta}=\frac{d\Phi(\varphi_{j-1},\theta)}{d\theta}=\frac{\partial\Phi(\varphi_{j-1},\theta)}{\partial\varphi_{j-1}}\frac{d\varphi_{j-1}}{d\theta}+\frac{\partial\Phi(\varphi_{j-1},\theta)}{\partial\theta}$

all the way down to $\varphi_{0}$. We get that

$\displaystyle\frac{dx^{(\delta
t)}}{d\theta}=\frac{d\varphi_{t}}{d\theta}=\frac{\partial\Phi(\varphi_{t-1},\theta)}{\partial\varphi_{t-1}}\Bigg{(}\frac{\partial\Phi(\varphi_{t-2},\theta)}{\partial\varphi_{t-2}}$
$\displaystyle\dots\Bigg{(}\frac{\partial\Phi(\varphi_{0},\theta)}{\partial\varphi_{0}}\frac{d\varphi_{0}}{d\theta}+\frac{\partial\Phi(\varphi_{0},\theta)}{\partial\theta}\Bigg{)}$
$\displaystyle+$
$\displaystyle\dots+\frac{\partial\Phi(\varphi_{t-2},\theta)}{\partial\theta}\Bigg{)}+\frac{\partial\Phi(\varphi_{t-1},\theta)}{\partial\theta}$

Here it is important to note that each partial gradient
$\frac{\partial\Phi(\varphi_{i},\theta)}{\partial\varphi_{i}}$ and
$\frac{\partial\Phi(\varphi_{i},\theta)}{\partial\theta}$ are matrices of size
$|\mathcal{X}|\times|\mathcal{X}|$ and $|\theta|\times|\mathcal{X}|$
respectively. Although in practice there will be nice factoring to simplify
the computation, the memory requirement is still dominated by
$|\mathcal{X}|^{2}$. This computation is hidden in line 17 of Algorithm 2 for
brevity.

##  Appendix B Videos

Legends for videos included in online supplementary information.

V1: A trajectory of an NCA with: 16 channels; identity, Sobel and Laplacian
kernels; relu activation, trained on the image morphing task. Annotated with
timestep $n$ in the top left corner

V2: A trajectory of the same NCA, but with hidden channels displayed as RGB in
the top right, bottom left and bottom right quadrants. Observable channels
displayed in the top left quadrant. Annotated with timestep $n$ in the top
left corner.

V3: Observable and hidden channels of an NCA trajectory, but this NCA has no
activation function (i.e. is linear). Annotated with timestep $n$ in the top
left corner.

V4: PDE trajectory corresponding to Figure 5. Annotated with timestep $n$ in
the top left corner.

V5: NCA trajectory corresponding to Figure 5. Annotated with timestep $n$ in
the top left corner.

V6: Difference between NCA and PDE trajectories in Figure 5. Annotated with
timestep $n$ in the top left corner.