研究人员打造“图表征状态空间”RNN，显著提升模型时空学习性能

英国爱丁堡大学博士毕业生、沙特阿卜杜拉国王科技大学博士后研究员李南伯一直在研究世界模型，这种模型能够根据事物的状态和背后规律预测未来，这一显著特征让李南伯深深着迷。

李南伯将预测理解为在时序上根据“背后规律”延展事物的“空间状态”。从这个角度来看，时空建模确实是世界模型面临的核心技术挑战之一。尤其在复杂“状态”描述（如通用高维向量表示）的场景下，研究其“时空规律”尤为复杂。这与传统学科研究非常相似，只不过各学科关注的是它们独有的“世界”。

在 AI 或机器学习领域，人们更关注如何让模型高效学习这些“时空规律”，而非规律本身。这本质上是一个计算问题。然而，现有的通用架构如 Transformers 和循环神经网络（RNN，Recurrent Neural Network）在计算上都有局限性。

Transformers 支持并行计算，但因其二次方复杂度问题难以扩展到无限长序列；而 RNNs 虽然可以用恒定内存扩展无限长序列，但并行能力不足（如训练时需串行计算）且存在长距离遗忘问题。近期广受关注的状态空间模型 Mamba 架构通过线性递归在 Transformers 和 RNNs 之间取得了一种折衷：它在时间/顺序轴上实现了高效压缩，因此更适合“时空学习”。但由于缺乏对输入数据的空间轴高效压缩和状态空间的充分利用，其时空学习潜力尚未完全释放。

为此，李南伯和所在团队提出一种简称为 FACTS 的新技术，针对这一问题重构了 Mamba 或通用结构状态空间模型的状态空间，将其结构化为一个图表征。该文章已被机器学习顶会 ICLR2025 接收。

换言之，FACTS 可以看作一个“图表征状态空间”的结构化状态空间模型，或一个“图表征状态空间”的 RNN。这种设计不仅显著提升了时空学习性能，还引入了处理图输入及建模因果关系（因果图）的能力。为了保留 Mamba 的并行计算特性，研究人员还提出了一种线性化的图状态机制。

对于相关论文审稿人表示：“本次论文提出了一种引入可置换内存结构的架构，能够灵活处理无序或动态变化的输入，同时通过高效压缩历史记录捕获长期依赖，从而在基线模型上实现了性能提升。这种架构采用内存输入路由机制，动态分配输入特征至潜在状态空间因子，解决了输入特征方差和动态关系建模的关键挑战。其设计既简单又高效，能够稳健处理输入顺序变化，同时简化高维数据的处理，并增强时空依赖的捕获能力，适用于实际应用场景。”

还有审稿人表示：“本次论文在多变量时间序列预测、时空图预测，以及以对象为中心的世界建模等任务中进行了广泛实验，结果表明 FACTS 在多项指标上始终优于或匹配当前最先进的模型。实验不仅验证了该方法在捕捉复杂时空动态方面的稳健性，还展示了其在不同数据集上的多功能性，进一步证明了其在现实世界中的应用潜力。”

作为一个通用时空模型的架构，除了传统时序预测应用领域，如金融、能源、交通、医疗等领域外，在多媒体领域，如视频，动画生成应该也有着很好的前景。FACTS 作为世界模型架构，李南伯也很期待看到其在其他学科研究中的应用，如物理、化学、生物以及相关社会科学。

正如之前提到的，李南伯对世界模型有着深深的执着，而 FACTS 是他在这一领域探索中的一项重要工作。李南伯在时空建模方面的思考，深受其导师“递归神经网络之父”、阿卜杜拉国王科技大学尤尔根·施密德胡伯（Jürgen Schmidhuber）教授提出的两个关键原则的影响：历史压缩（history compression）和可预测性最小化（predictability minimisation）。

基于这两条原则，李南伯对比了 LSTM、Transformers 和 Mamba 的优劣，并从中获得了 FACTS 的灵感。这个阶段的关键在于理论上的启发和明确方向。

2024 年 5 月，李南伯首次通过实验观察到重构“状态空间”后的“beta 版本”FACTS 在时间序列预测任务上的优越表现。不仅能够与当时最强基准模型匹敌，甚至在部分实验中有所超越。这一结果极大地增强了李南伯的信心，让他确信自己的思路是正确的，剩下的工作就是将想法付诸实践并优化模型。

模型迭代阶段是一个不断反思和完善的过程。从理论模型到实验实现，李南伯反复修改了大约 8 个版本的理论模型，实验实现更是多次推翻重来。汗水和咖啡因成为这段时间的主旋律，但因为方向明确，这段经历更多是令人愉快的挑战。这一阶段最大的难点在于设计可并行的内存更新机制。尽管 2024 年 5 月末的 beta 模型已经在时间序列预测上表现出很高的精度，但他希望在计算效率上不做过多妥协，时间复杂度至少要与 Mamba 同量级。这一目标推动他和所在团队最终开发出现有的可线性化的 FACTS，在性能与效率之间取得了平衡。

在这项研究的期间，李南伯收获了爱情与友情。虽然这看似与研究本身并没有直接的因果关系。但考虑到李南伯于 2024 年 3 月才刚刚加入施密德胡伯教授的实验室，一切都是全新的——新的环境、新的同事、新的研究。在短时间内结识志同道合的新朋友和研究伙伴，还幸运地收获了爱情，为李南伯能展开 FACTS 研究提供了坚实保障。其表示，每次赶截稿的过程都像被“扒了一层皮”，无比痛苦，但大家一起赶，在压力之下还能互相开玩笑，这确实是一种慰藉。

除了继续在世界模型方向上深入探索，他还计划对当前的 FACTS 模型进行扩展（scale up），以开展基于 FACTS 的基础模型（foundation model）应用研究。

参考资料：

1.https://arxiv.org/pdf/2410.20922

运营/排版：何晨龙