# Latent Neural Cellular Automata for Resource-Efficient Image Restoration

Andrea Menta
DEIB, Politecnico di Milano
Via Ponzio 34/5, 20133 Milan, Italy
<EMAIL_ADDRESS>
&Alberto Archetti
DEIB, Politecnico di Milano
Via Ponzio 34/5, 20133 Milan, Italy
<EMAIL_ADDRESS>
&Matteo Matteucci
DEIB, Politecnico di Milano
Via Ponzio 34/5, 20133 Milan, Italy
<EMAIL_ADDRESS>

###### Abstract

Neural cellular automata represent an evolution of the traditional cellular
automata model, enhanced by the integration of a deep learning-based
transition function. This shift from a manual to a data-driven approach
significantly increases the adaptability of these models, enabling their
application in diverse domains, including content generation and artificial
life. However, their widespread application has been hampered by significant
computational requirements. In this work, we introduce the Latent Neural
Cellular Automata (LNCA) model, a novel architecture designed to address the
resource limitations of neural cellular automata. Our approach shifts the
computation from the conventional input space to a specially designed latent
space, relying on a pre-trained autoencoder. We apply our model in the context
of image restoration, which aims to reconstruct high-quality images from their
degraded versions. This modification not only reduces the model’s resource
consumption but also maintains a flexible framework suitable for various
applications. Our model achieves a significant reduction in computational
requirements while maintaining high reconstruction fidelity. This increase in
efficiency allows for inputs up to 16 times larger than current state-of-the-
art neural cellular automata models, using the same resources.

_K_ eywords neural cellular automata $\cdot$ deep learning $\cdot$ latent
space $\cdot$ image restoration $\cdot$ optimization

##  1 Introduction

Cellular Automata (CA) [1] are a computational model for simulating the
dynamics of complex systems from simple local interactions. CAs operate on a
regular lattice in which each cell is in one of a finite number of states. At
each discrete time step, cells are updated using simple local rules based on
the state of neighboring cells. CAs are capable of generating complex patterns
from simple update rules. The most famous example is John Conway’s Game of
Life, which is Turing complete [2]. In general, the potential to generate
intricate global patterns from simple local rules has always attracted
interest in CA research.

The advent of Deep Learning (DL) [3] has revolutionized various domains, such
as computer vision and natural language processing, establishing itself as a
State-of-the-Art (SotA) methodology. Its inherent flexibility and
generalization power have allowed several extensions of existing models,
including CAs. Neural Cellular Automata (NCA) [4] exemplify this progression,
by integrating DL into CAs. NCAs differ from CAs by using neural networks as
their transition functions. This aspect allows the transition from manually
designed rules to data-driven learning processes that target specific
behaviors. Although many promising applications have been proposed, ranging
from content generation [5] to control [6], NCAs face significant limitations
due to their resource requirements. The need for numerous concurrent updates
introduces significant latency and memory overhead [7], similar to Recurrent
Neural Networks (RNNs) [8], hindering their practical application.

To address these challenges, we introduce the Latent Neural Cellular Automata
(LNCA) model. The core of LNCA is to move computations from the conventional
input space to a customized, compressed latent space. The idea is to distill
the essential information for the NCA to solve the problem into a lower-
dimensional manifold, leaving the accessory information to a secondary
component. For this purpose, an Autoencoder (AE) is used to compress the input
and learn an embedding function. Our research focuses primarily on Image
Restoration (IR), tackling denoising and deblurring tasks. We evaluate LNCA
against SotA solutions such as Restormer [9] and NAFNet [10], which are based
on standard deep learning techniques, as well as the best NCA model, called
ViTCA [7]. Our evaluation includes 10 datasets, both real and synthetic,
measuring reconstruction performance and computational efficiency.

To the best of our knowledge, LNCA is the first model applying NCAs in the
latent space of a DL autoencoder backbone. Thus, the contributions of this
work are twofold:

  * •

We propose a novel architecture called LNCA and the corresponding learning
process to mitigate the inherent limitations of NCA models by exploiting the
expressiveness of DL autoencoders.

  * •

We apply the LNCA architecture to IR tasks, achieving significant reductions
in latency and GPU memory usage while maintaining competitive reconstruction
performance compared to SotA models.

To facilitate further research, we have made the source code of the
experiments publicly available.

##  2 Related Work

This section describes the main works related to (neural) cellular automata
and image restoration.

###  2.1 Cellular Automata

Cellular Automata (CA) are defined as a tuple
$\mathcal{M}=(\mathcal{L},\mathcal{S},\eta,\phi)$, where $\mathcal{L}$ is a
$d$-dimensional lattice, $\mathcal{S}$ is the set of possible states,
$\eta:\mathcal{L}\rightarrow 2^{\mathcal{L}}$ is the neighborhood function and
$\phi:\mathcal{S}^{n}\rightarrow\mathcal{S}$ is the transition function. Let
$s^{t}_{i}\in\mathcal{S}$ indicate the state of cell $i$ at time step $t$ and
$\Omega^{t}_{i}=\left\\{s^{t}_{j}:j\in\eta(i)\right\\}$ the set of neighbors
of cell $i$ according to $\eta$. Then, the state update for cell $i$ at each
iteration $t$ is computed as

$s_{i}^{t+1}=\phi\left(\left\\{s_{i}^{t}\cup\Omega_{i}^{t}\right\\}\right)\text{.}$
(1)

Neural Cellular Automata (NCAs) extend the CA framework by replacing the
transition function $\phi$ with a learnable neural network $\phi_{\theta}$
with parameters $\theta$, preserving the essential feature of locality.
However, training the neural network $\phi_{\theta}$ requires some
modifications to standard CAs. First, the state representations are
continuous. Second, optimization is performed after $n$ processing iterations,
and a boolean mask is used during each iteration to allow asynchronous
updates, which improves training dynamics, as in [4]. This training
methodology typically involves a two-step process with a replay buffer to
perform pool sampling. During each odd step, the neural network updates a
newly initialized tensor $x$, and the resulting tensor $x^{\prime}$ is stored
in the replay buffer. For even steps, instead, the network retrieves
$x^{\prime}$ from the replay buffer, updates it, and then stores the new
tensor $x^{\prime\prime}$ back in the buffer. This strategy ensures that the
target output acts as an attractor in the training process, effectively
preventing the neural network from deviating from the desired behavior.

###  2.2 Image Restoration

Image Restoration (IR) is a central task in digital image processing that aims
to reconstruct corrupted images into their original form. This process
involves reducing distortions such as noise and blur. Over time, a plethora of
methods have been developed to tackle IR, ranging from basic filtering
techniques [11] to advanced deep learning models, each improving the quality
of the restoration. In our analysis, we focus on three state-of-the-art deep
learning-based architectures, called Restormer, NAFNet, and ViTCA.

Restormer [9], leverages a residual learning approach based on the classic
UNet [12] architecture, powered by Transformer blocks [13]. Each block is
composed of a Multi-Dconv Head Transposed Attention (MDTA) and a Gated-Dconv
Feed-forward Network (GDFN). MDTA offers a memory-efficient alternative to the
self-attention block by utilizing channel dimension, while GDFN incorporates a
gating mechanism to the classic transformer head.

NAFNet [10] is a non-linear activation-free network, based on a UNet
structure. The basic block of this architecture includes two components, the
Simple Gate (SG) and the Simplified Channel Attention (SCA). SG introduces
non-linearities in the block, while SCA performs an attention operation
leveraging a channel-wise product. Both SG and SCA derive from the gated
linear unit [14].

Finally, ViTCA [7] is the first application of the attention mechanism to the
NCA family. This CA-based architecture implements a transition function
composed of four stages: embedding, Multi-Head Self-Attention (MHSA), Multi-
Layer Perceptron (MLP), and head. The embedding takes the seeded input
(optionally including the positional encoding) and feeds it through a linear
projection. The MHSA, the major novelty of ViTCA, estimates the relative
importance of each token. It implements a localized self-attention method to
respect NCA locality constraints, limiting the attention span to each cell’s
neighborhood. This way, even though information propagates locally, it
achieves a global behavior. The MLP, composed of two fully connected layers
with GELU activation, completes the classic Vision Transformer architecture
[15], while the head reshapes the output obtaining the final update vector.
Lastly, this model is trained using pool sampling and curriculum learning, a
typical approach for NCA architectures.

##  3 Latent Neural Cellular Automata

This paper introduces a novel approach to address the resource constraints
associated with NCAs, called Latent Neural Cellular Automaton (LNCA). The core
concept of our methodology is to move the NCA computation – local information
passing and state transition – from the input space to a specially designed
latent space. The main goal motivating this shift is to limit the
dimensionality of the space in which the NCA operates, extracting only the
task-specific information required and reducing redundant operations. In the
context of IR, our approach is particularly advantageous. We aim to build an
embedding space featuring a distinct separation between the image semantics
and the information about its corruption. This way, the autoencoder focuses on
constructing an embedding space structured to decouple image and corruption
information, while the NCA performs the restoration task. This task division
significantly enhances the overall efficiency of the system. In the following,
we present the LNCA architecture and its training pipeline within the IR task.

###  3.1 Architecture

Figure 1: LNCA structure. $+$ is an element-wise addition. $S$ is a channel
switch. On the left, the Encoder, which is composed of adjusted convolution
blocks, outputs the skip-connection $\bar{x}$ and the latent tensor $\hat{x}$.
In the middle, the NCA block processes the latent tensor to remove the
corruption, obtaining $\hat{x}^{\prime}$. On the right, the Decoder uses both
the skip-connection and the latent tensor to reconstruct the final output
$x^{\prime}$. The switch path around the NCA is used to bypass the NCA
computation during the AE training, as described in Section 3.2.1.

The LNCA architecture is composed of two modules, the Autoencoder (AE) and the
Neural Cellular Automaton (NCA). The AE’s encoder module projects input images
from the input space to the latent space, while the AE’s decoder performs the
inverse transformation. The NCA, instead, performs the actual image
restoration operating in the latent space. Figure 1 shows the LNCA
architecture. Specifically, the input image $x\in\mathbb{R}^{H\times W\times
C}$ is translated by the encoder $\mathcal{E}$ of the AE into its latent
representation $\hat{x}\in\mathbb{R}^{\hat{H}\times\hat{W}\times\hat{C}}$.
Then, the NCA $\mathcal{N}$ performs $n$ transition steps starting from the
initial state $\hat{x}$ and obtaining a new latent representation
$\hat{x}^{\prime}$ with the same dimensionality. Finally, the decoder block
$\mathcal{D}$ of the AE transforms the latent denoised image
$\hat{x}^{\prime}$ to the original image space, obtaining the desired output
$x^{\prime}\in\mathbb{R}^{H\times W\times C}$.

####  3.1.1 Autoencoder

The AE’s core objective is to create a lower dimensional manifold of the input
space containing task-related information for the NCA while preserving
reconstruction fidelity. For this reason, our architecture opts for a
classical AE model instead of more sophisticated variants, such as variational
autoencoders [16], which may introduce additional reconstruction uncertainty.
As shown in Figure 1, we adopted a standard AE architecture with some minor
tweaks to obtain better convergence during training. First, we included a
variant of the convolutional block, called adjusted convolution. This block
comprises an initial convolution with $3\times 3$ kernel and no activation,
followed by batch normalization, and swish activation function. The only
exception is the output layer which uses a sigmoid activation. The same is
applied to the transposed version of convolution used in the decoder. In
general, this block grants better convergence properties and stability for the
AE.

To help the latent space construction, we made two major additions to the AE
architecture, namely a skip-connection and an encoder side-channel. As skip-
connections aid in reconstruction quality, they funnel image semantics
unrelated to the corruption directly to the decoder. At the same time, all the
restoration information flows into the AE’s bottleneck for NCA processing.
This behavior is enforced through a composite loss function, detailed in
Section 3.2.1.

####  3.1.2 Neural Cellular Automata

The NCA is the component devoted to the actual restoration task. We opted for
integrating a ViTCA-based NCA as the NCA core of our architecture, as it
stands as the SotA model within the NCA family. Additionally, this enables a
direct comparison with the original, non-latent ViTCA and our novel model in
terms of reconstruction performance and computational complexity.

To further reduce computational requirements and evaluate the impact of the
attention mechanism for NCA architectures, we developed an alternative
attention-less NCA model. This model, called NAFCA, is inspired by NAFNet, the
architecture that replaced self-attention with a lighter gating mechanism,
adjusted to satisfy the CA locality constraint. This is achieved by
substituting the transformer-based transition function with a simpler
architecture. Specifically, this model comprises three residual blocks: the
perception block, the update block, and the head. The perception block
performs the information extraction and it is composed of an initial layer
normalization, followed by an activation-free $1\times 1$ convolution and a
Simple Gate [10]. The classic attention mechanism is then substituted by the
following operation:

$x^{\prime}=x\circ\mathcal{C}(x)\text{,}$ (2)

where $\mathcal{C}$ is a single-filter $3\times 3$ convolution respecting the
Moore neighborhood and $\circ$ is the Hadamard product. Then, another $1\times
1$ convolution and a dropout layer complete the perception block. The update
block takes care of the transition process. Its structure is similar to the
perception block, except for the lack of attention mechanism and the number of
filters in the convolutional layers. Finally, the head handles the state
update. It is composed of a layer normalization and a $1\times 1$ convolution.

###  3.2 Training

In the context of NCA, the training procedure must be carefully designed to
allow model convergence. We chose a training strategy consisting of two phases
to allow for more robust convergence. The first phase is the AE training,
where the AE learns to build a latent space suitable for the NCA while
minimizing the reconstruction error. The second phase is NCA training, where
the NCA learns to reconstruct the image within the latent space. In each
phase, the non-trained part of the architecture (i.e., first the NCA, then the
AE) is kept frozen. This two-phase approach allows a clear separation of tasks
between the AE and the NCA, driving the architecture blocks to the desired
behavior, and avoiding contamination.

####  3.2.1 Autoencoder Training

The first training phase involves the AE block only. Thus, the NCA processing
is bypassed with a direct connection between the encoder and the decoder. This
phase leverages batches $x$ composed of three types of samples:

  * •

Anchor samples ($x_{A}$), the ground truth restored images.

  * •

Positive samples ($x_{P}$), corrupted versions of the ground truth.

  * •

Negative samples ($x_{N}$), derangement (permutation where no element appears
in its original position) of the ground truth.

Given a sample batch, the AE training procedure is composed of two separate
gradient descent steps with dedicated losses to ensure more stable
convergence.

The first step includes the sum of three losses: the reconstruction loss, the
distance loss, and the task loss. The goal of the reconstruction loss is to
minimize the AE’s reconstruction error and it is evaluated as the mean squared
error between the AE’s input and output tensor:

$\mathcal{L}_{\text{REC}_{\text{AE}}}(x,x^{\prime})=\text{MSE}(x,x^{\prime})\text{.}$
(3)

The distance loss takes inspiration from the Triplet loss [17] and its goal is
to enforce the similarity in the latent space between images with the same
subjects and the dissimilarity between images with different subjects. Given a
tolerance margin $\alpha$, this loss is calculated as

$\mathcal{L}_{\text{DIST}}(\hat{x})=\max\\{0,\text{MSE}(\hat{x}_{A},\hat{x}_{P})-\text{MSE}(\hat{x}_{A},\hat{x}_{N})+\alpha\\}\text{.}$
(4)

Finally, the task loss is used to obtain a faithful reconstruction of the
input corruption. To enforce this property, this loss takes into account only
the corrupted pixels. Specifically, it applies a boolean mask before computing
the mean squared error of the corrupted samples:

$\begin{split}\mathcal{L}_{\text{TASK}}(x,x^{\prime})&=\text{MSE}(x_{P}\circ
m_{A,P},x_{P}^{\prime}\circ m_{A,P})\\\ m_{A,P}&=\mathbf{1}(x_{A}-x_{P}\neq
0)\text{.}\end{split}$ (5)

Each of these losses is then weighted and summed to obtain the loss of the
first gradient descent step. The weights are determined through validation
over a finite set of values using grid-search. This tuning procedure ensures a
balance between reconstruction quality and latent space structure and helps
manage differences in loss magnitudes.

The second gradient descent step, instead, involves the equivalence loss,
inspired by the Swapping AE [18] design. The goal of this loss is to split the
information related to the image semantics and the image corruption between
the two encoder outputs. In particular, the skip-connection $\bar{x}$ must
help in the reconstruction of the image semantics, while the corruption
profile is encoded only in the latent tensor $\hat{x}$. This is obtained by
swapping $\hat{x}_{P}$ and $\hat{x}_{A}$, the latent tensors corresponding to
the corrupted images and the ground truth images, respectively. In fact, by
substituting $\hat{x}_{P}$ with its clean counterpart $\hat{x}_{A}$, while
keeping the corrupted skip-connection $\bar{x}_{P}$, we should obtain a clean
output $x_{A}$ and vice versa. This approach has the advantage of introducing
a specific behavior in the AE, which is essential for the subsequent training
procedure of the NCA. We enforce this property by minimizing the relative mean
squared error:

$\begin{split}\mathcal{L}_{\text{EQ}}(x,\bar{x},\hat{x})=&\,\text{MSE}(x_{A}\circ
m_{A,P},x^{\star}_{A}\circ m_{A,P})+\\\ &\,\text{MSE}(x_{P}\circ
m_{A,P},x^{\star}_{P}\circ m_{A,P})\\\
x_{A}^{\star}=&\,\mathcal{D}(\hat{x}_{A}+\epsilon,\bar{x}_{P})\\\
x_{P}^{\star}=&\,\mathcal{D}(\hat{x}_{P},\bar{x}_{A})\text{,}\end{split}$ (6)

where $\epsilon\sim N(0,v)$ is a smoothing factor that allows for a small
tolerance in the reconstruction.

####  3.2.2 Neural Cellular Automata Training

In this phase, we focus on the restoration process. Each input batch $x$ is
composed only of corrupted images and it is associated with a clean ground
truth $y$. Furthermore, the generated latent tensor $\hat{x}$ is processed by
the NCA module, obtaining $\hat{x}^{\prime}$.

Similar to the AE training, this procedure leverages multiple losses to learn
the IR task: the reconstruction loss, the latent loss, and the overflow loss.
Each loss has an associated weight, which is tuned similarly to the AE case.
First, the reconstruction loss minimizes the reconstruction error of the
overall architecture in the image space. It is computed as the mean squared
error between the predicted output $x^{\prime}$ and the ground truth $y$:

$\mathcal{L}_{\text{REC}_{\text{NCA}}}(y,x^{\prime})=\text{MSE}(y,x^{\prime})\text{.}$
(7)

The latent loss plays a similar role to
$\mathcal{L}_{\text{REC}_{\text{NCA}}}$ but in the latent space. It enforces
the similarity between the latent tensor of the ground truth, $\hat{y}$, and
the one reconstructed by the NCA processing, $\hat{x}^{\prime}$, starting from
the associated corrupted input:

$\mathcal{L}_{\text{LAT}}(\hat{y},\hat{x}^{\prime})=\text{MSE}(\hat{y},\hat{x}^{\prime})\text{.}$
(8)

Finally, the overflow loss regularizes the output magnitude of the NCA using
an L1 norm. As in [7], it is applied to both the output and the hidden
channels of the NCA output tensor $\hat{x}^{\prime}$:

$\mathcal{L}_{\text{OVER}}(\hat{x}^{\prime})=\frac{1}{C}\big{\|}\hat{x}^{\prime}-\text{clip}(\hat{x}^{\prime})\big{\|}_{1}\text{,}$
(9)

where $C$ is the number of channels considered. Output channels are clipped in
the interval $[0,1]$, while hidden channels are in the interval $[-1,1]$.

##  4 Experiments and Results

In this section, we evaluate the LNCA architecture in the IR task against
several SotA models. In addition to the restoration performance, the analyses
focus on time and memory resource efficiency, to assess the applicability of
LNCA compared with existing NCA solutions.

###  4.1 Datasets

This section collects the information related to the datasets involved in our
experiments. To allow for a comprehensive performance evaluation, we adopted
both synthetic and real-world datasets. Concerning synthetic datasets, we
started from CelebA [19], CIFAR-10 [20], and TinyImageNet [21] and applied
specific corruption patterns. For denoising, we applied an additive Gaussian
noise in a channel-coherent fashion. This procedure generates tone shifts
rather than completely out-of-context pixel values, making the task harder.
For deblurring, we use a camera motion blur. It generates a defocus along a
random direction coherent for all the elements in the image.

For the real datasets, instead, we selected some of the most widely
acknowledged in the image restoration field, namely Renoir [22] and SID [23]
for denoising, and GoPro [24] and RealBlur [25] for deblurring.

Synthetic datasets grant fine-grained control over the type of corruption
applied to the images, while real-world data allows for performance assessment
in realistic scenarios. Also, synthetic and real datasets differ in their
cardinality. On the one side, synthetic datasets contain many low-resolution
images, while real datasets are composed of a few high-resolution images.
Therefore, we need to apply different preprocessing to use these datasets
properly. Specifically, for all the real datasets, we applied a tile-shrink-
clean procedure, which converts the original image into a $n\times m$
patchwork, resizes each tile to $32\times 32$ pixels, and then performs a
cleaning step to remove redundant images. Conversely, for the synthetic
datasets, we directly resized each image to the desired shape.

###  4.2 Training Procedure

In the following, we describe the experimental procedure adopted to collect
the restoration performance and the time and memory requirements of LNCA and
the baseline models from the literature. Specifically, we decided to use
Restormer [9] and NAFNet [10] as representatives of classic deep learning
architectures, and ViTCA [7] as the NCA flagship model. On our side, as stated
in Section 3.1.2, we propose two versions of the LNCA architecture, with a
different NCA structure. The first version, referred to as LatentViTCA, adopts
ViTCA as the NCA module. Conversely, the second version, called LatentNAFCA,
has an attention-free NCA core, leading to a more lightweight architecture.

To keep the tests as fair as possible, we kept consistent training procedures
across all the models. Specifically, we applied an 80-20% split to obtain
training and test sets followed by another 80-20% split on the training set to
extract the validation set. When training on synthetic datasets, we employ
curriculum learning in the training set, feeding images of increasing
corruption, while validation and test sets use a fixed difficulty. For real
datasets, instead, we do not apply any curriculum learning.

Each model has been trained for $20$ epochs, reaching a performance plateau.
We opted for using Adam as an optimizer with cosine annealing scheduling [26]
for the learning rate. Furthermore, regarding CA-based architectures, we used
a pool size of $1024$ and a hidden-channel size of $32$, as in [7]. During
training, the number of CA iterations was uniformly sampled between $8$ and
$32$, while for validation and testing, we fixed it to $64$.

###  4.3 Results on Restoration Performance

This section comments on the results obtained by LNCA and the SotA IR
counterparts concerning pure restoration performance. The goal of our proposed
latent models is to improve the computational and memory efficiency of NCA
architectures at the expense of task utility. Therefore, we expect latent
architectures to perform slightly worse than the alternatives in pure
restoration performance.

To assess performance, we adopt the reference metric in IR, namely the
Structural Similarity Index Measure [27] (SSIM). SSIM is a perception-based
metric incorporating the concepts of luminance ($l$), contrast ($c$), and
structure ($s$), which are estimated as follows:

$\begin{split}l(x,y)=\frac{2\mu_{x}\mu_{y}+c_{1}}{\mu_{x}^{2}+\mu_{y}^{2}+c_{1}}\\\
c(x,y)=\frac{2\sigma_{x}\sigma_{y}+c_{2}}{\sigma_{x}^{2}+\sigma_{y}^{2}+c_{2}}\\\
s(x,y)=\frac{2\sigma_{xy}+c_{3}}{\sigma_{x}\sigma_{y}+c_{3}}\text{,}\end{split}$
(10)

where $\mu_{x}$ and $\mu_{y}$ are the sample means, $\sigma_{x}$ and
$\sigma_{y}$ the standard deviations, $\sigma_{xy}$ the cross-correlation,
$c_{1}$, $c_{2}$, and $c_{3}$ stabilization coefficients. These terms are then
weighted (using $\alpha$, $\beta$ and $\gamma$ exponents) and combined to
obtain the final formulation:

$SSIM(x,y)=l(x,y)^{\alpha}c(x,y)^{\beta}s(x,y)^{\gamma}\text{.}$ (11)

This metric leverages a kernel to extract the structural information of
images, rather than focusing on single-pixel errors. This leads to a more
robust metric compared to other alternatives like MSE and PSNR.

####  4.3.1 Denoising

Table 1: Denoising results on both synthetic (CelebA, CIFAR-10, and
TinyImageNet) and real (Renoir and SID) datasets, reporting the SSIM obtained
by each of the analyzed models. The best results are highlighted in bold,
while the second best are underlined.

Models | CelebA | CIFAR-10 | TinyImageNet | Renoir | SID
---|---|---|---|---|---
NAFNet | 0.958 | 0.924 | 0.837 | 0.962 | 0.752
Restormer | 0.967 | 0.918 | 0.880 | 0.965 | 0.781
ViTCA | 0.942 | 0.910 | 0.852 | 0.939 | 0.647
LatentViTCA | 0.889 | 0.855 | 0.802 | 0.877 | 0.646
LatentNAFCA | 0.837 | 0.833 | 0.776 | 0.880 | 0.652

The first set of experiments focuses on the image-denoising task. Examining
the results presented in Table 1, NAFNet and Restormer achieve the highest
denoising performance on all the tested datasets. Nevertheless, ViTCA obtains
comparable results. As expected, latent models, namely LatentViTCA and
LatentNAFCA, exhibit a decrease in performance due to the downsizing of the
input. Upon close examination, we notice a decrease of $\approx 5\%$ in
denoising efficiency for LatentViTCA, with respect to ViTCA, and a slightly
larger decrease of $\approx 7\%$ for LatentNAFCA. We emphasize that a task
utility downgrade is expected for latent architectures, as they are designed
to trade restoration performance for computational and memory efficiency. In
Section 4.4 we comment on this tradeoff by analyzing the resources required by
latent models.

####  4.3.2 Deblurring

Table 2: Deblurring results on both synthetic (CelebA, CIFAR-10, and
TinyImageNet) and real (GoPro and RealBlur) datasets, reporting the SSIM
obtained by each of the analyzed models. The best results are highlighted in
bold, while the second best are underlined.

Models | CelebA | CIFAR-10 | TinyImageNet | GoPro | RealBlur
---|---|---|---|---|---
NAFNet | 0.986 | 0.970 | 0.962 | 0.908 | 0.869
Restormer | 0.993 | 0.977 | 0.975 | 0.934 | 0.918
ViTCA | 0.941 | 0.903 | 0.870 | 0.773 | 0.814
LatentViTCA | 0.769 | 0.752 | 0.672 | 0.739 | 0.772
LatentNAFCA | 0.745 | 0.739 | 0.673 | 0.739 | 0.771

Results on the deblurring task, follow a similar trend to denoising, as shown
in Table 2. In particular, NAFNet and Restormer still achieve the best
deblurring results. ViTCA follows as the third best alternative, although the
performance gap is higher than the denoising task. Compared to ViTCA, latent
models exhibit a greater decrease in performance of $\approx 13\%$ for
LatentViTCA and $\approx 15\%$ for LatentNAFCA. This gap is less prominent in
real-world datasets – GoPro and RealBlur – which are affected by less critical
corruptions than synthetic datasets.

###  4.4 Efficiency Results

This section focuses on the efficiency analysis of the proposed solution,
LNCA, compared with the existing IR architectures. We measure and compare
three metrics: training memory usage, training latency, and inference latency.
These metrics summarize the resource usage and computational requirements of a
generic deep learning architecture. The objective of LNCA is to reduce these
metrics with respect to other NCA-based models. In this context, our primary
comparison is with ViTCA, which belongs to the NCA family. However, we will
also extend our tests to NAFNet and Restormer to ensure a more comprehensive
understanding of the results.

To enable testing on a wider range of input sizes, the following experiments
were performed on a cloud service using an Nvidia A100 GPU with $80$GB of
VRAM, which has been the main limiting factor for NCA architectures.

####  4.4.1 Training memory requirement

Table 3: Training memory requirements tests. Results are reported once the
memory requirements have stabilized after an initial warm-up phase, measured
in GB. If the configuration requires more than the available $\mathbf{80}$ GB
of VRAM, the measurement is not available and is indicated by "–". The best
results are highlighted in bold, while the second best are underlined.

Models | $32\times 32$ | $64\times 64$ | $128\times 128$
---|---|---|---
$8$ | $16$ | $32$ | $8$ | $16$ | $32$ | $8$ | $16$ | $32$
NAFNet | 4.1 | 4.1 | 4.2 | 4.2 | 8.2 | 8.8 | 8.7 | 19.6 | 35.0
Restormer | 4.4 | 4.4 | 8.5 | 8.8 | 16.8 | 32.9 | 32.8 | 65.6 | –
ViTCA | 8.7 | 16.9 | 33.8 | 33.8 | 65.6 | – | – | – | –
LatentViTCA | 0.5 | 1.1 | 1.2 | 2.2 | 4.4 | 4.7 | 8.7 | 16.9 | 17.4
LatentNAFCA | 0.5 | 0.5 | 1.1 | 1.1 | 2.1 | 4.1 | 4.1 | 8.2 | 16.4

Figure 2: Training memory trend of the tested models across the different
configurations.

This metric (Table 3 and Figure 2) is calculated as the highest VRAM
requirement for each model at a given input resolution, typically reached
after an initial warm-up period and then held constant throughout training.

Starting from the lowest image resolution tested, namely $32\times 32\times 3$
with a batch size of 8, we notice a reduction of $\approx 94\%$ for both our
solutions compared to ViTCA. This reduction remains nearly constant across all
input configurations. Furthermore, considering the minimum memory requirement
of ViTCA, i.e., $8.7$GB, we achieve a $16$x improvement in the training input
size by using its latent counterpart. Interestingly, the latent models are
also up to $\approx 88\%$ more efficient than the classical solutions, which
should have a natural advantage due to their non-recurrent structure. Overall,
we can see that our solutions (and NAFNet) cope well with the increase in
batch size, keeping the requirements almost stable, while the other models
turn out to be less scalable.

####  4.4.2 Training latency

In this experiment, we adopt a dummy dataset of $1000$ noisy samples ($80\%$
training, $20\%$ validation) and perform a $10$-epoch training. We compute the
average training latency discarding the values of the initial warm-up epoch.
For latent models, we sum the AE and NCA training times to get a fair estimate
of total training time.

As expected, with a fixed-size dataset, increasing the batch size leads to an
improvement in the overall training time for all models (Table 4 and Figure
3). At the lowest resolution ($32\times 32\times 3$), the difference between
ViTCA and the latent models is negligible, although the decreasing slope is
less pronounced in the former. Increasing the resolution, we notice the
degradation in performance of ViTCA and, subsequently, of Restormer, while the
remaining models keep stable timings. In our view, the point at which
performance begins to decline is closely tied to the memory usage of the
models. Specifically, processing latency deteriorates as the model starts to
saturate the GPU VRAM. Overall, latent models achieve up to $\approx 80\%$
reduction in training latency when compared to ViTCA, making them comparable
to the NAFNet.

Table 4: Training latency tests performed with on a dummy dataset of
$\mathbf{1000}$ samples ($\mathbf{80-20}$ split). The results are the average
of $\mathbf{10}$ training epochs (after the initial warm-up), measured in
seconds. If the configuration requires more than the available $\mathbf{80}$
GB of VRAM, the measurement is not available and is indicated by "–". The best
results are highlighted in bold, while the second best are underlined.

Models | $32\times 32$ | $64\times 64$ | $128\times 128$
---|---|---|---
$8$ | $16$ | $32$ | $8$ | $16$ | $32$ | $8$ | $16$ | $32$
NAFNet | 25.1 | 13.3 | 7.9 | 22.8 | 13.0 | 9.0 | 28.8 | 19.0 | 15.7
Restormer | 21.7 | 11.9 | 7.9 | 25.2 | 18.9 | 16.8 | 59.9 | 57.1 | –
ViTCA | 21.7 | 18.5 | 17.5 | 69.9 | 61.6 | – | – | – | –
LatentViTCA | 14.9 | 8.9 | 4.9 | 16.6 | 11.8 | 8.4 | 30.0 | 21.8 | 16.4
LatentNAFCA | 24.5 | 14.5 | 7.2 | 17.1 | 11.4 | 7.4 | 21.2 | 14.8 | 10.2

Figure 3: Training latency trend of the tested models across the different
configurations.

####  4.4.3 Inference latency

Regarding the inference latency, tests are performed using a single batch as
input and averaging over $10$ iterations. Similarly to the training latency,
we discard the first processing of each model as a warm-up.

We expect our solutions and ViTCA to be significantly slower than classical
models due to the NCA computation. In fact, we fix these models to perform
$t=64$ NCA processing steps instead of the single pass used by the non-NCA
models. Looking at the results (Table 5 and Figure 4), we can see that our
solutions and NAFNet maintain almost stable performance across all
resolutions, while the other models diverge once a memory breakpoint is
reached. As mentioned in the previous section, we believe that this trend is
directly related to VRAM usage. At the highest resolution ($128\times
128\times 3$), where memory requirements are demanding, our latent models are
respectively up to $\approx 72\%$ and $\approx 89\%$ faster than ViTCA, with
LatentNAFCA also $\approx 37\%$ faster than Restormer.

Table 5: Inference latency tests performed on a single batch. The results are
the average of $\mathbf{10}$ inference steps (discarding the initial warm-up),
measured in seconds. The best results are highlighted in bold, while the
second best are underlined.

Models | $32\times 32$ | $64\times 64$ | $128\times 128$
---|---|---|---
$8$ | $16$ | $32$ | $8$ | $16$ | $32$ | $8$ | $16$ | $32$
NAFNet | 0.25 | 0.24 | 0.25 | 0.24 | 0.24 | 0.24 | 0.24 | 0.25 | 0.25
Restormer | 0.56 | 0.53 | 0.54 | 0.54 | 0.55 | 0.58 | 0.58 | 0.68 | 0.93
ViTCA | 1.47 | 1.42 | 1.45 | 1.43 | 1.46 | 2.00 | 2.01 | 2.96 | 5.17
LatentViTCA | 1.45 | 1.45 | 1.46 | 1.43 | 1.43 | 1.41 | 1.43 | 1.44 | 1.42
LatentNAFCA | 0.59 | 0.60 | 0.60 | 0.59 | 0.60 | 0.59 | 0.59 | 0.60 | 0.59

Figure 4: Inference latency trend of the tested models across the different
configurations.

##  5 Conclusions

NCAs are fascinating models capable of learning complex system behavior from
simple local interactions. Yet, their practical application has been severely
limited by their computational constraints. This work focused on mitigating
this critical challenge, providing an NCA model, called LNCA, with low memory
requirements and latency bottlenecks. LNCA is a novel deep learning-based
architecture that moves the computationally demanding NCA computation from the
input space to a lower-dimensional manifold. This manifold is tailored by an
autoencoder to extract the essential information for the NCA to solve the
task. We validated LNCA in the field of image restoration, specifically
addressing denoising and deblurring tasks. The experiments showcased an
effective reduction in the computational requirements of NCA. With respect to
the SotA, LNCA achieved up to a $16\times$ improvement in the maximum input
size, and up to a $\approx 94\%$ and $\approx 80\%$ reductions in memory
requirements and processing latency, respectively, with the drawback of an
overall degradation in restoration performance of $\approx 10\%$. In
conclusion, LNCA successfully addressed some of the inherent limitations of
NCA models, providing a flexible and practical solution for the integration of
cellular automata into deep learning architectures.

## References

  * [1] J. von Neumann.  The general and logical theory of automata.  In L. A. Jeffress, editor, Cerebral Mechanisms in Behaviour. Wiley, 1951.
  * [2] Paul Rendell.  Turing Universality of the Game of Life, pages 513–539.  Springer London, London, 2002.
  * [3] Ian J. Goodfellow, Yoshua Bengio, and Aaron Courville.  Deep Learning.  MIT Press, Cambridge, MA, USA, 2016.  http://www.deeplearningbook.org.
  * [4] Alexander Mordvintsev, Ettore Randazzo, Eyvind Niklasson, and Michael Levin.  Growing neural cellular automata.  Distill, 2020.  https://distill.pub/2020/growing-ca.
  * [5] Shyam Sudhakaran, Djordje Grbic, Siyan Li, Adam Katona, Elias Najarro, Claire Glanois, and Sebastian Risi.  Growing 3d artefacts and functional machines with neural cellular automata, 2021.
  * [6] Alexandre Variengien, Stefano Nichele, Tom Glover, and Sidney Pontes-Filho.  Towards self-organized control: Using neural cellular automata to robustly control a cart-pole agent, 2021.
  * [7] Mattie Tesfaldet, Derek Nowrouzezahrai, and Chris Pal.  Attention-based neural cellular automata.  Advances in Neural Information Processing Systems, 35:8174–8186, 2022.
  * [8] David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams.  Learning representations by back-propagating errors.  nature, 323(6088):533–536, 1986.
  * [9] Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, and Ming-Hsuan Yang.  Restormer: Efficient transformer for high-resolution image restoration.  In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 5728–5739, 2022.
  * [10] Liangyu Chen, Xiaojie Chu, Xiangyu Zhang, and Jian Sun.  Simple baselines for image restoration.  In European Conference on Computer Vision, pages 17–33. Springer, 2022.
  * [11] M.R. Banham and A.K. Katsaggelos.  Digital image restoration.  IEEE Signal Processing Magazine, 14(2):24–41, 1997.
  * [12] Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In Nassir Navab, Joachim Hornegger, William M. Wells, and Alejandro F. Frangi, editors, Medical Image Computing and Computer-Assisted Intervention – MICCAI 2015, pages 234–241, Cham, 2015. Springer International Publishing.
  * [13] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [14] Yann N. Dauphin, Angela Fan, Michael Auli, and David Grangier.  Language modeling with gated convolutional networks.  In Proceedings of the 34th International Conference on Machine Learning - Volume 70, ICML’17, page 933–941. JMLR.org, 2017.
  * [15] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  In International Conference on Learning Representations, 2021.
  * [16] Diederik P. Kingma and Max Welling.  Auto-Encoding Variational Bayes.  In 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings, 2014.
  * [17] Florian Schroff, Dmitry Kalenichenko, and James Philbin.  FaceNet: A unified embedding for face recognition and clustering.  In 2015 IEEE Conference on Computer Vision and Pattern Recognition (CVPR). IEEE, jun 2015.
  * [18] Taesung Park, Jun-Yan Zhu, Oliver Wang, Jingwan Lu, Eli Shechtman, Alexei Efros, and Richard Zhang.  Swapping autoencoder for deep image manipulation.  Advances in Neural Information Processing Systems, 33:7198–7211, 2020.
  * [19] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang.  Deep learning face attributes in the wild.  In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
  * [20] Alex Krizhevsky and Geoffrey Hinton.  Learning multiple layers of features from tiny images.  Technical Report 0, University of Toronto, Toronto, Ontario, 2009.
  * [21] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei.  ImageNet Large Scale Visual Recognition Challenge.  International Journal of Computer Vision (IJCV), 115(3):211–252, 2015.
  * [22] Josue Anaya and Adrian Barbu.  RENOIR – a dataset for real low-light image noise reduction.  Journal of Visual Communication and Image Representation, 51:144–154, feb 2018.
  * [23] Chen Chen, Qifeng Chen, Jia Xu, and Vladlen Koltun.  Learning to see in the dark.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2018.
  * [24] Seungjun Nah, Tae Hyun Kim, and Kyoung Mu Lee.  Deep multi-scale convolutional neural network for dynamic scene deblurring.  In CVPR, July 2017.
  * [25] Jaesung Rim, Haeyun Lee, Jucheol Won, and Sunghyun Cho.  Real-world blur dataset for learning and benchmarking deblurring algorithms.  In Andrea Vedaldi, Horst Bischof, Thomas Brox, and Jan-Michael Frahm, editors, Computer Vision – ECCV 2020, pages 184–201, Cham, 2020. Springer International Publishing.
  * [26] Ilya Loshchilov and Frank Hutter.  Sgdr: Stochastic gradient descent with warm restarts.  In ICLR (Poster). OpenReview.net, 2017.
  * [27] Zhou Wang, A.C. Bovik, H.R. Sheikh, and E.P. Simoncelli.  Image quality assessment: from error visibility to structural similarity.  IEEE Transactions on Image Processing, 13(4):600–612, 2004.