CAX: Cellular Automata Accelerated in JAX

## Abstract

Abstract Cellular automata have become a cornerstone for investigating emergence and self-organization across diverse scientific disciplines, spanning neuroscience, artificial life, and theoretical physics.
However, the absence of a hardware-accelerated cellular automata library limits the exploration of new research directions, hinders collaboration, and impedes reproducibility.
In this work, we introduce CAX (Cellular Automata Accelerated in JAX), a high-performance and flexible open-source library designed to accelerate cellular automata research.
CAX offers cutting-edge performance and a modular design through a user-friendly interface, and can support both discrete and continuous cellular automata with any number of dimensions.
We demonstrate CAX’s performance and flexibility through a wide range of benchmarks and applications.
From classic models like elementary cellular automata and Conway’s Game of Life to advanced applications such as growing neural cellular automata and self-classifying MNIST digits, CAX speeds up simulations up to 2,000 times faster.
Furthermore, we demonstrate CAX’s potential to accelerate research by presenting a collection of three novel cellular automata experiments, each implemented in just a few lines of code thanks to the library’s modular architecture.
Notably, we show that a simple one-dimensional cellular automaton can outperform GPT-4 on the 1D-ARC challenge.

## 1 Introduction

Emergence is a fundamental concept that has captivated thinkers across various fields of human inquiry, including philosophy, science and art (Holland, ).
This fascinating phenomenon occurs when a complex entity exhibits properties that its constituent parts do not possess individually.
From the collective intelligence of ant colonies to the formation of snowflakes, self-organization and emergence manifest in myriad ways.
The study of self-organization and emergence holds the promise to unravel deep mysteries, from the origin of life to the development of conciousness.

Cellular automata (CA) are models of computation that exemplify how complex patterns and sophisticated behaviors can arise from simple components interacting through basic rules.
Originating from the work of Ulam and von Neumann in the 1940s (Neumann & Burks, ), these systems gained prominence with Conway’s Game of Life in the 1970s (Gardner, ) and Wolfram’s systematic studies in the 1980s (Wolfram, ).
The discovery that even elementary cellular automata can be Turing-complete underscores their expressiveness (Cook, ).
CAs serve as a powerful abstraction for investigating self-organization and emergence, offering insights into complex phenomena across scientific domains, from physics and biology to computer science and artificial life.

Figure: Figure 1: Cellular Automata types supported in CAX.
Refer to caption: x1.png

In recent years, the integration of machine learning techniques with cellular automata has opened new avenues for research in morphogenesis (Mordvintsev et al., ), self-organization (Randazzo et al., ; ), and developmental processes (Najarro et al., ).
The advent of Neural Cellular Automata (NCA) has significantly broadened the scope of CA research, yielding profound biological insights and showcasing the power of gradient-based optimization in studying emergence and self-organization.
NCAs extend traditional CAs by incorporating neural networks to learn update rules, allowing for more complex and adaptive behaviors. This approach enables the modeling of sophisticated phenomena such as pattern formation in biological systems and the evolution of artificial life forms.
This progress has not only deepened our understanding of complex systems but also underscored the growing computational demands of CA experiments, pointing to the potential for scaling through hardware-accelerated libraries inspired by advances in deep learning.

Despite their conceptual simplicity, cellular automata simulations can be computationally intensive, especially when scaling to higher dimensions with large numbers of cells or implementing backpropagation through time for NCAs.
Moreover, the implementation of CA in research settings has often been fragmented, with individual researchers frequently reimplementing basic functionalities, creating custom implementations across various deep learning frameworks such as TensorFlow, JAX, and PyTorch.
As the field continues to grow and attract increasing interest, there is a pressing need for a unified, robust library that facilitates collaboration, reproducibility, fast experimentation and exploration of new research directions.

In response to these challenges and opportunities, we present CAX: Cellular Automata Accelerated in JAX, an open-source library with cutting-edge performance, designed to provide a flexible and efficient framework for cellular automata research.
CAX is built on JAX (Bradbury et al., ), a high-performance numerical computing library, enabling to speed up cellular automata simulations through massive parallelization across various hardware accelerators such as CPUs, GPUs, and TPUs.
CAX is flexible and supports both discrete and continuous cellular automata with any number of dimensions, accommodating classic models like elementary cellular automata and Conway’s Game of Life, as well as modern variants such as Lenia and Neural Cellular Automata.

JAX offers efficient vectorization of CA rules, enabling millions of cell updates to be processed simultaneously.
It also provides automatic differentiation capabilities to backpropagate through time efficiently, facilitating the training of Neural Cellular Automata.
CAX can run experiments with millions of cell updates in minutes, reducing computation times by up to 2,000 times compared to traditional implementations in our benchmark.
This performance boost opens up new possibilities for large-scale CA experiments that were previously computationally prohibitive.

CAX’s flexibility and potential to accelerate research is showcased through three novel cellular automata experiments. Thanks to CAX’s modular architecture, each of these experiments is implemented in just a few lines of code (Appendix B), significantly reducing the barrier to entry for cellular automata research.
Notably, we show that a simple one-dimensional cellular automaton implemented with CAX outperforms GPT-4 on the 1D-ARC challenge (Xu et al., ), see Section 5.3.
Finally, to support users and facilitate adoption, CAX comes with high-quality, diverse examples and comprehensive documentation. The list of implemented CAs is detailed in Table 1.

## 2 Background

### 2.1 Cellular Automata

A cellular automaton is a simple model of computation consisting of a regular grid of cells, each in a particular state. The grid can be in any finite number of dimensions. For each cell, a set of cells called its neighborhood is defined relative to the specified cell. The grid is updated at discrete time steps according to a fixed rule that determines the new state of each cell based on its current state and the states of the cells in its neighborhood.

A CA is defined by a tuple $(\mathcal{L},\mathcal{S},\mathcal{N},\phi)$, where $\mathcal{L}$ is the $d$-dimensional lattice or grid with $c$ channels, $\mathcal{S}$ is the cell state set, $\mathcal{N}\subset\mathcal{L}$ is the neighborhood of the origin, and $\phi:\mathcal{S}^{\mathcal{N}}\rightarrow\mathcal{S}$ is the local rule.
A mapping from the grid to the cell state set $\mathbf{S}:\mathcal{L}\rightarrow\mathcal{S}$ is called a configuration or pattern. In this work, we will simply refer to it by the state of the CA. $\mathbf{S}(\mathbf{x})$ represents the state of a cell $\mathbf{x}\in\mathcal{L}$. Additionally, we denote the neighborhood of a cell $\mathbf{x}\in\mathcal{L}$ by $\mathcal{N}_{\mathbf{x}}=\left\{\mathbf{x}+\mathbf{n},\mathbf{n}\in\mathcal{N}\right\}$, and $\mathbf{S}(\mathcal{N}_{\mathbf{x}})=\left\{\mathbf{S}(\mathbf{n}),\mathbf{n}
\in\mathcal{N}_{\mathbf{x}}\right\}$.

The global rule $\Phi:\mathcal{S}^{\mathcal{L}}\rightarrow\mathcal{S}^{\mathcal{L}}$ applies the local rule uniformly to all cells in the lattice and is defined such that, for all $\mathbf{x}$ in $\mathcal{L}$, $\Phi(\mathbf{S})(\mathbf{x})=\phi(\mathbf{S}(\mathcal{N}_{\mathbf{x}}))$. A cellular automaton is initialized with a state $\mathbf{S}_{0}$. Then, the state is updated according to the global rule $\Phi$ at each discrete time step $t\in\mathbb{N}$, to give,

$$ $\Phi(\mathbf{S}_{0})=\mathbf{S}_{1},\Phi(\mathbf{S}_{1})=\mathbf{S}_{2},\dots$ $$

The close connection between CA and recurrent convolutional neural networks has been observed by numerous researchers (Gilpin, ; Wulff & Hertz, ; Mordvintsev et al., ; Chan, ). For example, the general NCA architecture introduced by Mordvintsev et al. () can be conceptualized as a “recurrent residual convolutional neural network with per-cell dropout”.

### 2.2 Controllable Cellular Automata

A controllable cellular automaton (CCA) is a generalization of CA that incorporates the ability to accept external inputs at each time step. CCAs formalize the concept of Goal-Guided NCA that has been introduced in the literature by Sudhakaran et al. (). The external inputs can modify the behavior of CCAs, offering the possibility to respond dynamically to changing conditions or control signals while maintaining the fundamental principles of cellular automata.

A CCA is defined by a tuple $(\mathcal{L},\mathcal{S},\mathcal{I},\mathcal{N},\phi)$, where $\mathcal{I}$ is the input set and $\phi:\mathcal{S}^{\mathcal{N}}\times\mathcal{I}^{\mathcal{N}}\rightarrow
\mathcal{S}$ is the controllable local rule.
A mapping from the grid to the input set $\mathbf{I}:\mathcal{L}\rightarrow\mathcal{I}$ is called the input. $\mathbf{I}(\mathbf{x})$ represents the input of a cell $\mathbf{x}\in\mathcal{L}$. Similarly to the state, we denote $\mathbf{I}(\mathcal{N}_{\mathbf{x}})=\{\mathbf{I}(\mathbf{n}),\mathbf{n}\in
\mathcal{N}_{\mathbf{x}}\}$.

The controllable global rule $\Phi:\mathcal{S}^{\mathcal{L}}\times\mathcal{I}^{\mathcal{L}}\rightarrow
\mathcal{S}^{\mathcal{L}}$ is defined such that, for all $\mathbf{x}$ in $\mathcal{L}$, $\Phi(\mathbf{S},\mathbf{I})(\mathbf{x})=\phi(\mathbf{S}(\mathcal{N}_{\mathbf{x
}}),\mathbf{I}(\mathcal{N}_{\mathbf{x}}))$. A controllable cellular automaton is initialized with an initial state $\mathbf{S}_{0}$. Then, the state is updated according to the controllable global rule $\Phi$ and a sequence of input $(\mathbf{I}_{t})_{t\geq 0}$ at each discrete time step $t\in\mathbb{N}$, to give,

$$ $\Phi(\mathbf{S}_{0},\mathbf{I}_{0})=\mathbf{S}_{1},\Phi(\mathbf{S}_{1},\mathbf {I}_{1})=\mathbf{S}_{2},\dots$ $$

As discussed in Section 2.1, CAs can be conceptualized as recurrent convolutional neural networks. However, traditional CAs lack the ability to take external inputs at each time step. CCAs extend the capabilities of traditional CAs by making them responsive to external inputs, akin to recurrent neural networks processing sequential data. CCAs bridge the gap between recurrent convolutional neural networks and cellular automata, opening up new possibilities for modeling complex systems that exhibit both autonomous emergent behavior and responsiveness to external control.

### 2.3 Related Work

The field of CA has spawned numerous tools and libraries to support research and experimentation, with CellPyLib (Antunes, ) emerging as one of the most popular and versatile options. This Python library offers a simple yet powerful interface for working with 1- and 2-dimensional CA, supporting both discrete and continuous states, making it an ideal baseline for comparative studies and further development. While it provides implementations of classic CA models like Conway’s Game of Life and Wireworld, CellPyLib is not hardware-accelerated and does not support the training of neural cellular automata. Golly is a cross-platform application for exploring Conway’s Game of Life and many other types of cellular automata. Golly’s features include 3D CA rules, custom rule loading, and scripting via Lua or Python. While powerful and versatile for traditional CA, Golly is not designed for hardware acceleration or integration with modern machine learning frameworks.

The recent surge in artificial intelligence has increased the availability of computational resources, and encouraged the development of sophisticated tools such as JAX (Bradbury et al., ), a high-performance numerical computing library with automatic differentiation and JIT compilation. A rich ecosystem of specialized libraries has emerged around JAX, such as Flax (Heek et al., ) for neural networks, RLax (DeepMind et al., ) for reinforcement learning, and EvoSax (Lange, ), EvoJax (Tang et al., ) and QDax (Chalumeau et al., ) for evolutionary algorithms.

In the realm of cellular automata, there have been efforts to implement specific CA models using JAX. For instance, EvoJax (Tang et al., ) and Leniax (Giraud, ) both provide a hardware-accelerated Lenia implementation. Biomaker CA (Randazzo & Mordvintsev, ), a specific CA model focusing on biological pattern formation, further demonstrates the potential of JAX in CA research. Finally, various GitHub repositories replicate results from neural cellular automata papers, but these implementations are typically narrow in focus. Recent advancements in continuous cellular automata research have also benefited from JAX-based implementations. These include Lenia (Chan, ) and Leniabreeder (Faldor & Cully, ), which have enabled large-scale simulations of open-ended evolution in continuous cellular automata (Chan, ).

While existing implementations demonstrate JAX’s potential in CA research, they also reveal significant gaps in the field. Current tools are often specialized for specific CA types (e.g., discrete, 1- and 2-dimensional), narrow in focus (e.g., replicating specific neural CA papers), or lack hardware acceleration. This limitation underscores the need for a comprehensive, flexible, and efficient library that can handle a broad spectrum of CA types while leveraging hardware acceleration. CAX aims to address this gap by providing a versatile, JAX-based tool to accelerate progress across the entire landscape of cellular automata research.

## 3 CAX: Cellular Automata Accelerated in JAX

Figure: Figure 2: High-level architecture of CAX, illustrating the modular design with perceive and update components. This flexible structure supports various CA types across multiple dimensions. (Adapted from Mordvintsev et al. () under CC-BY 4.0 license.)
Refer to caption: x2.png

CAX is a high-performance and flexible open-source library designed to accelerate cellular automata research. In this section, we detail CAX’s architecture, design and key features.
At its core, CAX leverages JAX and Flax (Heek et al., ), capitalizing on the well-established connection between CA and recurrent convolutional neural networks. This synergy, discussed in Section 2), allows CAX to harness advancements in machine learning to accelerate CA research.
CAX offers a modular and intuitive design through a user-friendly interface, supporting both discrete and continuous cellular automata across any number of dimensions. This flexibility enables researchers to seamlessly transition between different CA types and complexities within a single, unified framework (Table 1).
We have made our anonymized repository available at [github.com/b769eb6f/cax](https://github.com/879f4cf7/cax). We invite readers to experience CAX’s capabilities firsthand by accessing our curated examples as interactive notebooks in Google Colab, conveniently linked in the repository’s README.

### 3.1 Architecture and Design

CAX introduces a unifying framework for all cellular automata types, encompassing discrete, continuous, and neural models across any number of dimensions (Table 1). This flexible architecture is built upon two key components: the perceive module and the update module. Together, these modules define the local rule of the CA. At each time step, this local rule is applied uniformly to all cells in the grid, generating the next global state of the system, as explained in Section 2.1. This modular approach not only provides a clear separation of concerns but also facilitates easy experimentation and extension of existing CA models.

The architecture of CAX allows for easy composition of different perceive and update modules, enabling the creation of a wide variety of cellular automata models. This modular design also facilitates experimentation with new types of cellular automata by allowing users to define custom perceive and update modules while leveraging the existing infrastructure provided by the library.

#### 3.1.1 Perceive module

The perceive module in CAX is responsible for gathering information from the neighborhood of each cell. This information is then used by the update module to determine the cell’s next state. CAX provides several perception mechanisms, including Convolutional Perception, Depthwise Convolutional Perception and Fast Fourier Transform Perception. The perceive modules are designed to be flexible and can be customized for different types of cellular automata.

#### 3.1.2 Update module

The update module in CAX is responsible for determining the next state of each cell based on its current state and the information gathered by the perceive module. CAX provides several update mechanisms, including MLP Update, Residual Update and Neural Cellular Automata Update. Like the perceive modules, the update modules are designed to be flexible and can be customized for different cellular automata models.

### 3.2 Features

#### 3.2.1 Performance

Figure: Figure 3: Performance benchmarks of CAX. Left: Simulation speed comparison between CAX and CellPyLib for classical cellular automata. CAX demonstrates a 1,400x speed-up for Elementary Cellular Automata and a 2,000x speed-up for Conway’s Game of Life. Right: Training speed comparison between CAX and the official TensorFlow implementation for neural cellular automata experiments. CAX achieves a 1.5x speed-up on the Self-classifying MNIST Digits task.
Refer to caption: x3.png

CAX leverages JAX’s powerful vectorization and scan capabilities to achieve remarkable speed improvements over existing implementations. Our benchmarks, conducted on a single NVIDIA RTX A6000 GPU, demonstrate significant performance gains across various cellular automata models. For Elementary Cellular Automata, CAX achieves a 1,400x speed-up compared to CellPyLib. In simulations of Conway’s Game of Life, a 2,000x speed-up is observed relative to CellPyLib.

Furthermore, in the domain of Neural Cellular Automata, specifically the Self-classifying MNIST Digits experiment, CAX demonstrates a 1.5x speed-up over the official TensorFlow implementation. These performance improvements, illustrated in Figure 3, are made possible by JAX’s efficient vectorization and the use of its scan operation for iterative computations. The following code snippet exemplifies how CAX utilizes JAX’s scan function to optimize multiple CA steps:

This optimized approach allows for rapid execution of complex CA simulations, opening new possibilities for large-scale experiments and real-time applications.

#### 3.2.2 Utilities

CAX offers a rich set of utility functions to support various aspects of cellular automata research. A high-quality implementation of the sampling pool technique is provided, which is crucial for training stable growing neural cellular automata Mordvintsev et al. (). To facilitate the training of unsupervised neural cellular automata and enable generative modeling within the CA framework, CAX incorporates a variational autoencoder implementation. Additionally, the library provides utilities for handling image and emoji inputs, allowing for diverse and visually engaging CA experiments. These utilities are designed to streamline common tasks in CA research, allowing researchers to focus on their specific experiments rather than reimplementing standard components.

#### 3.2.3 Documentation and Examples

CAX prioritizes user experience and ease of adoption through comprehensive documentation and examples. The entire library is thoroughly documented, with typed classes and functions accompanied by descriptive docstrings. This ensures users have access to detailed information about CAX’s functionality and promotes clear, type-safe code. To help users get started and showcase advanced usage, CAX offers a collection of tutorial-style interactive Colab notebooks. These notebooks demonstrate various applications of the library and can be run directly in a web browser without any prior setup, making it easy for new users to explore CAX’s capabilities.

For easy access and integration into existing projects, CAX can be installed directly via PyPI, allowing users to quickly incorporate it into their Python environments. The library maintains high standards of code quality, with extensive unit tests covering a significant portion of the codebase. Continuous Integration (CI) pipelines ensure that all code changes are thoroughly tested and linted before integration. These features collectively make CAX not just a powerful tool for cellular automata research, but also an accessible and user-friendly library suitable for both novice and experienced researchers in the field.

## 4 Implemented Cellular Automata and Experiments

To showcase the versatility and capabilities of the library, we show that CAX supports a wide array of cellular automata, ranging from classical discrete models to advanced continuous CAs and including neural implementations. In this section, we provide an overview of these implementations, demonstrating the library’s flexibility in handling various dimensions and types (Table 1).

We begin with three classic models that highlight CAX’s ability to support both discrete and continuous systems across different dimensions. The Elementary CA, a foundational one-dimensional discrete model studied extensively by Wolfram (), demonstrates CAX’s efficiency in handling simple discrete systems. Conway’s Game of Life (Gardner, ), a well-known two-dimensional model, showcases CAX’s capability in simulating complex emergent behaviors in discrete space. Lenia (Chan, ), a continuous, multi-dimensional model, illustrates CAX’s flexibility in supporting more complex, continuous systems in arbitrary dimensions.

Furthermore, we have replicated four prominent NCA experiments that have gained significant attention in the field. The Growing NCA (Mordvintsev et al., ) demonstrates CAX’s ability to handle complex growing patterns and showcases the implementation of the sampling pool technique, crucial for stable growth and regeneration. The Growing Conditional NCA (Sudhakaran et al., ) utilizes CAX’s Controllable CA capabilities, as introduced in Section 2.2 allowing for targeted pattern generation. The Growing Unsupervised NCA (Palm et al., ) highlights CAX’s versatility in incorporating advanced machine learning techniques, specifically the use of a Variational Autoencoder within the NCA framework. The Self-classifying MNIST Digits (Randazzo et al., ) showcases CAX’s capacity for self-organizing systems with global coordination via local interactions, contrasting with growth-based tasks.

These implementations not only validate CAX’s performance and flexibility but also serve as valuable resources for researchers looking to build upon or extend these models. We complement these implementations with three novel experiments, which will be detailed in the following section.

**Table 1: Overview of Cellular Automata implemented in CAX**
| Cellular Automata | Reference | Type | Dimensions |
| --- | --- | --- | --- |
| Elementary Cellular Automata | Wolfram () | Discrete | 1D |
| Conway’s Game of Life | Gardner () | Discrete | 2D |
| Lenia | Chan () | Continuous | ND |
| Growing Neural Cellular Automata | Mordvintsev et al. () | Neural | 2D |
| Growing Conditional Neural Cellular Automata | Sudhakaran et al. () | Neural | 2D |
| Growing Unsupervised Neural Cellular Automata | Palm et al. () | Neural | 2D |
| Self-classifying MNIST Digits | Randazzo et al. () | Neural | 2D |
| Diffusing Neural Cellular Automata | Section 5.1 | Neural | 2D |
| Self-autoencoding MNIST Digits | Section 5.2 | Neural | 3D |
| 1D-ARC Neural Cellular Automata | Section 5.3 | Neural | 1D |

## 5 Novel Neural Cellular Automata Experiments

### 5.1 Diffusing Neural Cellular Automata

In this experiment, we introduce a novel training procedure for NCA, inspired by diffusion models.
Traditionally, NCAs have predominantly relied on growth-based training paradigms, where the state is initialized with a single alive cell and trained to grow towards a target pattern (Sudhakaran et al., ; Mordvintsev et al., ; Palm et al., ).
However, this approach often faces challenges in maintaining stability and achieving consistent results (Mordvintsev et al., ).

Figure: Figure 4: Inspired by diffusion models, the NCA learns to denoise images over a fixed number of steps. The process evolves from pure noise (left) to a target pattern (right).
Refer to caption: x4.png

The conventional NCA training method typically employs a ”sample pool” strategy to address stability issues and encourage the formation of attractors.
This approach involves maintaining a diverse pool of intermediate states, sampling from this pool for training, and periodically updating it with newly generated states.
By exposing the NCA to various intermediate configurations and consistently guiding them towards the target pattern, the sample pool method helps shape the system’s dynamics, making the desired pattern a more robust attractor in the state space.

Our proposed diffusion-inspired approach offers several advantages over the traditional growing mechanism. First, unlike the growing mechanism, our diffusion-based approach doesn’t require a sample pool, which simplifies the training process and reduces memory requirements, making it more efficient and scalable. Second, our diffusion-inspired approach naturally guides the NCA towards more stable dynamics, effectively creating a stronger attractor basin around the target pattern. In Figure 5, we compare the regeneration capabilities of growing NCAs with diffusing NCAs. We create an artificial damage by cutting the tail of the gecko and observe that diffusing NCA demonstrate emergent regenerating capabilities.

Figure: Figure 5: Diffusing NCAs demonstrate emergent regenerating capabilities compared to growing NCAs that are unstable if not trained explicitely to regenerate and recover from damage.
Refer to caption: x5.png

### 5.2 Self-autoencoding MNIST Digits

Figure: Figure 6: The 3D NCA is initialized with an MNIST digit (left). The NCA learns to reconstruct the digit on the opposite red face (right).
Refer to caption: extracted/5899433/assets/self_autoencoding_mnist.png

In this experiment, we draw inspiration from Randazzo et al. () where a NCA is trained to classify MNIST digits through local interactions. In their work, each cell (pixel) of an MNIST digit learns to output the correct digit label through local communication with neighboring cells. The NCA demonstrates the ability to reach global consensus on digit classification, maintain this classification over time, and adapt to perturbations or mutations of the digit shape. Their model showcases emergent behavior, where simple local rules lead to complex global patterns, analogous to biological systems achieving anatomical homeostasis.

Building upon this concept, we propose a novel experiment that could be termed “Self-autoencoding MNIST Digits”. In this setup, we utilize a three-dimensional NCA initialized with an MNIST digit on one face, see Figure 6. The objective of the NCA is to learn a rule that will replicate the MNIST digit on its opposite face (red face). However, we introduce a critical constraint: in the middle of the NCA, there is a mask where cells cannot be updated, effectively preventing direct communication between the two faces. Crucially, we allow for a single-cell wide hole in the center of this mask, creating a minimal channel for information transfer.

To successfully replicate the MNIST digit on the opposite face, the NCA must develop a sophisticated rule set that accomplishes two key tasks. First, it must encode the MNIST image into a compressed form that can pass through the single-cell hole. Second, it must then decode this information on the other side to accurately reconstruct the original digit.
A notable aspect of this result is that each cell in the NCA performs an identical local update rule, contributing to the system’s overall emergent behavior.
As shown in Figure 7, the NCA successfully reconstructs MNIST digits on the red face, demonstrating its ability to encode, transmit, and decode complex visual information through a minimal channel. This experiment highlights the power of NCAs in learning complex information processing tasks using simple, uniform rules, while demonstrating CAX’s ability to support sophisticated 3-dimensional CA rules.

Figure: Figure 7: The top row shows the original digits from the test set, while the bottom row displays the corresponding reconstructions on the red face of the NCA.
Refer to caption: x6.png

### 5.3 1D-ARC Neural Cellular Automata

In this experiment, we train a one-dimensional NCA on the 1D-ARC dataset (Xu et al., ). The 1D-ARC dataset is a novel adaptation of the original Abstraction and Reasoning Corpus (Chollet, ) (ARC), designed to simplify and streamline research in artificial intelligence and language models. By reducing the dimensionality of input and output images to a single row of pixels, 1D-ARC maintains the core knowledge priors of ARC while significantly reducing task complexity. For example, the tasks in 1D-ARC include ”Static movement by 3 pixels”, ”Fill”, and ”Recolor by Size Comparison”. For a full description of the dataset, see the [project page](https://khalil-research.github.io/LLM4ARC/).

Figure: Figure 8: 1D-ARC NCA space-time diagrams for each task. The top row of pixels in each image is the input. Subsequent rows of pixels show the NCA’s intermediate steps as it attempts to transform the input into the target. The bottom row of pixels represents the NCA’s final output after a fixed number of steps, which is compared to the target for task completion.
Refer to caption: x7.png

Our experiment focuses on training an NCA to solve the 1D-ARC tasks. Each input sample consists of a single row of colored pixels and a corresponding target row. The NCA’s objective is to transform the input into the target through successive applications of its rule. We consider a task successful if all pixels in the NCA’s output match the target pixels after a predetermined fixed number of steps.

**Table 2: GPT-4 and NCA accuracy in percentage on all tasks from the 1D-ARC test set. The GPT-4 values are direct-grid approach, directly taken from Xu et al. ().**
| Task | GPT-4 | NCA |
| --- | --- | --- |
| Move 1 | 66 | 100 |
| Move 2 | 26 | 100 |
| Move 3 | 24 | 100 |
| Move Dynamic | 22 | 12 |
| Move 2 Towards | 34 | 98 |
| Fill | 66 | 66 |
| Padded Fill | 26 | 28 |
| Hollow | 56 | 98 |
| Flip | 70 | 28 |
| Mirror | 20 | 6 |
| Denoise | 36 | 100 |
| Denoise Multicolor | 60 | 58 |
| Pattern Copy | 36 | 100 |
| Pattern Copy Multicolor | 38 | 100 |
| Recolor by Odd Even | 32 | 0 |
| Recolor by Size | 28 | 0 |
| Recolor by Size Comparison | 20 | 0 |
| Scaling | 88 | 88 |
| Total | 41.56 | 60.12 |

The primary goal of this experiment is for the NCA to learn from the training set, a generalizable rule for each task that can solve unseen examples from the test sets. This challenge tests the NCA’s ability to infer abstract patterns and apply them to new situations, a key aspect of human-like reasoning. Figure 8 illustrates the NCA’s “reasoning” on all 1D-ARC tasks. The visualization shows the input at the top, intermediate steps, and final output of the NCA at the bottom of each image, and is called a space-time diagram.

To evaluate the NCA’s performance, we compare it to GPT-4, a state-of-the-art language model, on the 1D-ARC test set. Table 2 presents the accuracy of the NCA and GPT-4 across 18 different task types. The GPT-4 values are direct-grid results, directly taken from Xu et al. () Appendix A. Notably, the NCA outperforms GPT-4 on several tasks, particularly those involving movement, pattern copying, and denoising. Overall, the NCA achieves a total accuracy of 60.12% compared to GPT-4’s 41.56%, as reported by Xu et al. ().

These results demonstrate the potential of NCAs in solving abstract reasoning tasks, even outperforming sophisticated language models in certain domains. The NCA’s success in tasks like ”Move 3” and ”Pattern Copy Multicolor” showcases its ability to learn complex spatial transformations and apply them consistently.

However, the NCA struggles with tasks involving more abstract concepts like odd-even distinctions or size comparisons. This limitation suggests areas for future improvement, possibly through the integration of additional priors or more sophisticated architectures. While the average of NCA outperforms GPT4, it is interesting to note that GPT4 performs equally in every task, while NCA completely fails on some of them (0% accuracy). This opens interesting questions for future work. This experiment not only highlights the capabilities of NCAs in abstract reasoning tasks but also demonstrates CAX’s flexibility in implementing and training NCA models for diverse applications.

## 6 Conclusion

In this paper, we introduce CAX: Cellular Automata Accelerated in JAX, an open-source library, designed to provide a high-performance and flexible framework to accelerate cellular automata research. CAX provides substantial speed improvements over existing implementations, enabling researchers to run complex simulations and experiments more efficiently.

CAX’s flexible architecture supports a wide range of cellular automata types across multiple dimensions, from classic discrete models to advanced continuous and neural variants. Its modular design, based on customizable perceive and update components, facilitates rapid experimentation and development of novel CA models, enabling efficient exploration of new ideas.

CAX’s comprehensive documentation, example notebooks, and seamless integration with machine learning workflows not only lower the barrier to entry but also promote reproducibility and collaboration in cellular automata research. We hope this accessibility will accelerate the pace of discovery by attracting new researchers.

In the future, we envision several exciting directions, such as expanding the model zoo to implement and optimize a wider range of cellular automata models, and exploring synergies between cellular automata and other approaches, such as reinforcement learning or evolutionary algorithms.

## Appendix A Hyperparameters

In this section, we detail the hyperparameters for the three novel neural cellular automata experiments presented in Section 5.

**Table 3: Diffusing Neural Cellular Automata**
| Parameter | Value |
| --- | --- |
| Spatial dimensions | $(72,72)$ |
| Channel size | $64$ |
| Number of kernels | $3$ |
| Hidden size | $256$ |
| Cell dropout rate | 0.5 |
| Batch size | $8$ |
| Number of steps | $128$ |
| Learning rate | $0.001$ |

**Table 4: Self-autoencoding MNIST Digits**
| Parameter | Value |
| --- | --- |
| Spatial dimensions | $(16,16,32)$ |
| Channel size | $32$ |
| Number of kernels | $4$ |
| Hidden size | $256$ |
| Cell dropout rate | 0.5 |
| Batch size | $8$ |
| Number of steps | $96$ |
| Learning rate | $0.001$ |
| Pool size | $1,024$ |

**Table 5: 1D-ARC Neural Cellular Automata**
| Parameter | Value |
| --- | --- |
| Spatial dimensions | $(128)$ |
| Channel size | $32$ |
| Number of kernels | $2$ |
| Hidden size | $256$ |
| Cell dropout rate | 0.5 |
| Batch size | $8$ |
| Number of steps | $128$ |
| Learning rate | $0.001$ |

## Appendix B Example Notebook

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: None.
soundness: 2
strengths: The issue the library tackles is impactful. CAs play a huge role in various areas of research and having efficient implementations is, of course, key to experimenters. Dramatic improvements in performance have been known to enable a way broader research community in other fields.

The library also unifies various scenarios in which CAs are used that have been fragmented over multiple platforms.

The foundation in JAX also seems to allow for comparatively easy integration with other approaches in the field of AI. (A discussion of that point might be helpful.)

The libraries capabilities are demonstrated on a well-selected range of applications.

Various interesting novel applications are defined and well presented in the paper.
summary: The paper present a new library that can model several types of cellular automata (CAs), ranging from simply Wolfram-style automata and Conway's game of life to continuous and/or complex neural cellular automata. This library is based on JAX and can thus utilize modern computation hardware better, leading to significant performance increases in some cases as the authors show on a few common examples for CA applications.
weaknesses: The paper stresses some standard library features (like documentation or examples) too much compared to more research-relevant parts.

The performance comparison should be broader. Various libraries exist. Why were only two tested in only very specific cases?

The performance difference to TensorFlow is left unexplained. TF should use hardware acceleration as well, so what is happening there?

The figures are not properly referenced in the text.

Some formal problems exist:
- superfluous ")" in "Section 2)" on page 4
- ill-formatted citation of "Mordvintsev et al." on page 6
- many wrong quotation marks on pages 8-10
- "doesn't" on page 8
- "a NCA" on page 8
- superfluous comma "training set," page 10
- ill-formatted reference to Appendix A on page 10

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 4
questions: Are the mentioned works feasible or simple to implement in CAX? Specifically ViTCA is interesting and could lead to an application of NCA on the full ARC set, but is the attention mechanism possible to implement in CAX?

How does the damage repair example of Figure 5 incorporate diffusion? Do you renoise the "damaged" image or is it simply through training on the diffusion process that it is able to recreate the original image? Could this process be used to train a generative NCA over a distribution of images instead of a single one?
soundness: 3
strengths: The article is clearly written and makes a compelling argument for cellular automata. It demonstrates their utility in image analysis and even reasoning tasks. The use of GPU acceleration through JAX is timely and relevant as a new exploration of GPU acceleration capabilities. As the main contribution is a software library, it provides some experimental comparisons with other open-source libraries.

The secondary contribution of the article are three novel applications of NCA. The most concrete of these is the application to the 1D ARC dataset. This demonstrates that NCA can follow patterns from a small number of examples, allowing them to rival GPT4's in-context learning on visual reasoning tasks. To my knowledge, this is the first study of NCA for reasoning; the demonstration that NCA can do low-data pattern matching, whether or not that is "reasoning", is very interesting.
summary: This paper presents a software library for cellular automata with GPU acceleration through the JAX library. It demonstrates this new library on MNIST digit classification and the 1D version of the Abstract Reasoning Challenge
weaknesses: A question I had in reviewing this was its relevancy to ICLR. It isn't the standard ICLR paper, so I consulted ICLR's official CFP, which includes:

+ generative models
+ causal reasoning
+ infrastructure, software libraries, hardware, etc.

I believe the article fits into these categories, however the match to ICLR and the general machine learning literature could be strengthened. For example, these works use NCA in more traditional machine learning settings:

Grattarola, Daniele, Lorenzo Livi, and Cesare Alippi. "Learning graph cellular automata." Advances in Neural Information Processing Systems 34 (2021): 20983-20994.

Tesfaldet, Mattie, Derek Nowrouzezahrai, and Chris Pal. "Attention-based neural cellular automata." Advances in Neural Information Processing Systems 35 (2022): 8174-8186.

Pajouheshgar, Ehsan, et al. "Dynca: Real-time dynamic texture synthesis using neural cellular automata." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2023.

None of these are cited, despite being high-profile studies of NCA in machine learning (NeurIPS, CVPR). As the main contribution of this article is a software library, it would be very valuable to know if these studies or similar ones could be reproduced using CAX. For submission to ICLR, this article could do more to relate the contribution to machine learning.

The second major weakness is the limited experimental study. The studies proposed in section 5 function more as demonstrations of the CAX library than new experiments due to their lack of rigor. Only the 1D-ARC experiment presents quantified results and a comparison with another method. The novel diffusion method presented in 5.1 is only demonstrated on a single example, and the self-encoding MNIST digits experiment in 5.2 is only validated through visual comparison of 1 image per digit. As such, claims like " As shown in Figure 7, the NCA successfully reconstructs MNIST digits on the red face, demonstrating its ability to encode, transmit, and decode complex visual information
through a minimal channel" are not fully supported. Experimental details and full reporting of experimental results are necessary to support such claims. As demonstrations of the CAX library, these examples are interesting, but as experiments, they fall short.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 4
questions: * Why not run proper hard and soft scaling tests in the performance comparison? It is always hard to show few numbers in a speed comparison. 
* Further, is the performance improvement only due to JAX or have you also written the code in a way that further improves parts of the algorithms? If so, that would make a key result in the paper and improve its novelty.
soundness: 3
strengths: The paper is very well written and easy to read. The authors have made presentation a clear priority in the paper and should be commended. 

The software itself is also a great idea. Reducing the time it takes for people to start working on these kinds of studies and to do so in an accelerated, modern way is very useful and this work presents a framework doing exactly that.
summary: The paper presents a new open-source library for performing Cellular Automata studies written in the JAX framework. The use of JAX allows for significant acceleration of such studies, and the framework reduces the amount of boilerplate code scientists need to write when starting such an investigation. 

Several benchmarks are shown to demonstrate the structure of the software API and even demonstrate a result when compared with common language models.
weaknesses: While the paper presents a valuable tool, it does so in a very surface-level way. Most software discussions are in the API demonstration, which shows users how to put it into practice. This is important, but in an academic paper presenting the framework, the framework should take centre stage. Further, benchmarks against common test systems are very welcome in a software paper, but the comparison with GPT-4, while certainly interesting, is more of a research result in and of itself and not directly related to the software. I am sure there is more to the code that replacing numpy with JAX, but that should be laid out and explained in the manuscript. 

My biggest concern is that ICLR isn't quite the correct platform for such a paper. I understand the desire for it to go there as it is related to ML but it is, at the end of the day, software. If the paper were to present a novel approach to structuring CA algorithms and packaging that up, it would be a different story. However, this decision is likely up to conference organizers. I can see the benefit of using ICLR as a platform to spread this software and perhaps help a lot of scientists. But from a newness perspective, it is not a breakthrough.

---

confidence: 5
contribution: 4
initial_rating: 8
presentation: 3
questions: I don't have any major questions regarding this package other than the reason for the import crashes in the notebook I listed in the weaknesses.
soundness: 4
strengths: Originality: While this paper does not necessarliy re-invent new methods, they have created an original and unified tool to be used by practitioners of cellular automata in a highly efficient way for a broad set of possible implementations. This is a tool that will surely accelerate the study of emergence and self-organization in cellular automata.

Quality/Clarity: The paper is clearly written with sufficient background context to motivate the necessity of this tool as well as citing the relevant literature that has sparked recent interest in deeper neural cellular automata architectures.

Significance: Related to the originality of this tool, it is significant for practitioners of this field and lowers the barrier of entry for researchers to produce highly efficient code and run larger scale simulations and experiments. As the field of cellular automata is deeply intertwined with the study of emergence, scaling such simulations is a crucial requirement to push the field to new grounds. With such a tool, such experiments become much more feasible and cheaper for researchers.
summary: The authors have created an open-sourced python library called CAX (cellular automata accelerated in JAX) that implements a number of recent developments/experiments in the field of cellular automata as JAX models which allows for highly efficient and parallel computing. The authors also implement tools to customize these implementation for future researchers to run their own experiments in order to have a unified library to be used for efficient CA computations. Furthermore, they implement 3 new experiments with this library to demonstrate these customization tools (a diffusion CA model, an auto-encoding CA model, and neural CA model trained on 1D-ARC tasks).
weaknesses: - the performance increases plotted in figure 3 I feel are not sufficiently demonstrating the utility of such a library. I say this as a user of JAX and PyTorch and know that there are many performance increases when using JAX, and so there needs to be a much more comprehensive case being made here that JAX is actually introducing large speed ups. This is also increasingly important as there doesn't seem to be a massive speed up in the "Growing cellular automata" experiment, so perhaps the authors should compute performance changes for a larger set of experiments. (For example for the [self-organizing textures](https://distill.pub/selforg/2021/textures/) experiment, or [cellular automata controlling cartpole](https://arxiv.org/abs/2106.15240) experiment. 

- as someone that has had to write JAX code for my own custom CA implementations, it would be great if the appendix included some annotated code that explains how efficient JAX implementations can be written. I know that curious readers can always read the code on github, and there are some examples shown such as on line 281-289 as well as the example notebook in the appendix (but it has minimal comments), but it would be extremely helpful for new programmers to understand the thought process for writing such a library. This is doubly important because when trying to implement this library for custom experiments, a certain degree of programming in JAX will be required. This is possibly slightly outside the scope or responsibility of the authors, but I think it would go a long way for allowing researchers to adopt this library if there were tutorials or detailed annotations of code that explain this process (I know for myself this would have been very helpful a few years ago).

- I very quickly tried to run one of the colab notebooks for the 1D-ARC tasks and had problems with importing. I did not test with all the notebooks but this should probably be double-checked across the board. I recieved two errors:
   - (on the second cell: `%pip install "cax[examples] @ git+https://github.com/879f4cf7/cax.git"`): after installing most packages, I recieved the following message: `ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
gcsfs 2024.10.0 requires fsspec==2024.10.0, but you have fsspec 2024.9.0 which is incompatible.`
    - (on the third cell: `from cax.core.update.residual_update import ResidualUpdate`): 
~~~
---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
<ipython-input-3-7bc79b5c20c3> in <cell line: 11>()
      9 from cax.core.perceive.depthwise_conv_perceive import DepthwiseConvPerceive
     10 from cax.core.perceive.kernels import grad_kernel, identity_kernel
---> 11 from cax.core.update.residual_update import ResidualUpdate
     12 from flax import nnx
     13 from tqdm.auto import tqdm

1 frames
/usr/local/lib/python3.10/dist-packages/cax/core/update/mlp_update.py in <module>
      5 import jax.numpy as jnp
      6 from flax import nnx
----> 7 from flax.nnx.nnx.nn import initializers
      8 from flax.nnx.nnx.nn.linear import default_kernel_init
      9 

ModuleNotFoundError: No module named 'flax.nnx.nnx'
~~~

Meta-review: **Summary:**
This paper proposes CAX, which is an open-source python library for cell automaton (CA) research. It makes use of GPU acceleration provided by the JAX library, which would enable significant speeding-up of CA simulations.

**Strengths:**
The library has implemented several different types of CA, ranging from classic ones (such as elementary CA and Conway's game of life), through a variety of neural CAs, to quite new instances (those showcased in Section 5). Massive speedup against two benchmarks has also been demonstrated in Section 3.2.1.

**Weaknesses:**
Some reviewers (5JEh, BH4L, 7WC1) raised concern about relative weakness of the research-relevant contents of this paper compared with its documentational nature.

**Reasons:**
Despite the concern on research-relevant contents of this paper, all the reviewers agreed that the strengths of this paper outweigh the weakness, recommending its acceptance.

Decision: Accept (Oral)