RiteeshGrandhi/sciclaimeval-shared-task huggingface dataset:
Vision Transformer Cellular Automata (ViTCA)
Building upon NCAs and ViTs, we propose a new class of attention-based NCAs formed using a spatially localized—yet globally organized—self-attention scheme. We detail an instance of this class, ViTCA, by first reviewing its backbone ViT architecture before describing the “pool sampling”-based training process for the ViTCA update rule (see overview in Fig. 3 ).
Update rule training procedure
To train the ViTCA update rule, we follow a “pool sampling”-based training process [ 7 , 30 ] along with a curriculum-based masking/noise schedule when corrupting inputs. During odd training iterations, we uniformly initialize a minibatch of cells \mathbf{Z} \!=\!(\mathbf{Z}_{1},...,\mathbf{Z}_{b}) with constant values (0.5 for output channels, 0 for hidden—see Appendix A.2 for alternatives), then inject the masked input \texttt{mask}(\mathbf{X}) (see Sec. 4.1 ). After input injection, we asynchronously update cells ( \sigma\!=\!50\% update rate) using F_{\theta} for T\!\sim\!\mathcal{U}\{8,32\} recurrent iterations. We retrieve output \mathbf{Z}_{o} from the cell grid and apply an L_{1} loss against the ground truth \mathbf{X} . We also apply overflow losses to penalize cell output values outside of [0,1] and cell hidden values outside of [-1,1]. We use L_{2} normalization on the gradient of each parameter in \theta . After backpropagation, we append the updated cells and their ground truths to a pool \mathcal{P} which we then shuffle and truncate up to the first N_{\mathcal{P}} elements. During even training iterations, we retrieve a minibatch of cells and their ground truths from \mathcal{P} and process them as above. This encourages F_{\theta} to guide cells towards a stable fixed-point. Alg. 1 in Appendix A details this process.
Experiments
Here we examine ViTCA through extensive experiments. We begin with experiments for denoising autoencoding, then an ablation study followed by various qualitative analyses, before concluding with 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 .
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. Specifically, we examine the impact of 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 we examine the impact of various methods of incorporating positional information into the model.

Specifically, we examine the use of: (1) a xy-coordinate-based positional encoding concatenated (“injected”) to cells, and; (2) a Transformer-based positional encoding (or embedding, if learned) added into embed . These two categories are subdivided into:
(1a) sincos5—consisting of handcrafted Fourier features [ 29 ] with four doublings of a base frequency, i.e . , {\setBold[0.75]\gamma}\unsetBold \!=\!(\sin{2^{0}\pi p}, \cos{2^{0}\pi p}, ...,\sin{2^{J-1}\pi p}, \cos{2^{J-1}\pi p})\!\in\mathbb{R}^{N\times(4JP_{H}P_{W})} where J\!=\!5 and p is the pixel coordinate (normalized to [-1,1]) for each pixel the cell is situated on (one pixel since P_{H}\!=\!P_{W}\!=\!1 );
(1b) sincos5xy—consisting of both Fourier features and explicit xy-coordinates concatenated;
(1c) xy—only xy-coordinates;
(2a) handcrafted (our baseline approach)—sinusoidal encoding {\setBold[0.75]\gamma}\unsetBold \!\in\!\mathbb{R}^{N\times d} similar to (1a) but following a Transformer-based approach [ 41 ] , and;
(2b) learned—learned embedding {\setBold[0.75]\gamma}\unsetBold \!\in\!\mathbb{R}^{N\times d} following the original ViT approach [ 13 ] .
To further test the self-organizing capabilities of ViTCA, we also include: (3) none—no explicit positioning provided, where we let the cells localize themselves.

As shown in Tab. 2 , ViTCA benefits from an increase to most CA and Transformer-centric parameters, 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 likely due to the vast increase in parameter count necessitating more training. In the original ViT, multiple encoding blocks were needed before the model could exhibit performance equivalent to their baseline CNN [ 13 ] , as verified in our ablation with our ViT. However, for ViTCA we notice an inverse relationship of the effect of Transformer depth, causing a divergence in cell state. It is not clear why this is the case, as we have observed that the LN layers and overflow losses otherwise encourage a contractive F_{\theta} . This is an investigation we leave for future work. Despite the benefits of increasing h , we use h\!=\!4 for our baseline to optimize runtime performance. Finally, we show that ViTCA does not dramatically suffer when no explicit positioning is used—in contrast to typical Transformer-based models—as cells are still able to localize themselves by relying on their stored hidden information.
Cell state analysis
Here we provide an empirically-based qualitative analysis on the effects ViTCA and UNetCA have on cell states through several experiments with our pre-trained models (Fig. 4 (a,b,c)). We notice that in general, ViTCA indefinitely maintains cell state stability while UNetCA typically induces a divergence 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\!\times\!W/2 patch of cells with random values \sim\!\mathcal{U}(-1,1) twice in succession. ViTCA is able to maintain cell stability despite not being trained to deal with such noise, while UNetCA induces a divergence. Note both models are simultaneously performing the typical denoising task. We also note that ViTCA’s inherent damage resilience 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 is able 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 the hidden states of converged cell grids for all examples in FashionMNIST [ 42 ] . The clusters suggest some linear separability in the learned representation, motivating our probing experiments in Sec. 4.2 .
Investigating update rule inductive biases
Here we investigate the inductive biases inherent in ViTCA and UNetCA by testing their adaptation to various environmental changes (Fig. 4 (d,e,f,g,h)).

Adaptation to varying update rates. Despite being trained with a \sigma\!=\!50\% cell update rate, ViTCA is able to adapt to varying rates (Fig. 4 (d)). Higher rates result in a proportionally faster rate of cell state convergence, and equivalently with lower rates. UNetCA exhibits a similar relationship, although is unstable at \sigma\!=\!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 configurations of noise it has not been trained on. Fig. 4 (e; left-to-right ): 4\!\times\!1 and 1\!\times\!4 patches of Gaussian noise at 65% coverage. In contrast, UNetCA induces a cell state divergence (see Appendix A.3 ).

Adaptation to changing inputs. At various moments during cell updates, we re-inject cells with new masked inputs (Fig. 4 (f)). ViTCA is able to consistently adapt cells to new inputs while UNetCA experiences difficulty past a certain point ( e.g . , at 464 iterations in the figure).

Effects of not vs . completely masking input. Fig. 4 (g; left ): ViTCA is able to perform autoencoding despite 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\!\times\!32 using various types of positioning to generate 128\!\times\!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 ] , Fourier features 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 of spatial interpolation due to being a 1D positioning, while, as expected, the 2D encodings make it capable. Surprisingly, removing Fourier features and using only xy-coordinates results in a higher fidelity interpolation. We believe this to be caused by the distracting amount of positional information Fourier features provide to cells, as cells can instead rely on their hidden states to store higher frequency positional information. Finally, with no explicit positioning, ViTCA is still able to perform high-quality interpolation—even exceeding using Fourier features—by taking advantage of its self-organizing nature. As a side note, we point attention to the fact that ViTCA is simultaneously denoising at a scale space it has not been trained on, exemplifying its generalization capabilities.
Investigating hidden representations via linear probes
Here we examine the learned representations of our models pre-trained for denoising. We freeze model parameters and learn linear classifiers on each of their learned representations: converged cell 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 ] . Classification results on MNIST, FashionMNIST, and CIFAR10 are shown in Tab. 4.2 and we use the same training setup as for denoising, but without any noise. For comparison, we also provide results using a linear classifier and two 2-layer MLPs of varying complexity, all trained directly on raw pixel values. Correlations between denoising performance in Tab. 4.1 and classification performance in Tab. 4.2 can be observed. Linear classification accuracy on ViTCA-based features typically exceeds classification accuracy using 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 ViTCA on a variety of datasets under a denoising autoencoding framework. We have demonstrated the superior denoising performance and robustness of our model when compared to a U-Net-based CA baseline (UNetCA) and ViT, as well as its generalization capabilities under a variety of environmental changes 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 complexity by spatially localizing it within ViTCA—there remains the same memory limitations inherent to all recurrent models: multiple recurrent iterations are required for each training iteration, resulting in larger 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 potential of NCAs (self-organization, inherent distributivity, etc . ), we encourage follow-up work to address this limitation. Adapting recent techniques using implicit differentiation is one avenue to circumvent these issues [ 2 , 32 ] . Also, as mentioned in our ablation (Sec. 4.1.1 ), we hope to further investigate the instabilities caused by increasing the depth of ViTCA.
Acknowledgments and Disclosure of Funding
First and foremost, M.T. thanks their former supervisor and mentor, Konstantinos (Kosta) G. Derpanis, for his invaluable support throughout the project.
M.T. also thanks Martin Weiss for his helpful feedback on implementing the linear probe experiments (Sec. 4.2 );
Olexa Bilaniuk for his assistance in investigating the gradient differences caused by PyTorch’s gradient checkpointing implementation (see Appendix A.2 ), and;
the Mila Innovation, Development, and Technology (IDT) team for their overall technical support, particularly, their tireless efforts maintaining cluster reliability during the crucial moments preceding the submission deadline.

M.T. is partially supported by the Natural Sciences and Engineering Research Council of Canada (NSERC) Canada Graduate Scholarship – Doctoral [application number CGSD3-519428-2018].
D.N. and C.P. are each partially supported by an NSERC Discovery Grant [application IDs 5011360 and 5018358, respectively]. D.N. thanks Samsung Electronics Co. Ldt. for their support. C.P. thanks CIFAR for their support under the AI Chairs Program.
A.1 Training on high-resolution imagery with fusion and mitosis
As an alternative to gradient checkpointing for reducing memory usage, we briefly experimented with a downsampling scheme inspired by cell fusion and mitosis when training on CelebA at 64\!\times\!64 . Specifically, we split the T applications of the update rule (within a training iteration) into multiple stages: 1) We apply the update rule twice so that cells will have, at minimum, some amount of knowledge of their neighbours. 2) We stash the masked input for a later re-injection. 3) Fusion —we apply a 2\!\times\!2 average pooling with a stride of 2 across the cell grid, combining 2\!\times\!2 groups of cells into singular cells. 4) We apply the update rule T-4 times at this 32\!\times\!32 downsampled cell grid resolution. 5) Mitosis —we perform a 2\!\times\!2 duplication of cells (each cell is duplicated to its right, bottom-right, and bottom). 6) We re-inject the stashed masked input. 7) We apply the update rule twice to adapt the cells to the 64\!\times\!64 resolution and to fill in any missing information.

We found that performing this fusion and mitosis scheme decreased training memory consumption to levels similar to our gradient checkpointing scheme ( \sim\!50\% memory reduction) while having a \sim\!70\% faster backward pass. Loss-wise, we observed a \sim\!33\% increase in the average validation reconstruction loss during training, which can qualitatively be observed in the example provided in Fig. 6 ( bottom ). Although the results shown are not ideal— i.e . , we did not perform a hyper-parameter search here, for example, finding the optimal number of iterations preceding fusion and following mitosis—this brief experiment tests the feasibility of reducing memory consumption while maintaining denoising capability and avoiding gradient checkpointing. As shown in the figure, ViTCA with fusion and mitosis is able to successfully denoise the input despite applying updates at two different scales. This scale agnostic behaviour reveals potentially interesting research directions beyond the scope of this work, such as allowing an NCA update rule to dynamically and locally modify cell grid resolution based on a compute budget, which could see applications in signal (image, video, or audio) compression.

[0.4]

[0.56] PSNR \uparrow SSIM \uparrow LPIPS \downarrow # Params. Pool size 128 26.51 0.914 0.065 92.5K 256 26.40 0.912 0.067 92.5K 512 26.61 0.915 0.064 92.5K 1024 26.53 0.913 0.066 92.5K 2048 26.54 0.915 0.064 92.5K 4096 26.48 0.912 0.066 92.5K 8192 26.30 0.910 0.069 92.5K Cell init. constant 26.53 0.913 0.066 92.5K random 25.90 0.905 0.074 92.5K Patch size 1 \times 1 26.53 0.913 0.066 92.5K 2 \times 2 25.85 0.906 0.076 96.0K 4 \times 4 24.54 0.882 0.113 109.8K 8 \times 8 21.62 0.803 0.212 165.3K 16 \times 16 18.71 0.687 0.279 387.0K
A.2 Extended ablation study
Here we present an extension of our ablation study in Sec. 4.1.1 , using the baseline ViTCA model as our reference. As before, the ablation examines the effects certain training configuration parameters have on test performance.
A.3 Extended analysis of cell state and update rule inductive biases
Here we present an extension of the analyses provided in Sec. 4.1.2 and Sec. 4.1.3 .
A.4 Runtime analysis of ViTCA
Here we provide a brief analysis of ViTCA’s runtime performance and memory usage while training on a minibatch of random 32\times 3\times H\times W images through measurements of forward pass duration (ms), backward pass duration (ms), and training memory usage (GB), with and without using gradient checkpointing. We use T\!=\!32 ViTCA iterations and 16 checkpoint segments. Results are shown in Tab. 7 . Gradient checkpointing provides substantial memory savings at the cost of proportionally increasing the duration of the backward pass.
Figure 1: ViT vs . ViTCA for denoising Tiny ImageNet [ 49 ] validation set images with 2 \times 2 pixel masks covering 75\% of the image. Top-to-bottom : noisy input, ViT, ViTCA, and ground truth.; ViT vs . ViTCA for denoising Tiny ImageNet [ 49 ] validation set images with 2 \times 2 pixel masks covering 75\% of the image. Top-to-bottom : noisy input, ViT, ViTCA, and ground truth.
Figure 2: Global self-organization manifested within localized self-attention . Despite operating in spatially local neighbourhoods about a cell, over time the localized (multi-head) self-attention in ViTCA experiences a global self-organization admitted by its NCA nature. This circumvents the quadratic complexity of explicit global self-attention (w.r.t. input size) with a linear amortization over time (recurrent CA iterations), enabling effective per-pixel dense processing. Middle : visualizing local attention maps about each cell as colour-coded “splats” blended together in overlapping regions, producing a “splat map” [ 58 ] . Left, right : ViTCA iterations on a cell grid, updated from a seed state to a converged state, given a noisy input image to denoise. For each head of the cells’ local attention maps, there is global agreement on the types of features to attend to ( e.g . , foreground contours, noise, background). Enveloping ViT by the NCA paradigm dramatically improves its output fidelity.; Global self-organization manifested within localized self-attention . Despite operating in spatially local neighbourhoods about a cell, over time the localized (multi-head) self-attention in ViTCA experiences a global self-organization admitted by its NCA nature. This circumvents the quadratic complexity of explicit global self-attention (w.r.t. input size) with a linear amortization over time (recurrent CA iterations), enabling effective per-pixel dense processing. Middle : visualizing local attention maps about each cell as colour-coded “splats” blended together in overlapping regions, producing a “splat map” [ 58 ] . Left, right : ViTCA iterations on a cell grid, updated from a seed state to a converged state, given a noisy input image to denoise. For each head of the cells’ local attention maps, there is global agreement on the types of features to attend to ( e.g . , foreground contours, noise, background). Enveloping ViT by the NCA paradigm dramatically improves its output fidelity.
Figure 3: Computational overview . NCAs use a stateful lattice of cells , each storing information along channels, to promote desired behaviour over the course of an evolutionary cycle. Starting from an initial seed, each cell state evolves at discrete time steps according to a homogeneous, learned update rule applied either synchronously or asynchronously ( \sigma ). This update depends on the current cell state and that of its neighbours (pictured is the Moore neighbourhood [ 1 ] ). In ViTCA , each cell is represented as a vector where the first C_{i}P_{H}P_{W} channels contain a P_{H}\times P_{W} noisy input image patch ( \texttt{mask}(\mathbf{x}) ), the next C_{o}P_{H}P_{W} channels contain the current output patch ( \mathbf{z}^{t}_{o} ), the following C_{h} channels contain undefined data hidden from the loss that can be used to encode additional information ( \mathbf{z}^{t}_{h} ), and (optionally) the remaining C_{\gamma}P_{H}P_{W} channels contain positional information ( {\setBold[0.75]\gamma}\unsetBold ). The update rule ( F_{\theta} ) is a modified ViT [ 13 ] whose self-attention mechanism is locally constrained to each cell’s neighbourhood ( localize ).; Computational overview . NCAs use a stateful lattice of cells , each storing information along channels, to promote desired behaviour over the course of an evolutionary cycle. Starting from an initial seed, each cell state evolves at discrete time steps according to a homogeneous, learned update rule applied either synchronously or asynchronously ( \sigma ). This update depends on the current cell state and that of its neighbours (pictured is the Moore neighbourhood [ 1 ] ). In ViTCA , each cell is represented as a vector where the first C_{i}P_{H}P_{W} channels contain a P_{H}\times P_{W} noisy input image patch ( \texttt{mask}(\mathbf{x}) ), the next C_{o}P_{H}P_{W} channels contain the current output patch ( \mathbf{z}^{t}_{o} ), the following C_{h} channels contain undefined data hidden from the loss that can be used to encode additional information ( \mathbf{z}^{t}_{h} ), and (optionally) the remaining C_{\gamma}P_{H}P_{W} channels contain positional information ( {\setBold[0.75]\gamma}\unsetBold ). The update rule ( F_{\theta} ) is a modified ViT [ 13 ] whose self-attention mechanism is locally constrained to each cell’s neighbourhood ( localize ).
Figure 4: Qualitative results. Gold boxes are inputs, green ground truths, purple ViTCA outputs, and blue UNetCA outputs. We analyze the effects of ViTCA and UNetCA on cell states in terms of: (a) damage resilience; (b) convergence stabilility, and; (c) hidden state PCA visualizations of converged cell grids for all examples in FashionMNIST [ 42 ] . We also investigate update rule inductive biases in terms of adapting to: (f) varying inputs during cell updates; (d) varying cell update rates; (e) noise configurations unseen during training; (g) unmasked and completely masked inputs, and; (h) spatial interpolation enabled by our various methods of incorporating cell positioning.; Qualitative results. Gold boxes are inputs, green ground truths, purple ViTCA outputs, and blue UNetCA outputs. We analyze the effects of ViTCA and UNetCA on cell states in terms of: (a) damage resilience; (b) convergence stabilility, and; (c) hidden state PCA visualizations of converged cell grids for all examples in FashionMNIST [ 42 ] . We also investigate update rule inductive biases in terms of adapting to: (f) varying inputs during cell updates; (d) varying cell update rates; (e) noise configurations unseen during training; (g) unmasked and completely masked inputs, and; (h) spatial interpolation enabled by our various methods of incorporating cell positioning.
Figure 5: Qualitative results using cell fusion and mitosis as an alternative to gradient checkpointing. Gold boxes are inputs, green ground truths, purple ViTCA outputs, and blue ViTCA w. fusion and mitosis outputs. Outputs are after 64 CA iterations.; Quantitative ablation on pool size N_{\mathcal{P}} , cell initialization method, and patch size P_{H}\times P_{W} for denoising autoencoding with ViTCA on CelebA. Boldface and underlining denote best and second best results. Italicized items denote baseline configuration settings.
Figure 7: Qualitative results showcasing UNetCA’s inductive biases in terms of adapting to: (a) varying cell update rates; (b) noise configurations unseen during training, and; (c) unmasked and completely masked inputs. Gold boxes are inputs, green ground truths, and blue UNetCA outputs.; Qualitative results showcasing UNetCA’s inductive biases in terms of adapting to: (a) varying cell update rates; (b) noise configurations unseen during training, and; (c) unmasked and completely masked inputs. Gold boxes are inputs, green ground truths, and blue UNetCA outputs.