CAX: Cellular Automata Accelerated in JAX (2410.02651v2)

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.

• The paper delivers a JAX-based library that accelerates CA simulations up to 2,000x, significantly reducing computational barriers for large-scale experiments.
• The paper introduces a modular design with distinct perception and update modules, enabling flexible configuration for both traditional and neural cellular automata.
• The paper demonstrates novel experiments, including diffusing NCAs and 1D-ARC tasks, highlighting CAX's superior performance over conventional CA implementations.

Overview of "CAX: Cellular Automata Accelerated in JAX"

The paper introduces CAX, an open-source library engineered to expedite and enhance cellular automata (CA) research. Cellular automata have long served as a foundational model for exploring concepts of emergence and self-organization, impacting fields such as neuroscience, artificial life, and theoretical physics. The authors address a significant barrier in CA research: the absence of a versatile, hardware-accelerated library that facilitates collaboration and reproducibility.

Core Contributions

CAX stands out by delivering cutting-edge performance through JAX, a high-performance numerical computing library that leverages hardware acceleration across CPUs, GPUs, and TPUs. The library streamlines CA simulations, promising speed improvements up to 2,000 times faster compared to prior implementations. This significant performance gain empowers researchers to conduct extensive CA experiments that were previously computationally prohibitive.

The paper provides robust evidence of CAX's efficiency through various benchmarks and applications. Benchmark results indicate that for Elementary Cellular Automata and Conway's Game of Life, simulations achieve substantial accelerations. Moreover, CAX outperforms existing implementations in training neural cellular automata (NCA), enhancing both speed and feasibility of complex experimental setups.

Architectural Innovations

CAX's architecture is inherently modular, comprising two primary components: the perceive and update modules. This segmentation allows for flexible and scalable CA configurations, accommodating diverse models from simple elementary automata to sophisticated neural variants. The library supports both discrete and continuous CA types across multiple dimensions, making it highly adaptable for various research applications.

The integration with JAX is particularly noteworthy as it marries CA research with state-of-the-art advancements in machine learning, facilitating rapid prototyping and deployment of high-performance simulations. CAX’s modular design ensures that researchers can easily extend the library's capabilities by developing custom perception and update components.

Novel Experiments

The authors elaborate on several novel experiments facilitated by CAX, demonstrating the library's flexibility and potential for discovery:

1. Diffusing Neural Cellular Automata: This experiment introduces a training paradigm inspired by diffusion models. Unlike traditional growing mechanisms, the diffusion-inspired approach enhances stability and leads to robust attractor basins.
2. Self-autoencoding MNIST Digits: Utilizing a three-dimensional NCA structure, this experiment showcases the system’s ability to encode and reconstruct complex visual patterns in an innovative manner, highlighting the potential of NCAs in processing intricate information through simple, uniform rules.
3. 1D-ARC Neural Cellular Automata: The paper reveals NCAs’ capability to outperform GPT-4 on several tasks within the 1D-ARC dataset, especially those involving spatial transformations. This positions NCAs as a compelling alternative for abstract reasoning tasks, challenging contemporary AI models.

Implications and Future Directions

The implementation of CAX offers profound implications for theoretical and practical research in cellular automata. By enabling faster and more flexible CA simulations, CAX not only enhances current research methodologies but also invites exploration into novel CA types and interactions with other domains such as reinforcement learning and evolutionary algorithms.

Future work may focus on further expanding CAX’s library to encompass a wider variety of cellular automata models and optimize algorithms for additional high-dimensional experiments. Furthermore, integrating CAX with other AI techniques could yield innovative approaches to modeling complex systems, providing insights into emergent behaviors and self-organization.

In conclusion, CAX represents a significant contribution to the field of cellular automata, offering unprecedented speed, flexibility, and ease of use. Its deployment promises to accelerate research, fostering advancements across numerous scientific disciplines.