CAX: Cellular Automata Accelerated in JAX
Maxence Faldor, Antoine Cully
TL;DR
CAX is a high-performance, flexible open-source library that accelerates cellular automata simulations in JAX, enabling faster research and new applications across scientific disciplines.
Contribution
We introduce CAX, a novel hardware-accelerated, modular library for cellular automata in JAX, supporting both discrete and continuous models in arbitrary dimensions.
Findings
CAX speeds up simulations by up to 2,000 times.
CAX enables rapid implementation of complex cellular automata experiments.
A simple 1D cellular automaton can outperform GPT-4 on the 1D-ARC challenge.
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…
Peer Reviews
Decision·ICLR 2025 Oral
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 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
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
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,
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.
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 direc
Code & Models
Videos
Taxonomy
TopicsCellular Automata and Applications · Simulation Techniques and Applications
MethodsAttention Is All You Need · Dense Connections · Adam · Linear Layer · Residual Connection · Position-Wise Feed-Forward Layer · Label Smoothing · Dropout · Byte Pair Encoding · Absolute Position Encodings