Attention-based Neural Cellular Automata
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Recent extensions of Cellular Automata (CA) have incorporated key ideas frommodern deep learning, dramatically extending their capabilities and catalyzinga new family of Neural Cellular Automata (NCA) techniques.
Inspired by
Transformer-based architectures
, our work presents a new class of attention-based
NCAs formed using a spatially localized—yet globally organized—self-attentionscheme.
We introduce an instance of this class named Vision Transformer Cellular
Automata (ViTCA) .
We present quantitative and qualitative results on denoisingautoencoding across six benchmark datasets, comparing ViTCA to a U-Net, a
U-Net-based CA baseline (UNetCA), and a Vision Transformer (ViT).
When comparing across architectures conﬁgured to similar parameter complexity, ViTCAarchitectures yield superior performance across all benchmarks and for nearlyevery evaluation metric.
We present an ablation study on various architecturalconﬁgurations of ViTCA, an analysis of its effect on cell states, and an investigationon its inductive biases.
Finally, we examine its learned representations via linearprobes on its converged cell state hidden representations, yielding, on average,superior results when compared to our U-Net, ViT, and UNetCA baselines.
Introduction
Recent developments at the intersection of twofoundational ideas—Artiﬁcial Neural Networks(ANNs) and Cellular Automata (CA)—have ledto new approaches for constructing Neural Cellular Automata (NCA).
These advances haveintegrated ideas such as variational inference[7], U-Nets [26] and Graph Neural Networks(GNNs) [15] with promising results on problems ranging from image synthesis [7, 20, 21]to
Reinforcement Learning (RL) [6, 22].
Transformers are another signiﬁcant development indeep learning [41], but, until now, have not beenexamined under an NCA setting.
Vision Transformers (ViTs) [13] have emergedas a competitive alternative to Convolutional
Neural Network (CNN) [56] architectures forcomputer vision, such as Residual Networks(ResNets) [45].
ViTs leverage the self-attentionmechanisms of original Transformers [41], which have emerged as the dominant approach forsequence modelling in recent years.
Our work combines foundational ideas from Transformers and
ViTs, leading to a new class of NCAs: Vision Transformer Cellular Automata (ViTCA) .
An effective and ubiquitous Transformer-based learning technique for Natural Language Processing(NLP) pre-training is the unsupervised task of Masked Language Modelling (MLM), popularizedby the BERT language model [34].
The success of MLM-based techniques has similarly inspiredrecent work re-examining the classical formulation of Denoising Autoencoders (DAEs) [51], butfor ViTs [3, 13, 28], introducing tasks such as Masked Image Encoding [16] and Masked Feature
Prediction [24] for image and video modelling, respectively.
This simple yet highly-scalable strategyof masked-based unsupervised pre-training has yielded promising transfer learning results on visionbased downstream tasks such as object detection and segmentation, image classiﬁcation, and actiondetection, even outperforming supervised pre-training [16, 24].
We examine training methodologiesfor ViTCA within a DAE setting and perform extensive controlled experiments benchmarking theseformulations against modern state of the art architectures, with favourable outcomes, e.g ., Fig. 1.
Our contributions are as follows: ﬁrst —to the best of our knowledge—our work is the ﬁrst toextend NCA methodologies with key Transformer mechanisms, i.e ., self-attention and positionalencoding (and embedding), with the beneﬁcial side-effect of circumventing the quadratic complexityof self-attention; second , our ViTCA formulation allows for lower model complexity (by limiting
ViT depth) while retaining expressivity through CA iterations on a controlled state—all with thesame encoder weights.
This yields a demonstrably more parameter-efﬁcient [20] ViT-based model.
Importantly, ViTCA mitigates the problems associated with the explicit tuning of ViT depth originallyneeded to improve performance ( i.e ., we use a depth of 1).
With ViTCA, we simply iterate until cellstate convergence.
Since ViT (and by extension, ViTCA) employs Layer Normalization (LN)
[43]at each stage of its processing, it is a fairly contractive model capable of ﬁxed-point convergenceguarantees [32].
In relation to our ﬁrst contribution, ViTCA respects CA requirements, most importantly that computations remain localized about a cell and its neighbourhood.
As such, we modify the globalself-attention mechanisms of a ViT to respect this locality requirement (Fig. 2).
Localized selfattention is not a new idea [4, 12, 19, 27]; however, because cells contain state information thatdepends on its previous state, over CA iterations the effective receptive ﬁeld of ViTCA’s localizedself-attention grows increasingly larger until eventually incorporating information implicitly acrossall cells.
Thus, admitting global propagation of information from spatially localized self-attention.
Moreover, due to the self-organizing nature of NCAs, self-organization also manifests itself withinthe localized self-attention, resulting in a globally agreed-upon arrangement of local self-attention.
Thus, circumventing the quadratic complexity of explicit global self-attention (w.r.t.
the input size) through a linear amortization over time, and increasing the feasibility of per-pixel dense processing(as we demonstrate).
This globally consistent and complex behaviour, which arises from strictly localinteractions, is a unique feature of NCAs and confers performance beneﬁts which we observe bothqualitatively and quantitatively when comparing ViT and ViTCA for denoising autoencoding.
Background and related work
Neural Cellular Automata.
Cellular Automata are algorithmic processes motivated by the biological behaviours of cellular growth and, as such, are capable of producing complex emergent(global) dynamics from the iterative application of comparatively simple (localized) rules [60].
Neural Cellular Automata present a more general CA formulation, where the evolving cell states arerepresented as (typically low-dimensional) vectors and the update rule dictating their evolution is adifferentiable function whose parameters are learned through backpropagation from a loss, ratherthan a handcrafted set of rules [30, 35, 59].
Neural net-based formulations of CAs in the NeurIPScommunity can be traced back to the early work of [59], where only small and simple modelswere examined.
Recent formulations of NCAs have shown that when leveraging the power of deeplearning techniques enabled by advances in hardware capabilities—namely highly-parallelizabledifferentiable operations implemented on GPUs—NCAs can be tuned to learn surprisingly complexdesired behaviour, such as semantic segmentation [31]; common RL tasks such as cart-pole balancing[22], 3D locomotion [6], and Atari game playing [6]; and image synthesis [7, 20, 21].
Althoughthese recent formulations rely on familiar compositions of convolutions and non-linear functions, it isimportant to highlight that NCAs are fundamentally not equivalent to “very-deep” CNNs ( vs .
[35]),or any other feedforward architecture ( e.g ., ResNets [45]), particularly, in the same way that a Recurrent Neural Network (RNN) is not equivalent: CNNs and other feedforward architectures inducea directed acyclic computation graph ( i.e ., a ﬁnite impulse response),
whereas NCAs (and RNNs)induce a directed cyclic computation graph ( i.e ., an inﬁnite impulse response), where stateful datacan additionally be manipulated using (learned) feedback loops and/or time-delayed controls.
Assuch, NCAs can be viewed as a type of RNN, and both (N)CAs and RNNs are known to be Turingcomplete [11, 54, 57, 59].
Vision Transformers.
Vision Transformers [13] are an adaptation of Transformers [41] to visionbased tasks like image classiﬁcation.
In contrast to networks built from convolutional layers, ViTs relyon self-attention mechanisms operating on tokenized inputs.
Speciﬁcally, input images are dividedinto non-overlapping patches, then fed to a Transformer after undergoing a linear patch projectionwith an embedding matrix.
While ViTs provide competitive image classiﬁcation performance, the quadratic computational scaling of global self-attention limits their applicability in high-dimensionaldomains, e.g ., per-pixel dense processing.
Recent developments have attempted to alleviate suchefﬁciency limitations [9, 10, 14, 17], one notable example being Perceiver IO [5, 8] with its use ofcross-attention.
We refer interested readers to a comprehensive survey on ViTs [18].
Vision Transformer Cellular Automata (ViTCA)
Building upon NCAs and ViTs, we propose a new class of attention-based NCAs formed using aspatially localized—yet globally organized—self-attention scheme.
We detail an instance of this class,
ViTCA, by ﬁrst reviewing its backbone ViT architecture before describing the “pool sampling”-basedtraining process for the ViTCA update rule (see overview in Fig.
Input tokenization.
ViT starts by dividing a C i ⇥ H ⇥ W input image X into N non-overlapping 112 P
H ⇥ P W patches ( 16 ⇥ 16 in the original work [13]), followed by a linear projection of the ﬂattenedimage patches with an embedding matrix E 2 R L ⇥ d (Fig. 3 embed ), where L = C i P H P W , toproduce initial tokens T 0 2 R N ⇥ d .
Next, a handcrafted positional encoding [41] or learned positionalembedding  2 R N ⇥ d [13] is added to tokens to encode positional information and break permutationinvariance.
Finally, a learnable class token is appended to the token sequence, resulting with
T 2 R ( N +1) ⇥ d .
For the purposes of our task, we omit this token in all ViT-based models.
In
ViTCA, the input to the embedding is a ﬂattened cell grid Z 2 R N ⇥ L where L = C P P H P W + C h ,
C P = C i + C
o + C
, C h is the cell hidden size, C o is the number of output image channels (oneor three for grayscale or RGB), and C  is the positional encoding size when positional encoding is(optionally) concatenated to each cell rather than added to the tokens [29].
Multi-head self-attention (MHSA).
Given a sequence of tokens T , self-attention estimates therelevance of one token to all others ( e.g ., which image patches are likely to appear together in animage) and aggregates this global information to update each token.
This encodes each token in termsof global contextual information, and does so using three learned weight matrices: W Q 2 R d ⇥ d ,
W K 2 R d ⇥ d , and W V 2 R d ⇥ d .
T is projected onto these weight matrices to obtain Queries Q =
TW Q , Keys K = TW K , and Values V = TW V .
The self-attention layer output SA 2 R N ⇥ d is:
[Equation]
Multi-head self-attention employs many sets of weight matrices, { W Q i , W K i , W V
i 2 R d ⇥ ( d/h ) |i = 0 , ..., ( h  1) } .
The outputs of h self-attention heads are concatenated into ( SA 0 , ..., SA h  1 ) 2 R N ⇥ d and projected onto a weight matrix W 2 R d ⇥ d to produce MHSA 2 R N ⇥ d .
Self-attention explicitly models global interactions and is more ﬂexible than grid-based operators( e.g ., convolutions) [33, 38], but its quadratic cost in time and memory limits its applicability to highresolution images.
Spatially localizing self-attention.
The global nature of self-attention directly conﬂicts with thespatial locality constraint of CAs; in response, we limit the connectivity structure of the attentionoperation to each cell’s neighbourhood.
This can be accomplished by either masking each head’sattention matrix ( A = softmax ( · · · ) 2 R N ⇥ N in Eq. 1) with a banded matrix representing localconnectivity ( e.g ., Fig.
3 localize ), or more efﬁciently,
[Equation]
SA
[Equation]
3)with ( V ) jl where i = { 0 , ..., ( N  1) } , j = { ( i + n w + n h ) , ..., i, ..., ( i  n w  n h ) } , and l = { 0 , ..., ( d 1) } , and with n w = { b N W / 2 c , ..., 0 , ..., b N W / 2 c } and n h = {  W b N H / 2 c , ..., 0 , ..., W b N H / 2 c } .
Here, we assume top-left-to-bottom-right input ﬂattening.
Instead of explicitly computing the globalself-attention matrix A 2 R N ⇥ N then masking it, this approach circumvents the O ( N 2 d ) computationin favour of an O ( NMd ) alternative that indexes the necessary rows and columns during self-attention.
The result is a localized self-attention matrix A ? 2 R N ⇥ M , where M = N H N W ⌧ N .
As we show inour experiments, ViTCA is still capable of global self-attention despite its localization, by leveragingstored state information across cells and their global self-organization during CA iterations (Fig. 2).
Following MHSA is a multilayer perceptron (Fig. 3 MLP ) with two layers and a GELU non-linearity.
We apply Layer Normalization (LN) [43] before MHSA and MLP , and residual connections afterwards,forming a single encoding block.
We use an MLP head (Fig. 3 head ) to decode to a desired output, with LN applied to its input, ﬁnalizing the ViTCA update rule F ✓ .
In our experiments, ViT’s headdecodes directly into an image output whereas ViTCA decodes into update vectors added to cells.
3.1 Update rule training procedure
To train the ViTCA update rule, we follow a “pool sampling”-based training process [7, 30] along witha curriculum-based masking/noise schedule when corrupting inputs.
During odd training iterations,we uniformly initialize a minibatch of cells Z
= ( Z 1 , ..., Z b ) with constant values (0.5 for outputchannels, 0 for hidden—see Appendix A.2 for alternatives), then inject the masked input mask ( X )(see Sec. 4.1).
After input injection, we asynchronously update cells (  = 50% update rate) using
F ✓ for T ⇠ U{ 8 , 32 } recurrent iterations.
We retrieve output Z o from the cell grid and apply an L 1 161 loss against the ground truth X .
We also apply overﬂow losses to penalize cell output values outsideof [0,1] and cell hidden values outside of [-1,1].
We use L 2 normalization on the gradient of eachparameter in ✓ .
After backpropagation, we append the updated cells and their ground truths to a pool
P which we then shufﬂe and truncate up to the ﬁrst N P elements.
During even training iterations,we retrieve a minibatch of cells and their ground truths from P and process them as above.
Thisencourages F ✓ to guide cells towards a stable ﬁxed-point.
Alg. 1 in Appendix
A details this process.
Experiments
Here we examine ViTCA through extensive experiments.
We begin with experiments for denoisingautoencoding, then an ablation study followed by various qualitative analyses, before concludingwith linear probing experiments on the learned representations for MNIST [50], FashionMNIST [42],and CIFAR10 [53].
We provide an extension to our experiments in Appendix A.
Baseline models and variants.
Since we are performing pixel level reconstructions, we create a
ViT baseline in which the class token has been removed.
This applies identically for ViTCA.
For our ViT and ViTCA baselines, we use a patch size of 1 ⇥ 1 ( P H =
P W =1 ), and only a single encodingblock with h = 4 MHSA heads, embed size d = 128 , and MLP size of 128 .
For ViTCA, we choose N H =3 and N W =3 ( i.e ., the Moore neighbourhood [1]).
We also compare with a U-Net baselinesimilar to the original formulation [48], but based on the speciﬁc architecture from [37].
Since mostof our datasets consist of 32 ⇥ 32 (resampled) images, we only have two downsampling steps asopposed to ﬁve.
We implement a U-Net-based CA (UNetCA) baseline consisting of a modiﬁedversion of our U-Net with 48 initial output feature maps as opposed to 24 and with all convolutionsexcept the ﬁrst changed to 1 ⇥ 1 to respect typical NCA restrictions [7, 30].
4.1 Denoising autoencoding
We compare between our baseline models and a number of ViTCA variants in the context ofdenoising autoencoding.
We present test set results across six benchmark datasets: a land coverclassiﬁcation dataset intended for representation learning (LandCoverRep) [25], MNIST, CelebA[47], FashionMNIST, CIFAR10, and Tiny ImageNet (a subset of ImageNet [49]).
All datasets consistof 32 ⇥ 32 resampled images except Tiny ImageNet, which is at 64 ⇥ 64 resolution.
During testing,we use all masking combinations, chosen in a ﬁxed order, and we update cells using a ﬁxed numberof iterations ( T =64 ).
See Tab. 1 for quantitative results.
Brieﬂy mentioned in Sec.
3.1, we employ a masking strategy inspired by Curriculum Learning (CL)[23, 52] to ease training.
This schedule follows a geometric progression of difﬁculty—tied to trainingiterations—maxing out at 10K training iterations.
Speciﬁcally, masking starts at covering 25% of theinput with 1 ⇥ 1 patches of noise (dropout for RGB inputs, Gaussian for grayscale), then at each shift indifﬁculty, new masking conﬁgurations are added to the list of available masking conﬁgurations in thefollowing order: (2 0 ⇥ 2 0 , 50%) , (2 0 ⇥ 2 0 , 75%) , (2 1 ⇥ 2 1 , 25%) , (2 1 ⇥ 2 1 , 50%) , (2 1 ⇥ 2 1 , 75%) , ..., (2 2 ⇥2 2 , 75%) .
Masking conﬁgurations are randomly chosen from this list.
We initialize weights/parameters using He initialization [46], except for the ﬁnal layer of CA-basedmodels, which are initialized to zero [30].
Unless otherwise stated, we train for I =100 K iterations,use a minibatch size b =32 , AdamW optimizer [36], learning rate ⌘ =10  3 with a cosine annealingschedule [40], pool size N P = 1024 , and cell hidden channel size C h = 32 .
In the case of Tiny
ImageNet, b = 8 to accommodate training on a single GPU (48GB Quadro RTX 8000).
Trainingtypically lasts a day at most, depending on the model.
Due to the recurrent iterations requiredper training step, CA-based models take the longest to train.
To alleviate memory limitations forsome of our experiments, we use gradient checkpointing [44] during CA iterations at the cost of
MNIST
[Equation]
[Table]
backpropagation duration and slight variations in gradients due to its effect on round-off propagation.
We also experiment with a cell fusion and mitosis scheme as an alternative.
See Appendix A fordetails on runtime performance, gradient checkpointing, and fusion and mitosis.
Amongst baselines, ViTCA outperforms on most metrics across the majority of datasets used (10 outof 18).
Exceptions include LandCoverRep, where UNetCA universally outperforms by a small margin,likely due to the texture-dominant imagery being amenable to convolutions.
Notably, ViTCA stronglyoutperforms on MNIST.
Although MNIST is a trivial dataset for common tasks such as classiﬁcation,our masking/noise strategy turns it into a challenging dataset for denoising autoencoding, e.g ., itis difﬁcult for even a human to classify a 32 ⇥ 32 MNIST digit 75% corrupted by 4 ⇥ 4 patches of
Gaussian noise.
We hypothesize that when compared to convolutional models, ViTCA’s weakerinductive biases (owed to attention [5, 8]) immediately outperform these models when there are largeregions lacking useful features, e.g ., MNIST digits cover a small space in the canvas.
This is notthe case with FashionMNIST, where the content is more ﬁlled out.
Between baselines and ViTCAvariants, ViTCA-32 (32 heads) and 32xy (xy-coordinate positional encoding) outperform all modelsby large margins, demonstrating the beneﬁts of multi-head self-attention.
We also experiment with aparameter-reduced (by ⇠ 60% ), inverted bottleneck variant where d =64 and MLP size is 256, oftenwith a minimal reduction in performance.
4.1.1 Ablation study
In Tab. 2 we perform an ablation study using the baseline ViTCA model above as reference on
CelebA. Results are ordered in row-wise blocks, top-to-bottom.
Speciﬁcally, we examine the impactof varying the cell hidden size C h ; the embed size d ; the number of MHSA heads h ; the depth (#encoders), comparing both ViTCA (used throughout the table) with ViT; and in the last block weexamine the impact of various methods of incorporating positional information into the model.
[Table]
Speciﬁcally, we examine the use of: (1)a xy-coordinate-based positional encodingconcatenated (“injected”) to cells, and;(2)
a Transformer-based positional encoding (or embedding, if learned) added intoembed .
These two categories are subdivided into: (1a) sincos5—consisting of handcrafted Fourier features [29] with four doublings of a base frequency, i.e .,  =(sin 2 0 ⇡ p,cos 2 0 ⇡ p, ..., sin 2 J  1 ⇡ p, cos 2 J  1 ⇡ p ) 2 238 R N ⇥ (4 JP H P W ) where J =5 and p is the pixelcoordinate (normalized to [-1,1]) for each pixelthe cell is situated on (one pixel since P H =
P W =1 ); (1b) sincos5xy—consisting of both
Fourier features and explicit xy-coordinatesconcatenated; (1b) xy—only xy-coordinates;(2a) handcrafted (our baseline approach)—sinusoidal encoding  2 R N ⇥ d similar to (1a)but following a Transformer-based approach[41], and; (2b) learned—learned embedding 2 R N ⇥ d following the original ViT approach[13].
To further test the self-organizing capabilities of ViTCA, we also include: (3) none—noexplicit positioning provided, where we let thecells localize themselves.
As shown in Tab. 2, ViTCA beneﬁts from anincrease to most CA and Transformer-centricparameters, at the cost of computational complexity and/or an increase in parameter count.
A noticeable decrease in performance is observed when embed size d =512 , most likelydue to the vast increase in parameter count necessitating more training.
In the original ViT,multiple encoding blocks were needed beforethe model could exhibit performance equivalent to their baseline CNN [13], as veriﬁed inour ablation with our ViT. However, for ViTCAwe notice an inverse relationship of the effectof Transformer depth, causing a divergence incell state.
It is not clear why this is the case, aswe have observed that the LN layers and overﬂow losses otherwise encourage a contractive
F ✓ .
This is an investigation we leave for futurework.
Despite the beneﬁts of increasing h , weuse h =4 for our baseline to optimize runtime performance.
Finally, we show that ViTCA does notdramatically suffer when no explicit positioning is used—in contrast to typical Transformer-basedmodels—as cells are still able to localize themselves by relying on their stored hidden information.
4.1.2 Cell state analysis
Here we provide an empirically-based qualitative analysis on the effects ViTCA and UNetCA haveon cell states through several experiments with our pre-trained models (Fig. 4 (a,b,c)).
We noticethat in general, ViTCA indeﬁnitely maintains cell state stability while UNetCA typically induces adivergence past a certain point.
An extended analysis is available in Appendix A.3.
Damage resilience.
Shown in Fig.
4 (a), we damage a random H/ 2 ⇥ W/
2 patch of cells withrandom values ⇠ U (  1 , 1) twice in succession.
ViTCA is able to maintain cell stability despite notbeing trained to deal with such noise, while UNetCA induces a divergence.
Note both models aresimultaneously performing the typical denoising task.
We also note that ViTCA’s inherent damageresilience is in contrast to recent NCA formulations that required explicit training for it [7, 30].
Convergence stability.
Fig.
4 (b) shows denoising results after 2784 cell grid updates.
ViTCA isable to maintain a stable cell grid state while UNetCA causes cells to diverge.
Hidden state visualizations.
Fig.
4 (c) shows 2D and 3D PCA dimensionality reductions on thehidden states of converged cell grids for all examples in FashionMNIST [42].
The clusters suggestsome linear separability in the learned representation, motivating our probing experiments in Sec. 4.2.
4.1.3 Investigating update rule inductive biases
Here we investigate the inductive biases inherent in ViTCA and UNetCA by testing their adaptationto various environmental changes (Fig. 4 (d,e,f,g,h)).
Adaptation to varying update rates.
Despite being trained with a  =50% cell update rate, ViTCAis able to adapt to varying rates (Fig. 4 (d)).
Higher rates result in a proportionally faster rate of cellstate convergence, and equivalently with lower rates.
UNetCA exhibits a similar relationship, althoughis unstable at  =100%
(see Appendix A.3).
For details comparing training with a synchronous vs .asynchronous cell grid update, see Appendix A.2.
Generalization to noise unseen during training.
ViTCA is capable of denoising conﬁgurations ofnoise it has not been trained on.
Fig at % coverage.
contrast, induces a (see Appendix A
Adaptation to changing inputs.
At various moments during cell updates, we re-inject cells with newmasked inputs (Fig.
(f)).
ViTCA is able to consistently adapt cells to new inputs while UNetCAexperiences difﬁculty past a certain point ( e.g ., at 464 iterations in the ﬁgure).
Effects of not vs .
completely masking input.
Fig.
4 (g; left ): ViTCA is able to perform autoencodingdespite not being trained for it.
UNetCA induces a cell grid divergence (see Appendix A.3).
Fig.
4 (g;right ): Interestingly, when the input is completely masked, ViTCA outputs the median image [37].
UNetCA does not exhibit such behaviour and instead causes cells to diverge (see Appendix A.3).
Spatial interpolation.
We use ViTCA models trained at 32 ⇥ 32 using various types of positioningto generate 128 ⇥ 128 outputs during inference, assuming an identical cell grid resolution.
Fig.4 (h; top-to-bottom of outputs ): xy-coordinates, no positioning, Fourier features [29], Fourierfeatures concatenated with xy-coordinates, and a Transformer-based handcrafted positional encoding(baseline) [41].
Results are ordered from best to worst.
The baseline approach is not capable ofspatial interpolation due to being a 1D positioning, while, as expected, the 2D encodings make itcapable.
Surprisingly, removing Fourier features and using only xy-coordinates results in a higherﬁdelity interpolation.
We believe this to be caused by the distracting amount of positional information
[Table]
All baselines and variants were pre-trained for denoising autoencoding and kept ﬁxed during probing.
A linear classiﬁer and 2-layer Multilayer Perceptrons (MLP) were trained on raw image inputs.
Boldface and underlined values denote the best and second best results, respectively.
Interestingly, CA-based models trained for denoising autoencoding on increasingly challenging datasets produce an increasingly more useful self-supervised representation for image classiﬁcation compared to non-CA-based models.
Fourier features provide to cells, as cells can instead rely on their hidden states to store higherfrequency positional information.
Finally, with no explicit positioning, ViTCA is still able to performhigh-quality interpolation—even exceeding using Fourier features—by taking advantage of its selforganizing nature.
As a side note, we point attention to the fact that ViTCA is simultaneouslydenoising at a scale space it has not been trained on, exemplifying its generalization capabilities.
4.2 Investigating hidden representations via linear probes
Here we examine the learned representations of our models pre-trained for denoising.
We freezemodel parameters and learn linear classiﬁers on each of their learned representations: convergedcell hidden states for CA-based models, bottleneck features for U-Net, and LN’d tokens for ViT.
This is a common approach used to probe learned representations [28].
Classiﬁcation results on
MNIST, FashionMNIST, and CIFAR10 are shown in Tab.
and we use the same training setup as fordenoising, but without any noise.
For comparison, we also provide results using a linear classiﬁerand two 2-layer MLPs of varying complexity, all trained directly on raw pixel values.
Correlationsbetween denoising performance in Tab. 1 and classiﬁcation performance in Tab. 3 can be observed.
Linear classiﬁcation accuracy on ViTCA-based features typically exceeds classiﬁcation accuracyusing other model-based features or raw pixel values, even outperforming the MLPs in most cases.
Discussion
We have performed extensive quantitative and qualitative evaluations of our newly proposed ViTCAon a variety of datasets under a denoising autoencoding framework.
We have demonstrated thesuperior denoising performance and robustness of our model when compared to a U-Net-based CAbaseline (UNetCA) and ViT, as well as its generalization capabilities under a variety of environmentalchanges such as larger inputs ( i.e ., spatial interpolation) and changing inputs during cell updates.
Despite the computation savings—owed to our circumvention of self-attention’s quadratic complexityby spatially localizing it within ViTCA—there remains the same memory limitations inherent to allrecurrent models: multiple recurrent iterations are required for each training iteration, resulting inlarger memory usage than a feedforward approach.
This limits single-GPU training accessibility.
We have experimented with gradient checkpointing [44] but found its trade-off for increased backpropagation duration (and slightly different gradients)
less than ideal.
To fully realize the potentialof NCAs (self-organization, inherent distributivity, etc .), we encourage follow-up work to addressthis limitation.
Adapting recent techniques using implicit differentiation is one avenue to circumventthese issues [2, 32].
Also, as mentioned in our ablation (Sec. 4.1.1), we hope to further investigatethe instabilities caused by increasing the depth of ViTCA.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’scontributions and scope?
[Yes](b) Did you describe the limitations of your work?
[Yes] See Sec. 5.(c)
Did you discuss any potential negative societal impacts of your work?
[No] Althoughwe feel our work demonstrates the potential of NCAs as viable alternatives to commonrecurrent network architectures (ViTCA being our evidential contribution), our experiments intentionally tend towards the direction of optimizing model efﬁciency (andsingle-GPU training accessibility) rather than towards the increasingly popular directionof scaling upwards.
However, as much as our work demonstrates the downward-scalingcapabilities of NCAs, we also acknowledge that this similarly applies going upward,and as such, can be abused ( e.g ., creating a “deepfake”-capable ViTCA).(d)
Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are including theoretical results...(
a) Did you state the
full set of assumptions of all theoretical results [N/A] [Table]
(b) Did you include complete proofs of all theoretical results?
[N/A]
[Table]
If you ran experiments...
[Table]
(a) Did you include the code, data, and instructions needed to reproduce the main exper-
[Table]
imental results (either in the supplemental material or as a URL)?
[Yes] Code andinstructions to reproduce results are included in the supplemental material. [Table]
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
[Table]
were chosen)?
[Yes]
[Table]
(c) Did you report error bars (e.g., with respect to the random seed after running ex-
periments multiple times)?
[No] Given the combination of time and computational
restrictions and our exhaustive list of experiments, we opted to prioritize experiment
variety and dataset coverage as an implicit substitute for re-running experiments under
different random seeds.
For all experiments, we kept a ﬁxed random seed, even point-
ing out (deterministic) differences caused by gradient checkpointing when used (seeAppendix A).
(d) Did you include the total amount of compute and the type of resources used (e g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[N/A] Licensed frameworks used such as
PyTorch (BSD-style) and Hydra (MIT) will be mentioned in acknowledgements.
(c) Did you include any new assets either in the supplemental material or as a URL?
[No] No new assets—aside from code and training our models—were created for the
purposes of this work.
[Table]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[No] We used publicly available datasets.
[Table]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[No] Although not discussed in the manuscript,
we
would like to point that the datasets we used that could potentially contain personally
identiﬁable information (CelebA, CIFAR10, Tiny ImageNet) each have restrictions
and/or acknowledgements of such potential issues.
Also, our work is not focused on
classifying persons and ViTCA is not a generative model, e.g ., it can not generate new
faces
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]