Research map / Architectures and efficiency / RNNs and normalization
CAX: Cellular Automata Accelerated in JAX
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces CAX, a JAX-based open-source library for hardware-accelerated discrete and continuous cellular automata and neural cellular automata, with speedups up to 2,000x and three novel NCA experiments.
Problem
Cellular automata research lacks a hardware-accelerated, flexible, unified library, which makes simulations of CA and neural CA computationally expensive, fragmented and hard to reproduce.
Approach
A modular JAX/Flax framework that decomposes any CA (discrete, continuous, neural; 1D to ND) into perceive and update modules, using JAX vectorization and scan for speed, and demonstrated with new experiments such as diffusing NCA, self-autoencoding MNIST, and 1D-ARC NCA.
Key results
- CAX achieves a 1,400x speed-up over CellPyLib for Elementary Cellular Automata and a 2,000x speed-up for Conway's Game of Life.
- On the 1D-ARC test set, the 1D NCA reaches 60.12% total accuracy versus 41.56% for GPT-4 (direct-grid results taken from prior work).
- CAX achieves a 1.5x training speed-up over the official TensorFlow implementation on Self-classifying MNIST Digits.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use CAX instead of CellPyLib | Elementary Cellular Automata simulation · Elementary Cellular Automata | speed: 1,400x speed-up | Helped |
| use CAX instead of CellPyLib | Conway's Game of Life simulation · Conway's Game of Life | speed: 2,000x speed-up | Helped |
| use CAX instead of the official TensorFlow implementation | Self-classifying MNIST Digits experiment · Self-classifying MNIST Digits | training speed: 1.5x speed-up | Helped |
| use a 1D-ARC NCA instead of GPT-4 | 1D-ARC · 1D-ARC test set | accuracy: 41.56 -> 60.12 | Helped |
| use a 1D-ARC NCA instead of GPT-4 | 1D-ARC · 1D-ARC test tasks | accuracy | Mixed |
Surprising findings
- A simple 1D neural cellular automaton outperforms GPT-4 on average on 1D-ARC, yet scores 0% on several recoloring tasks.
- The 1D NCA fails completely on tasks requiring odd-even or size-comparison distinctions.
Limitations
- The NCA fails completely (0% accuracy) on several 1D-ARC tasks such as recolor by odd-even and size comparison.
- The GPT-4 1D-ARC numbers are taken from prior work rather than re-run.
- Speed-up comparisons are against a limited set of baselines (CellPyLib, one TensorFlow implementation).
Experimental setup
Datasets: MNIST 1D-ARC
Benchmarks: 1D-ARC MNIST
Metrics: accuracy speed-up factor training time
Baselines: CellPyLib official TensorFlow NCA implementation GPT-4
Builds on: JAX Flax CellPyLib Growing Neural Cellular Automata Lenia
Closest papers by meaning
- Learning Graph Cellular Automata (NeurIPS 2021 · similarity 0.85)
- JaxMARL: Multi-Agent RL Environments and Algorithms in JAX (NeurIPS 2024 · similarity 0.81)
- The Automated LLM Speedrunning Benchmark: Reproducing NanoGPT Improvements (NeurIPS 2025 · similarity 0.81)
- Deep Neural Cellular Potts Models (ICML 2025 · similarity 0.81)
- FlashEvolve: Accelerating Agent Self-Evolution with Asynchronous Stage Orchestration (NeurIPS 2026 · similarity 0.80)
- Bridging Expressivity and Scalability with Adaptive Unitary SSMs (NeurIPS 2025 · similarity 0.80)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.