ICLR2025 | CAX：基于JAX加速的元胞自动机

Cellular automata 已成为跨多个科学研究涌现现象和自组织的基石。然而，由于缺乏硬件加速的元胞自动机库，限制了新研究方向的探索，阻碍了合作，也不利于研究的可重复性。从而推出了CAX（基于JAX加速的元胞自动机），这是一个高性能且灵活的开源库，旨在加速元胞自动机的研究。CAX通过硬件加速实现了前沿的性能表现，同时凭借其模块化架构、直观的应用程序编程接口（API），以及对任意维度的离散和连续元胞自动机的支持，保持了高度的灵活性。

---

个人观点：元胞自动机应该是通过定义个体相互作用的局部规则，用到全局元胞上展现出系统的演化，在数据驱动时代下，如果真正要应用，应该是可自己学习规则的元胞自动机，这是不是其实可以被GNN替代（只要定义好邻居，给消息传递设置多时间步就行）

---

来自：CAX: CELLULAR AUTOMATA ACCELERATED IN JAX, ICLR 2025
地址：https://github.com/maxencefaldor/cax

背景概述

涌现是一个基本概念。当一个复杂实体展现出其组成部分单独并不具备的属性时，这种引人入胜的现象就会出现。从蚁群的集体智慧到雪花的形成，自组织和涌现以无数种方式呈现。对自组织和涌现的研究有望揭开从生命起源到意识发展等诸多深层次的谜团。

元胞自动机（CA）是一种计算模型，它展示了复杂的模式和复杂的行为如何从简单的组件通过基本规则相互作用而产生。CA起源于20世纪40年代乌拉姆（Ulam）和冯·诺伊曼（von Neumann）的研究，在20世纪70年代随着康威的生命游戏以及20世纪80年代斯蒂芬·沃尔夫勒姆（Stephen Wolfram）的系统研究而受到广泛关注。甚至基本CA也具有图灵完备性这一发现，凸显了它们的表达能力。CA是研究自组织和涌现现象的有力工具，有助于理解从物理学、生物学到计算机科学和 AI 生命等多个科学领域的复杂现象。

近年来，机器学习技术与 CA 的融合为形态发生和自组织的研究开辟了新途径。神经-元胞自动机（NCA）是该领域的一项关键进展，它使用可微神经网络通过基于梯度的优化来学习更新规则，而不是依赖手动设计。最近的研究将NCA与CNN、GNN和ViT等先进架构相结合。机器学习和元胞自动机研究的融合，不仅加深了我们对复杂系统的理解，也凸显了CA实验日益增长的计算需求。

尽管元胞自动机在概念上很简单，但它的模拟计算量可能很大，尤其是在扩展到具有大量细元胞的更高维度，或者为 NCA 进行时间反向传播时。此外，在研究环境中，CA 的实现往往是碎片化的，个别研究人员经常重新实现基本功能，在诸如TensorFlow、JAX和PyTorch等各种深度学习框架中进行定制化实现。随着该领域的不断发展并吸引越来越多的关注，迫切需要一个统一、强大的库，以促进合作、提高研究的可重复性、加快实验进程，并助力探索新的研究方向。

CAX（基于JAX加速的CA），这是一个具有前沿性能的开源库，旨在为CA研究提供灵活且高效的框架。JAX是一个高性能数值计算库，它能通过在CPU、GPU和TPU等各种硬件加速器上进行大规模并行计算，加快 CA 的模拟速度。CAX十分灵活，支持任意维度的离散和连续CA，既适用于基本CA和康威生命游戏等经典模型，也能兼容像Lenia和NCA这样的现代变体（图1）。

JAX能够对 CA 规则进行高效向量化处理，使数百万个元胞的更新能够同时进行。它还具备自动求导功能，可以有效地进行时间反向传播，这为 NCA 的训练提供了便利。在基准测试中，CAX能够在几分钟内完成涉及数百万元胞更新的实验，与传统实现方式相比，计算时间最多可缩短2000倍。CAX的灵活性和加速研究的潜力通过三个新颖的 CA 实验得以展现。得益于CAX的模块化架构，这些实验中的每一个都只需几行代码就能实现，显著降低了 CA 研究的入门门槛。

CA

元胞自动机是一种简单的计算模型，由规则排列的元胞网格组成，每个元胞（cell）都处于特定状态。网格可以是任意有限维度。对于每个cell，相对于该指定cell定义了一组被称为其邻域的cells。网格在离散的时间步长上根据固定规则进行更新，该规则依据每个cell的当前状态及其邻域内cell的状态来确定其新状态。

一个CA被定义为元组$(L,S,N,\varphi )$，其中有$c$个channels的$d$维grid。$S$是 cell 状态集合。$N\subset L$为邻域，$\varphi :S^N\to S$是局部规则。从grid到cell状态集的映射是$\text{S}:L\to S$，被称为配置或者模式。一个cell被记为$\text{x}$。$\text{S}(\text{x})$表示一个cell $\text{x}\in L$的状态。记$\text{x}$的邻居为 N x = { x + n , n ∈ N } N_{\textbf{x}}=\left\{\textbf{x}+\textbf{n},\textbf{n}\in N\right\} Nx​={x+n,n∈N}， S ( N x ) = { S ( n ) , n ∈ N x } \textbf{S}(N_{\textbf{x}})=\left\{\textbf{S}(\textbf{n}),\textbf{n}\in N_{\textbf{x}}\right\} S(Nx​)={S(n),n∈Nx​}。

---

channel的意义：

• 在细胞状态描述中的作用：当提到细胞自动机的网格由“具有$c$个通道的$d$维晶格或网格”组成时，意味着每个细胞除了其在$d$维空间中的位置信息外，还有$c$个不同的状态维度。例如在图像处理的细胞自动机应用中，一个像素（可类比为细胞）除了表示颜色的RGB值（这可看作一种状态），还可能有透明度等其他属性，这些额外的属性就可以通过“通道”来表示。每个通道都可以有自己独立的状态变化规则，与其他通道相互作用，共同决定细胞的整体行为和状态演变。
• 对系统复杂度和表达能力的影响：通道数量$c$的增加能够显著提升细胞自动机的复杂度和表达能力。以简单的一维细胞自动机为例，如果只有一个通道，细胞可能仅表示简单的“开”或“关”两种状态；而增加通道后，细胞可以携带更多信息，如不同通道分别表示细胞的生长速度、繁殖概率等，使得细胞自动机可以模拟更复杂的现象，如生物的生长发育、生态系统的演变等过程。

---

全局规则$\Phi :S^L\to S^L$表示统一应用局部规则到所有cells。对于所有$\text{x}$，$\Phi (\text{S})(\text{x})=\varphi (\text{S}(N_{\text{x}}))$。一个元胞自动机被初始化为状态${\text{S}}_0$，然后，状态在每个时间步$t$更新：$${\text{S}}_1=\Phi ({\text{S}}_0),{\text{S}}_2=\Phi ({\text{S}}_1),...$$

可控CA：CCA

可控元胞自动机（CCA）是CA的一种推广形式，它具备在每个时间步接受外部输入的能力。CCA将Sudhakaran等人（2022年）在文献中提出的目标导向神经细胞自动机概念进行了形式化。外部输入能够改变CCA的行为，使其在遵循细胞自动机基本原理的同时，具备对变化的条件或控制信号做出动态响应的可能性。

一个CCA定义为元组$(L,S,I,N,\varphi )$，其中$I$是输入集合，并且$\varphi :S^N\times I^N\to S$是可控局部规则。一个从grid到输入集合的映射$\text{I}:L\to I$被称为输入。$\text{I}(\text{x})$表示给$\text{x}$的输入。类似于状态，记 I ( N x ) = { I ( n ) , n ∈ N x } \textbf{I}(N_{\textbf{x}})=\left\{\textbf{I}(\textbf{n}),\textbf{n}\in N_{\textbf{x}}\right\} I(Nx​)={I(n),n∈Nx​}。

全局控制规则为$\Phi :S^L\times I^L\to S^L$。对于所有$\text{x}$，$\Phi (\text{S},\text{I})(\text{x})=\varphi (\text{S}(N_{\text{x}}),\text{I}(N_{\text{x}}))$。一个可控CA被${\text{S}}_0$初始化，每个$t$有输入${\text{I}}_t$，更新为：$${\text{S}}_1=\Phi ({\text{S}}_0,{\text{I}}_0),{\text{S}}_2=\Phi ({\text{S}}_1,{\text{I}}_1),...$$

细胞自动机可以被概念化为循环卷积神经网络。然而，传统的细胞自动机缺乏在每个时间步接受外部输入的能力。可控细胞自动机通过使其能够对外界输入做出响应，扩展了传统细胞自动机的功能，这类似于循环神经网络处理序列数据。可控细胞自动机弥合了循环卷积神经网络和元胞自动机之间的差距，为对既表现出自主涌现行为又能对外界控制做出响应的复杂系统进行建模开辟了新的可能性。

相关工作

元胞自动机领域催生了大量工具和库，以支持相关研究与实验。其中，CellPyLib 成为最受欢迎且功能最为多样的选择之一。这个Python库为处理一维和二维元胞自动机提供了简洁而强大的接口，支持离散和连续状态，是比较研究和进一步开发的理想基准。尽管它实现了诸如康威生命游戏等经典元胞自动机模型，但CellPyLib不支持硬件加速，也无法用于训练NCA。Golly是一款跨平台应用程序，用于探索康威生命游戏及许多其他类型的元胞自动机。Golly的功能包括三维元胞自动机规则、自定义规则加载，以及通过Lua或Python进行脚本编写。虽然Golly在传统元胞自动机研究方面功能强大且用途广泛，但它并非为硬件加速或与现代机器学习框架集成而设计。

JAX是一个高性能的数值计算库，具备自动求导和即时编译功能。在CA领域，已经有人尝试使用JAX来实现特定的CA模型。例如，EvoJax和Leniax都提供了经过硬件加速的Lenia模型的实现方式。GitHub上的各种代码库重现了NCA相关论文中的研究成果，但这些实现方式通常关注点较为狭窄。近年来，连续CA研究的进展也得益于基于JAX的实现方式。它们使得对连续CA中开放式进化的大规模模拟成为可能。

尽管现有的实现方式展示了JAX在CA研究中的潜力，但它们也暴露出了该领域的一些重大不足。当前的工具往往专门针对特定类型的CA（例如离散型、一维和二维的CA），关注点较为狭窄（比如只是重现特定NCA相关论文中的成果），或者缺乏硬件加速功能。这种局限性凸显了对一个全面、灵活且高效的库的需求，该库应能够处理各种各样的CA类型，同时利用硬件加速技术。CAX旨在通过提供一个基于JAX的多功能工具来填补这一空白。

CAX架构

CAX为所有类型的CA引入了一个统一的框架。这个灵活的架构基于两个关键组件构建：感知模块和更新模块，见图2。这些模块共同定义了CA的局部规则。在每个时间步长，这个局部规则会统一应用于grid中的所有元胞，从而生成系统的下一个全局状态。

对于CCA，CAX架构将NCA方法推广应用到循环卷积神经网络架构中，实现了外部输入和控制信号的无缝集成。通过使用Flax中的标准机器学习组件来实现这两个模块（感知模块和更新模块），CAX使得在保持CA范式的同时，对各种神经架构进行实验变得简单直接 —— 从简单的卷积层到复杂的注意力机制均可轻松实验。

CAX的架构便于将不同的感知模块和更新模块进行组合，从而能够创建各种各样的CA模型。这种模块化设计还为新型CA的实验提供了便利，用户可以在利用该库所提供的现有基础架构的同时，自定义感知模块和更新模块。

CAX中的感知模块负责从每个cell的邻域收集信息。然后，更新模块会利用这些信息来确定该细胞的下一个状态。CAX提供了多种感知机制，包括卷积感知、深度卷积感知和快速傅里叶变换感知。这些感知模块的设计具有灵活性，并且可以针对不同类型的CA进行定制。

CAX中的更新模块负责根据每个cell的当前状态以及感知模块收集到的信息，来确定该cell的下一个状态。CAX提供了多种更新机制，包括多层感知机（MLP）更新、残差更新以及NCA更新。与感知模块一样，更新模块的设计也具有灵活性，能够针对不同的CA模型进行定制。