对涌现、自组织或形态发生感兴趣吗？

介绍 CAX：基于 JAX 加速的元胞自动机，一个高性能且灵活的库，旨在加速元胞自动机研究。🦎

📄：https://t.co/tOu4ejIBYu 🌟：https://t.co/0T5i2CRuJC 🧵1/12 https://t.co/Xfhgvkgti4

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

来源：https://arxiv.org/html/2410.02651 Maxence Faldor 帝国理工学院计算机系 英国伦敦 [email protected] &Antoine Cully 帝国理工学院计算机系 英国伦敦 [email protected]

摘要

元胞自动机已成为跨学科研究涌现与自组织的基石。然而，缺乏硬件加速的元胞自动机库限制了新研究方向的探索，阻碍了协作，并影响了可复现性。为此，我们提出 CAX (https://github.com/maxencefaldor/cax)（基于 JAX 加速的元胞自动机），一个高性能、灵活的开源库，旨在加速元胞自动机研究。CAX 通过硬件加速提供顶尖性能，同时通过模块化架构、直观的 API 以及对任意维度的离散和连续元胞自动机的支持保持灵活性。我们通过广泛的基准测试和应用展示了 CAX 的性能与灵活性。从基础模型如初等元胞自动机和康威生命游戏，到高级应用如生长神经元胞自动机和自分类 MNIST 数字，CAX 可将模拟速度提升高达 2000 倍。此外，我们通过展示三个新颖的元胞自动机实验来证明 CAX 加速研究的潜力，每个实验借助库的模块化架构仅需少量代码即可实现。值得注意的是，我们展示了一个简单的一维元胞自动机在 1D-ARC 挑战上超越了 GPT-4。

1 引言

涌现是一个基本概念，一直吸引着人类探究各领域（包括哲学、科学和艺术）的思想家 (Holland,2000 (https://arxiv.org/html/2410.02651v2#bib.bib16))。这一迷人现象发生在复杂实体展现出其组成部分单独不具备的属性时。从蚁群的集体智慧到雪花的形成，自组织和涌现以多种方式显现。对自组织和涌现的研究有望解开深奥的谜团，从生命起源到意识发展。

元胞自动机 (CA) 是计算模型，展示了复杂的模式和精细的行为如何从通过基本规则相互作用的简单组件中涌现。该模型起源于 20 世纪 40 年代 Ulam 和 von Neumann 的工作 (Neumann & Burks,1966 (https://arxiv.org/html/2410.02651v2#bib.bib20))，后随着 20 世纪 70 年代康威生命游戏 (Gardner,1970 (https://arxiv.org/html/2410.02651v2#bib.bib11)) 和 20 世纪 80 年代 Wolfram 的系统性研究 (Wolfram,2002 (https://arxiv.org/html/2410.02651v2#bib.bib29)) 而声名鹊起。即便是初等元胞自动机也能图灵完备，这一发现凸显了其表现力 (Cook,2004 (https://arxiv.org/html/2410.02651v2#bib.bib8))。CA 作为研究自组织和涌现的强大抽象，为理解跨科学领域（从物理学、生物学到计算机科学和人工生命）的复杂现象提供了洞见。

参见图注

图 1：CAX 支持的元胞自动机类型。

近年来，机器学习技术与元胞自动机的融合为形态发生 (Mordvintsev et al.,2020 (https://arxiv.org/html/2410.02651v2#bib.bib18)) 和自组织 (Randazzo et al.,2020 (https://arxiv.org/html/2410.02651v2#bib.bib24);2021 (https://arxiv.org/html/2410.02651v2#bib.bib25); Najarro et al.,2022 (https://arxiv.org/html/2410.02651v2#bib.bib19)) 研究开辟了新途径。神经元胞自动机 (NCA) 是该领域的关键进展，使用可微神经网络通过基于梯度的优化学习更新规则，而非依赖手工设计。近期工作已将 NCA 与高级架构结合，如卷积神经网络 (Pajouheshgar et al.,2023 (https://arxiv.org/html/2410.02651v2#bib.bib21))、图神经网络 (Grattarola et al.,2021 (https://arxiv.org/html/2410.02651v2#bib.bib14)) 和视觉 Transformer (Tesfaldet et al.,2022 (https://arxiv.org/html/2410.02651v2#bib.bib28))。机器学习与元胞自动机研究的融合不仅加深了我们对复杂系统的理解，也凸显了 CA 实验中日益增长的计算需求。

尽管概念简单，元胞自动机模拟可能计算密集，尤其是当扩展到更高维度和大量细胞时，或对 NCA 实施时间反向传播时。此外，研究环境中 CA 的实现往往碎片化：研究人员常自行重新实现基本功能，在 TensorFlow、JAX 和 PyTorch 等各类深度学习框架中创建定制实现。随着该领域的持续发展和关注度的增加，迫切需要一个统一、健壮的库，以促进协作、可复现性、快速实验和新研究方向的探索。

为应对这些挑战和机遇，我们提出 CAX：基于 JAX 加速的元胞自动机，一个开放源代码的库，具有尖端性能，旨在为元胞自动机研究提供灵活高效的框架。CAX 构建于 JAX (Bradbury et al.,2018 (https://arxiv.org/html/2410.02651v2#bib.bib2))——一个高性能数值计算库之上，能够通过 CPU、GPU 和 TPU 等各类硬件加速器的大规模并行化来加速元胞自动机模拟。CAX 灵活且支持任意维度的离散和连续元胞自动机，兼容经典模型如初等元胞自动机和康威生命游戏，以及现代变体如 Lenia 和神经元胞自动机（图 1 (https://arxiv.org/html/2410.02651v2#S1.F1)）。

JAX 提供高效的 CA 规则向量化，可同时处理数百万个细胞更新。它还提供自动微分功能，可高效地通过时间进行反向传播，便于训练神经元胞自动机。CAX 能在数分钟内运行数百万个细胞更新的实验，在我们的基准测试中，相比传统实现将计算时间减少高达 2000 倍。这一性能提升为之前计算上不可行的大规模 CA 实验开辟了新的可能性。

CAX 的灵活性和加速研究的潜力通过三个新颖的元胞自动机实验得以展示。得益于 CAX 的模块化架构，每个实验仅需少量代码即可实现，显著降低了元胞自动机研究的入门门槛。值得注意的是，我们展示了一个使用 CAX 实现的简单一维元胞自动机在 1D-ARC 挑战上超越了 GPT-4 (Xu et al.,2024 (https://arxiv.org/html/2410.02651v2#bib.bib31))，详见第 5.3 节 (https://arxiv.org/html/2410.02651v2#S5.SS3)。最后，为支持用户并促进采用，CAX 附带高质量、多样化的示例和全面的文档。已实现的 CA 列表详见表 1 (https://arxiv.org/html/2410.02651v2#S4.T1)。

2 背景

2.1 元胞自动机

元胞自动机是一种简单的计算模型，由规则的细胞网格组成，每个细胞处于特定状态。网格可以是任意有限维数。对于每个细胞，定义一组称为其邻域的细胞，相对于指定细胞。网格根据固定规则在离散时间步更新，该规则基于每个细胞的当前状态及其邻域中细胞的状态来确定其新状态。

一个 CA 由元组 (L,S,N,φ) 定义，其中 (L) 是 (d) 维格子或网格，具有 (c) 个通道，(S) 是细胞状态集，(N \subset L) 是原点的邻域，(φ: S^N → S) 是局部规则。从格子到细胞状态集的映射 (S: L → S) 称为配置或模式。在本文中，我们简称其为 CA 的“状态”。(S(x)) 表示细胞 (x \in L) 的状态。此外，我们将细胞 (x \in L) 的邻域记为 (N_x = {x+n, n \in N})，并将 (S(N_x) = {S(n), n \in N_x})。

全局规则 (Φ: S^L → S^L) 将局部规则均匀应用于格子中的所有细胞，定义为对所有 (x \in L)，(Φ(S)(x) = φ(S(N_x)))。元胞自动机以状态 (S_0) 初始化。然后，在每个离散时间步 (t \in \mathbb{N})，根据全局规则 (Φ) 更新状态，得到：

(S_1 = Φ(S_0), S_2 = Φ(S_1), \dots)

多位研究者已注意到 CA 与循环卷积神经网络之间的紧密联系 (Gilpin,2019 (https://arxiv.org/html/2410.02651v2#bib.bib12); Wulff & Hertz,1992 (https://arxiv.org/html/2410.02651v2#bib.bib30); Mordvintsev et al.,2020 (https://arxiv.org/html/2410.02651v2#bib.bib18); Chan,2020 (https://arxiv.org/html/2410.02651v2#bib.bib5))。例如，Mordvintsev 等人 (2020 (https://arxiv.org/html/2410.02651v2#bib.bib18)) 引入的通用 NCA 架构可概念化为“带有逐细胞丢弃的循环残差卷积神经网络”。

2.2 可控元胞自动机

可控元胞自动机 (CCA) 是 CA 的泛化，其在每个时间步能够接受外部输入。CCA 形式化了文献中由 Sudhakaran 等人 (2022 (https://arxiv.org/html/2410.02651v2#bib.bib26)) 引入的目标引导 NCA 概念。外部输入可以修改 CCA 的行为，提供了动态响应变化条件或控制信号的可能性，同时保持元胞自动机的基本原理。

一个 CCA 由元组 (L,S,I,N,φ) 定义，其中 (I) 是输入集，(φ: S^N × I^N → S) 是可控制局部规则。从格子到输入集的映射 (I: L → I) 称为输入。(I(x)) 表示细胞 (x \in L) 的输入。类似于状态，我们记 (I(N_x) = {I(n), n \in N_x})。

可控制全局规则 (Φ: S^L × I^L → S^L) 定义为对所有 (x \in L)，(Φ(S,I)(x) = φ(S(N_x), I(N_x)))。可控元胞自动机以初始状态 (S_0) 初始化。然后，在每个离散时间步 (t \in \mathbb{N})，根据可控制全局规则 (Φ) 和输入序列 ((I_t)_{t≥0}) 更新状态，得到：

(S_1 = Φ(S_0, I_0), S_2 = Φ(S_1, I_1), \dots)