CAX: Cellular Automata Accelerated in JAX

2024-10-03

arXiv.org

AI Summary

Existing cellular automata (CA) research is hindered by the absence of a general-purpose, hardware-accelerated, open-source library, limiting reproducibility, collaboration, and exploration of novel paradigms. Method: We introduce the first JAX-based, hardware-accelerated CA metacompiler—designed functionally with native JAX primitives—supporting arbitrary dimensions, discrete/continuous states, and hybrid state types. It leverages XLA compilation, vmap/pmap parallelization, and automatic differentiation. Its modular CA kernel and dynamic-dimension mechanism lower experimental barriers significantly. Contribution/Results: Three novel CA experiments—including 1D-ARC outperforming GPT-4—are implementable in just a few lines of code. On multi-task benchmarks (elementary CAs, neural CAs, MNIST self-classification), it achieves up to 2000× speedup over conventional implementations. This infrastructure enables efficient, scalable, and reproducible CA foundational research and AI–CA interdisciplinary exploration.

Technology Category

• Machine Learning: Hardware-aware ML
• Search and Optimization: Metareasoning and Metaheuristics
• Cognitive Modeling & Cognitive Systems: Agent Architectures

Application Category

• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Cellular automata have become a cornerstone for investigating emergence and self-organization across diverse scientific disciplines, spanning neuroscience, artificial life, and theoretical physics. 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 offers cutting-edge performance and a modular design through a user-friendly interface, and can support both discrete and continuous cellular automata with any number of 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.

Problem

• Lack of hardware-accelerated cellular automata library
• Hindered exploration, collaboration, and reproducibility in research
• Need for high-performance, flexible cellular automata simulation tool

Innovation

• Hardware-accelerated cellular automata library
• Modular architecture for flexible simulations
• Supports discrete and continuous automata