CAX: Cellular Automata Accelerated in JAX

Thank you for your constructive feedback. We appreciate the opportunity to address your concerns.

---

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.

Thank you for this insightful suggestion about providing more detailed code explanations. To address this need and help researchers effectively use and contribute to CAX, we have added a comprehensive CONTRIBUTING.md file to the repository. This guide includes not only standard contribution guidelines but also a dedicated section on CAX architecture design principles.

Beyond the usual content about bug reports, pull requests, and coding style, we've added a "Designing Efficient CAX Architectures" section that explains:

• How to structure CA models using perceive and update modules
• Best practices for leveraging JAX/Flax functionality
• An example structure
• A common training loop for training neural cellular automata
• Common patterns and pitfalls when implementing custom CAs

This documentation should help both new users getting started with CAX and experienced developers looking to implement custom experiments efficiently.

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 experiment, or cellular automata controlling cartpole experiment.

Thank you for this feedback about performance analysis. In CAX, cellular automata models inherit from Flax's nnx.Module and use standard components like convolution layers, linear layers, and dropout mechanisms. Once defined, these models can be trained like any other recurrent architecture. As such, CAX naturally inherits the performance characteristics and scaling properties of the JAX/Flax ecosystem, which are already well-documented. A CA model in CAX is defined through perceive and update modules. The overall performance depends on the efficiency of these modules - when they are well-implemented using optimized Flax components, the resulting CA will be efficient.

Additionally, since CAX is the first hardware-accelerated library for cellular automata, running more extensive benchmarks or scaling studies would provide limited insight due to the absence of comparable baselines. The value of CAX lies not just in its performance but in providing a structured, modular framework for implementing cellular automata with hardware acceleration.

However, we agree with the reviewer that we could compute performance changes for a larger set of experiments. Our vision for CAX includes expanding its scope to encompass a broader range of experiments, including the notable self-organizing textures implementation. While we aim to add the self-organizing textures comparison before the end of the rebuttal period, we want to be transparent that, given our concurrent commitments as authors and reviewers of multiple papers, we cannot guarantee its completion within this limited timeframe.

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.

Thank you for reporting these installation issues. We apologize for the difficulties you encountered running our notebooks.

The first error regarding fsspec versions can be safely ignored.

The second error occurred due to an API change in Flax between versions 0.9.0 and 0.10.0 that happened between the submission and the rebuttal period.

We have updated the anonymous repository to be compatible with the latest Flax version, and the notebook is now running correctly. Please, let us know if you still have any issues.

---

We believe these responses address your concerns. We welcome any further questions or clarifications you may need.

I can see the benefit of using ICLR as a platform to spread this software and perhaps help a lot of scientists.

Thank you for this comment. Indeed, we believe that presenting CAX at ICLR would benefit the research community. As the first hardware-accelerated cellular automata library, CAX represents a significant technical advance that makes sophisticated CA experiments accessible to researchers. Its publication at ICLR would help disseminate these capabilities widely and inspire future innovations in the field.

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.

Thank you for your valuable feedback. In line with your remark, we have updated the manuscript to put the framework more at the center stage of the paper. In particular, we have expanded our description of the framework and its versatility, highlighting how our perceive/update architecture enables implementation of recent important NCA papers.

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

Thank you for this feedback. We have added clarifications about these key implementation details and their contributions to performance in both the introduction and discussion sections of the paper.

The performance improvements in CAX stem from both architectural design choices and extensive leveraging of JAX/Flax capabilities that are particularly well-suited for cellular automata computations.

We extensively utilize JAX's vectorization capabilities through vmap, which is especially powerful for CA since the same rule is applied across all cells. This naturally maps to hardware acceleration patterns. Additionally, we make extensive use of Flax neural network modules (nnx.Module) like convolutions, MLPs, and dropout layers, which are highly optimized for modern hardware and use JAX's scan operator for the sequential update of the CA state.

While we build upon previous work in combining CA with machine learning architectures, we also contribute novel optimized implementations. For example, we provide an efficient JAX implementation of Sample Pooling and introduce dropout-based stochastic cell updates. By leveraging existing Flax nnx.Module components wherever possible, we benefit from their high-quality hardware-accelerated implementations.

So while JAX's capabilities are foundational to CAX's performance, our careful attention to implementation details and systematic use of optimized Flax modules further enhance the library's efficiency.

---

We hope we have answered all your questions and concerns. Please do not hesitate to bring up any additional questions.

Thank you for your thoughtful feedback and the opportunity to clarify these points.

---

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.

Thank you for this comment. We agree that the GPT-4 comparison is more of a research result and not directly related to the software evaluation. As stated in our abstract and introduction, the experiments in Section 6 were designed to demonstrate the library's capabilities, not as full-fledged research contributions. Given that the paper's primary contribution is the library itself, we prioritized explaining its architecture and design principles.

Regarding benchmarks, we provide comparisons against CellPyLib on two canonical test systems: Conway's Game of Life and Elementary Cellular Automata. While these show significant speed-ups (2000x), we acknowledge that such comparisons have limited value since CellPyLib serves different use cases and isn't hardware-accelerated. We include these mainly to give CA practitioners an indication of potential performance gains when moving to hardware-accelerated implementations.

For Neural CA experiments, we could only compare against the original papers' implementations since no unified NCA library exists. While more extensive benchmarking would be valuable, the absence of other hardware-accelerated CA libraries makes this impossible at present. We have expanded our discussion of the library's engineering aspects, particularly how we leverage Flax's nnx.Module for hardware-accelerated implementations of CA components.

For NCA experiments, we would have liked to provide a more thorough benchmark and comparison against a more extensive range of library. However, at the moment, to the best of our knowledge, no NCA library exists so we could only compare against fragmented official implementations on two NCA experiments coming from the original papers.

My biggest concern is that ICLR isn't quite the correct platform for such a paper. If the paper were to present a novel approach to structuring CA algorithms and packaging that up, it would be a different story.

First, we would like to note that "Software Libraries" is explicitly listed in ICLR's official Call for Papers, and the conference specifically welcomes papers exploring implementation issues, parallelization, software platforms, and hardware. The present work falls primarily under the area of infrastructure, software libraries, hardware, and systems.

Moreover, NCAs represent a fundamental approach to machine learning where complex behaviors emerge from learning simple, local rules and representations. This aligns directly with ICLR's core focus on representation learning - instead of learning global solutions, NCAs learn local update rules that, when applied iteratively, lead to emergent problem-solving capabilities. As a machine learning architecture, NCAs have demonstrated state-of-the-art performance in various domains, such as dynamic texture synthesis, and ViTCA architectures have shown superior performance across all benchmarks and nearly every evaluation metric on denoising autoencoding tasks.

While CAX's main contribution is a general library to implement various CA methods, the key way we enable this is precisely by introducing a novel approach to structuring CA algorithms through our controllable CA framework, as illustrated in Figure 3 of our paper.

Specifically, we introduce a fundamental architectural division between Perceive and Update modules. This separation enables strong modularity, component reuse, and clear division of responsibilities, as detailed in Section 3. The Perceive module handles neighborhood observations while the Update module determines state transitions, creating a flexible yet powerful structure that can support diverse CA types across multiple dimensions.

This architectural contribution represents a novel framework for organizing CA algorithms, making them more maintainable, reusable, and easier to implement while maintaining high performance through hardware acceleration. The table below demonstrates how this architecture effectively unifies different CA methods under a common structure:

The studies proposed in section 5 function more as demonstrations of the CAX library than new experiments...

Thank you for this comment. As stated in both our abstract and introduction, these experiments were specifically designed to demonstrate and showcase the library's capabilities: "We demonstrate CAX's performance and flexibility through a wide range of benchmarks and applications" and "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".

We acknowledge that these experiments are not intended as full-fledged research contributions, and we would be happy to revise any language in the manuscript that might suggest otherwise.

We faced a fundamental trade-off in the paper's focus and limited space. A deeper experimental analysis would have come at the expense of thoroughly presenting the library's core innovations - the controllable CA framework and the perception/update design pattern. Each concept discussed in Section 6 could merit its own paper. Given that this paper's primary contribution is the library itself rather than the experiments, we chose to prioritize explaining the framework's architecture and design principles.

While the experimental validation could be more extensive, we prioritized transparency and reproducibility. All experiments are available as notebook examples that can be run in minutes using Colab. We believe this practical accessibility, allowing readers to directly verify and build upon our work, provides significant value - potentially more than extensive experimental results without accessible code or straightforward replication paths.

We are delighted that the reviewers find these experiments interesting, and we hope they can inspire and stimulate further research in cellular automata. We look forward to seeing how the community builds upon these initial experiments to develop them into full research contributions.

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?

Thank you for this question about the damage repair example. The reconstruction occurs simply by applying the learned CA rule to the damaged image, without any renoising step. Noise is only added during training.

The key insight is that the diffusion-based training procedure naturally endows the NCA with robust regenerative capabilities. By learning to denoise images at various noise levels during training, the CA develops dynamics with strong attractor basins around target patterns. When damage occurs, these learned dynamics guide the system back toward the original pattern.

This emergent regenerative behavior is a direct consequence of the diffusion training process - we do not explicitly train for damage repair. Rather, the CA learns stable dynamics that can both denoise images and repair damage, demonstrating how diffusion training leads to robust pattern formation capabilities.

Could this process be used to train a generative NCA over a distribution of images instead of a single one?

Excellent question. While the training procedure could theoretically be extended to handle a distribution of images rather than a single target, determining whether it would yield good results remains an open question. This is an exciting direction for future research.

---

We hope we have answered all your questions and concerns. Please let us know if you have any additional questions or concerns.

Thank you for your thorough review. We appreciate the opportunity to address your concerns.

---

I believe the article fits into these categories, however the match to ICLR and the general machine learning literature could be strengthened.

Thank you for highlighting these important references. We have incorporated the suggested citations and expanded our discussion to better position CAX within the ML literature.

NCAs represent a fundamental approach to machine learning where complex behaviors emerge from learning simple, local rules and representations. This aligns directly with ICLR's core focus on representation learning - instead of learning global solutions, NCAs learn local update rules that, when applied iteratively, lead to emergent problem-solving capabilities.

Your feedback has helped us better articulate these connections and strengthen the paper's relevance to the ICLR community.

Are the mentioned works feasible or simple to implement in CAX?

We appreciate your interest in implementing advanced CA architectures. Yes, all the mentioned works are feasible and should be straightforward to implement in CAX. We have added an explanation in the manuscript saying that these papers can easily be implemented with CAX.

Since all CA in CAX inherit from the Flax's nnx.Module, implementing attention mechanisms or any other modern ML architecture is fully supported. In fact, any architecture that can be implemented in Flax can be naturally integrated into CAX. This includes attention mechanisms, graph neural networks, and other sophisticated neural architectures from the ML literature.

Our vision for the library is to serve as a comprehensive platform that continuously grows with more CA architectures and experiments, spanning from classical models to cutting-edge neural approaches. We are actively expanding our collection of implementations and welcome contributions from the community.

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?

Thank you for the comment regarding ViTCA. Indeed, ViTCA is a very interesting approach and we confirm that the attention mechanism is fully implementable within the CAX framework.

We have implemented and added a proof of concept for ViTCA in the anonymous repository. We encourage the reviewers to explore the example Colab notebook that is now available in the README.

In CAX, we implemented ViTCA through a new perceive module called ViTPerceive, available in cax/core/perceive/vit_perceive.py. The ViTPerceive class processes cells by feeding them into a multi-head self-attention mechanism with localized attention, specifically using the Moore neighborhood.

While this implementation serves as a proof of concept and doesn't implement all details from the original paper, it successfully demonstrates that ViTCA can be integrated into CAX. For transparency, our current implementation has two notable simplifications compared to the original paper:

• We do not currently use layer normalization
• We assume a one-to-one mapping between pixels and patches

These simplifications allow us to demonstrate the core concepts while maintaining clarity in the implementation. The proof of concept successfully shows that the attention mechanism can be effectively implemented within CAX's architectural constraints.

We appreciate your comments and the opportunity to address your concerns. We address each point below.

---

The paper stresses some standard library features (like documentation or examples) too much compared to more research-relevant parts.

Thank you for this feedback. As stated in both our abstract and introduction, the goal of this paper is to introduce this new library and highlight its benefits to the research community. This is why the paper emphasizes certain aspects of the library like documentation and examples, as we believe they are essential for successful adoption by researchers.

The more research-relevant parts of the paper, particularly the experiments in Section 6, were specifically designed to demonstrate and showcase the library's capabilities: "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". As such, these experiments are not intended as full-fledged research contributions. Given that this paper's primary contribution is the library itself rather than the experiments, we chose to prioritize explaining the framework's architecture and design principles.

That being said, we agree with the reviewer that the research-relevant parts are also very valuable and deserve thorough treatment. Therefore, we have expanded Section 3 to provide a more comprehensive explanation of the framework's foundations and innovative aspects.

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

Thank you for this comment. While you are right that a few CA libraries exist, to the best of our knowledge, CAX is actually the first hardware-accelerated cellular automata library.

Existing hardware-accelerated NCA implementations are always tied to specific research papers and are not designed as general-purpose libraries. Therefore, in the paper, the NCA training benchmark compares CAX against these individual paper implementations, rather than against another general-purpose NCA library, since no such library exists. These paper implementations are typically focused on reproducing specific experiments and cannot be easily adapted for other NCA tasks or research directions. The fragmented nature of existing NCA implementations makes comprehensive library comparisons impossible. While non-hardware-accelerated libraries like CellPyLib and Golly exist, as outlined in our related work section, they serve different purposes. For instance, Golly is a cross-platform cellular automata simulator but isn't designed for hardware acceleration or integration with modern machine learning frameworks.

In the paper, we compare CAX with CellPyLib on the two most canonical cellular automata - Game of Life and Elementary CA to give users a concrete indication of the speed-ups they can expect.

However, we believe that a broader comparison with non-hardware-accelerated libraries would not be informative as they serve a different purpose and would not be able to compete in terms of runtime performance due to the lack of hardware acceleration. Moreover, common CA libraries are usually limited to two-dimensional and discrete cellular automata and not designed for training neural cellular automata.

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

The performance difference between CAX and TensorFlow implementations stems largely from our optimized handling of the Sample Pool mechanism, which was originally introduced in the Growing Neural Cellular Automata paper by Mordvintsev et al.

In CAX, we provide an implementation of this mechanism that has been developed with a strong attention to runtime optimization. It avoids costly data transfers between CPU and GPU, and enables efficient batched operations through JAX's vmap and compilation through jit.

Thank you for pointing out that this was not explained in the text. We have changed the manuscript to make this point clearer.

The figures are not properly referenced in the text. Some formal problems exist...

We sincerely appreciate your meticulous attention to typographical and formal issues. Thank you for helping us improve the paper!

---

We hope we have answered all your questions and concerns. Please do not hesitate to bring up any additional questions.