支撑当今大语言模型的 Transformer（深度学习中的注意力模型架构） 架构挑战了我们对 AI 能力的既有认知。编写代码、语言翻译、从网络各个角落综合信息——这些任务如今的大语言模型似乎都能轻松应对。短短几年间，它已经改变了我们的工作方式。

但基于 Transformer 的模型存在固有局限，研究人员在试图扩展其能力并降低计算成本时不断遭遇这些瓶颈。Transformer 的核心缺陷之一——至少在理论上——在于其对序列数据建模的能力有限。这会使某些计数问题变得出奇地困难——比如统计“strawberry”中有几个“r”。

通过思维链（CoT）提示将问题分解为多个步骤，或切换到具有更多顺序处理能力的“通用”架构，基于 Transformer 的模型可以得出正确答案。但这类变通方法往往会给本已昂贵的流程增加大量时间和成本。

Transformer 通过自注意力机制彻底改变了自然语言处理，该机制使其能够一次性处理大块文本。注意力可以让 Transformer 更丰富地理解词语之间的关联以及如何传达含义。但通过并行、无序地处理数据，Transformer 也在一种称为状态跟踪的能力上表现不佳，而状态跟踪涉及随着情况变化逐步更新其世界观。

状态跟踪是我们如何在长时间对话中追踪细节，或评估一段代码、或国际象棋比赛中对手上一步棋的方式。Transformer 会记住每一个新观察以供后续检索，AI 研究员 Albert Gu 将这种策略比作数据库。Transformer 可以精确回忆事实，但当历史变得过长时也会耗尽内存。

基于循环神经网络（RNN）架构的早期语言模型工作原理略有不同。它们逐词处理文本，将过去的输入总结为一个压缩状态，在新文本流入时可以参考该状态。它们的记忆可能更模糊，但由于经过了压缩，它们可以回溯更久远的时间。

为了提升推理速度，IBM 和其他公司已将类 RNN 处理集成到其大语言模型中。IBM 的 Granite 系列模型将 Transformer 层与具有类 RNN 结构的状态空间模型（SSM）层交替排列。Granite 的混合 Bamba 架构受到 Nvidia 在 Mamba2 中将 SSM 与 Transformer 融合的启发，显著提升了效率，但缺乏通用状态跟踪能力。

在 NeurIPS 2025 的 spotlight poster 中，IBM 研究人员深入研究了其代数原理，并提出了一种构建 SSM 转移矩阵的新方法以实现状态跟踪。在实验中，他们改进的模型——称为 PD-SSM——在通用状态跟踪任务上显著优于其他 SSM 变体。当研究人员将其方法集成到混合 Transformer-SSM 模型中时，他们发现该模型能够处理更复杂的任务，包括让模型预测用自然语言描述的一系列按时间排序的事件中接下来会发生什么。

这项工作的意义可能超越研究中的学术练习。有证据表明，状态跟踪对代码生成尤为重要，而这是 IBM 及其他公司一直积极追求的 AI 应用方向。“我们很兴奋能进一步探索这些影响，”该研究的第一作者 Aleksandar Terzic 表示，他是一位专注于 AI 数学基础的 IBM 研究员。

煤矿中的金丝雀

20 世纪 50 年代，语言理论家 Noam Chomsky 提出了一种按“文法”复杂度来组织所有语言的方式。最简单的语言可以用有限状态机来建模，这种机器内存有限，像红绿灯或旋转门一样，根据定义的输入从一个状态转换到另一个状态——红绿灯靠定时器，旋转门靠令牌。Chomsky 层级体系的最高层留给了图灵完备的“无限制”文法，可以计算任何可计算的内容。

甚至在 ChatGPT 席卷全球之前，AI 研究人员就已经在使用乔姆斯基的方法论来探究深度神经网络的局限性。2020 年，斯坦福大学研究员 Michael Hahn 从数学上证明，Transformer 在乔姆斯基层级的最低层次上执行状态跟踪任务的能力有限。

这些任务中最著名的、任何学过计算理论入门课程的人都熟悉的，或许是奇偶校验问题：给定一串 1 和 0，要求你计算 1 的个数加起来是奇数还是偶数。Hahn 证明，Transformer 在理论上会在某个位长度上无法通过奇偶校验测试。

“如果 Transformer 无法计算奇偶校验，它们也无法准确评估逻辑公式，”他写道，这预示了多年后互联网对‘草莓’问题的痴迷。

IBM 研究员 Shawn Tan 在魁北克 Mila AI 研究所读研究生时，就被奇偶校验问题所触动，当时他在博客中称之为“煤矿中的金丝雀”。如今正在训练 IBM Granite 模型的 Tan 继续关注着这场辩论。“如果一个模型无法顺序跟踪两个状态，我们怎么能期望它解决更具挑战性的问题？”他最近说道。

状态跟踪问题激发了其他研究人员不仅对 Transformer 进行检验，还对 RNN、LSTM 和现代 SSM 等序列模型进行了检验。Terzic 在 NeurIPS 上的 spotlight 论文是该系列中的最新一篇。

重构 SSM 的状态转移矩阵

SSM 几十年来一直用于建模动态系统，2021 年 Gu 及其斯坦福大学的同事引入 S4 模型时，SSM 跨入了深度学习领域。凭借固定的内存占用，S4 能够比 RNN 更巧妙地处理长序列，并且比 Transformer 快得多。

对角矩阵简化了计算并减少了内存开销，后来被引入，帮助 SSM 扩展规模并在复杂任务上与 Transformer 竞争。如今大多数 SSM 使用对角矩阵，但性能和效率的提升通常以牺牲状态跟踪为代价。对角 SSM 已被证明在奇偶校验和类似的抽象数学问题上具有与 Transformer 相同的局限性。

如果一个模型无法顺序跟踪两个状态，我们怎么能期望它解决更具挑战性的问题？

在他们的新论文中，IBM 研究人员提出放弃对角矩阵，以恢复 SSM 中的状态跟踪能力。他们在奇偶校验及相关状态跟踪任务上测试了多种 SSM，使用的字符串长度最多为模型训练时的六倍，以衡量其泛化能力。他们的 PD-SSM 模型以至少 15 个百分点的优势超越了其他模型。PD-SSM 的平均得分 98.5% 仅次于 LSTM，即长短期记忆网络，是一种循环神经网络（RNN），旨在通过学习和记忆长期依赖关系来识别序列数据中的模式，LSTM 擅长状态跟踪但难以扩展。

IBM 的 PD-SSM 模型在涉及长时间序列数据集的实际任务中也表现出色——例如心跳分类和乙醇需求预测。他们甚至在奇偶校验的单词版本上测试了一个混合模型，令人联想到 IBM 的 Bamba 架构，发现它能够计算出长达 40 个字符的正确答案，而其他 SSM 在 15 个字符后就力不从心了。

研究人员现在面临的挑战是将这些理论洞见转化为实际改进。“我们的团队全力专注于将 PD-SSM 的优势迁移到 Granite-4 中，”领导这项工作的 IBM 研究员 Abbas Rahimi 说道。

如果大语言模型能够掌握草莓问题和奇偶校验问题，谁知道它们还能取得哪些更大的成就？

注释

• 注释 1：LSTM，即长短期记忆网络，是一种循环神经网络（RNN），旨在通过学习和记忆长期依赖关系来识别序列数据中的模式 ↩︎

转载声明：本文为转载整理，原文版权归原作者及原站点所有。

原文来源：IBM Research Blog

原创地址：https://research.ibm.com/blog/state-tracking-for-state-space-models