CAX: Cellular Automata Accelerated in JAX

TL;DR

CAX addresses the need for a hardware-accelerated, flexible CA framework by delivering a JAX-based library that unifies discrete, continuous, and neural CA across dimensions. Its modular architecture (Perceive/Update) enables rapid experimentation and seamless integration of external inputs, while achieving substantial speedups (up to ~$2{,}000\times$) over prior implementations. The work demonstrates broad applicability through classic models (Elementary CA, Life), continuous CA (Lenia), and neural CA tasks, including three novel experiments and a GPT-4 comparison on 1D-ARC, underscoring the potential for large-scale, reproducible CA research. By coupling high performance with rich documentation and tutorials, CAX lowers barriers to exploration and accelerates advancement in emergent computation and neural CA research.

Abstract

Cellular automata have become a cornerstone for investigating emergence and self-organization across diverse scientific disciplines. 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 delivers cutting-edge performance through hardware acceleration while maintaining flexibility through its modular architecture, intuitive API, and support for both discrete and continuous cellular automata in arbitrary 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.

Figures (8)

• Figure 1: Cellular Automata types supported in CAX.
• Figure 2: High-level architecture of CAX, illustrating the modular design with perceive and update modules. This flexible structure supports various CA types across multiple dimensions. (Adapted from mordvintsev2020growing under CC-BY 4.0 license.)
• 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.
• 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).
• 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.