💬 NLP

Learning State-Tracking from Code Using Linear RNNs

本文通过将置换组合转换为基于代码的 REPL 轨迹，弥合了状态追踪研究与下一标记预测之间的鸿沟，并证明了虽然线性 RNN 在该任务上优于 Transformer，但在状态揭示是确定性而动作并非完全可观测的情况下，它们可能会逊于非线性 RNN。

原作者： Julien Siems, Riccardo Grazzi, Korbinian Pöppel, Kirill Kalinin, Hitesh Ballani, Babak Rahmani

发布于 2026-06-26

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Julien Siems, Riccardo Grazzi, Korbinian Pöppel, Kirill Kalinin, Hitesh Ballani, Babak Rahmani

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

大局观：代码的“猜盅戏”

想象一下你正在看魔术师表演“猜盅戏”（Shell Game）。桌上有三个杯子，一个球被藏在其中一个下面。魔术师不断交换杯子的位置。你的任务是追踪球的位置。

球：计算机程序中的一个变量（比如存储在盒子里的一个数字）。
交换： 移动变量的代码指令。
目标： 在经过一系列长长的交换后，准确知道球在哪里。

长期以来，研究人员一直使用一种特定的设置来测试 AI 在这个“猜盅戏”中的表现：他们向 AI 展示交换列表，然后问：“现在球在哪里？”论文指出，这就像要求一个学生背诵一整部电影剧本，然后复述结局。这并不能测试 AI 是否真的在过程中“理解”了故事。

新方法：“实时解说”

作者改变了测试方式，使其更符合真实 AI（比如你正在与之对话的这个 AI）的学习方式：下一个标记预测（Next-Token Prediction）。

他们不再一次性展示整个交换列表，而是给 AI 提供一份正在运行的计算机程序的实时转录文本，逐行进行。

第 1 行： “将杯子 A 移到 B。”
第 2 行： “查看杯子 A 下面！”（计算机打印出结果）。
第 3 行： “交换杯子 B 和 C。”
第 4 行： “查看杯子 C 下面！”

AI 必须预测转录文本中的下一个词。为了做到这一点，它必须像人类跟随故事一样，在阅读过程中在脑海中追踪杯子的状态。

竞争对手：“线性模型” vs. “Transformer”

论文将两种类型的 AI 架构进行了对比：

Transformers（“摄影式记忆”）： 这是目前 AI 领域的冠军（比如背后驱动这类聊天的模型）。如果信息就在眼前，它们非常擅长记忆事实和寻找模式。
线性 RNN（“笔记记录员”）： 这些是更新、更快的模型，旨在按顺序处理信息，就像逐字阅读一本书一样。

结果：

当“猜盅戏”完全透明时（AI 能看到每一次交换和每一次揭晓），线性 RNN（特别是被称为 DeltaNet 的一种类型）表现得非常出色。即使游戏变得很长，它们也能完美地追踪球的位置。
Transformers 则显得有些吃力。它们需要非常频繁地看到状态的揭晓才能跟上进度。如果“揭晓”的间隔较长，它们就会迷失方向。

转折点：当游戏变得“模糊”时

论文接着提出了一个问题：如果游戏不再那么清晰明确，会发生什么？

在真实的计算机代码中，情况并不总是确定性的。有时，代码会做出随机选择，或者一个变量取决于 AI 看不到的东西（比如隐藏的环境变量）。

作者创建了一个场景，让 AI 必须根据概率来猜测状态（例如：“有 50% 的概率球向左移动，有 50% 的概率留在原处”）。

线性 RNN 的问题：
论文发现了线性 RNN 在处理这种“模糊”不确定性时的根本弱点。

类比： 想象你正试图保持一叠纸的平衡。每当你得到一个新的线索（一次“揭晓”）时，你都必须重新整理这叠纸。
在线性 RNN 中，用于更新这叠纸的数学运算是“线性”的。这就像一个漏水的桶。每当你得到一个部分线索时，你的一点点“信心”（数学上的质量）就会流失出去。
如果你在没有得到“完全重置”（即清晰、完整的揭晓）的情况下，经历了一系列长期的部分线索，你的信心就会呈指数级萎缩。最终，这个数字会变得如此之小，以至于计算机将其视为零。AI 会忘掉一切。

“对抗性”陷阱：
作者展示了如何用特定的移动序列来欺骗这些线性 RNN：

随机混合杯子（制造不确定性）。
揭示其中一个特定杯子的位置（给出一个部分线索）。
重复此过程。

每当这种情况发生时，线性 RNN 追踪其他杯子的能力就会减弱一点。经过足够多次的重复后，线性 RNN 对其他杯子位置的内部“信念”会完全消失，尽管人类仍然可以根据逻辑推断出答案。

结论

线性 RNN 非常擅长在规则清晰且路径确定（如完美的猜盅戏）的情况下追踪状态。如果训练设置得当，它们甚至可以击败 Transformer。
线性 RNN 在处理现实世界中具有概率性或部分隐藏特征的代码时表现挣扎。 它们的数学结构导致它们在处理不确定性时，随着时间的推移会“遗忘”细节，因为它们缺乏一种机制，能在不破坏其线性结构的前提下进行“重新归一化”或修复其信心水平。

简而言之：线性 RNN 非常擅长遵循清晰的剧本，但当剧本变得模糊和随机时，它们往往会失去理智。

技术摘要：利用线性 RNN 从代码中学习状态追踪

问题陈述

状态追踪是理解程序执行、游戏状态和世界模型的一项基本能力。虽然理论研究已经确立了一个分歧，即 Transformer 擅长关联召回（associative recall），而循环神经网络（RNN）在状态追踪方面表现出色，但在这些模型如何被评估与如何被训练之间存在显著差距。

大多数现有的状态追踪基准测试（例如排列组合）采用的是序列到序列（sequence-to-sequence）的设置：将一系列动作直接映射到一系列状态。这与现代大语言模型（LLM）使用的下一 Token 预测范式有着本质区别，在后者中，模型必须预测一个交织了动作与部分状态揭示（例如 print 语句）的流中的下一个 Token。此外，现实世界的代码执行通常涉及部分可观测性和概率性转换（例如基于隐藏状态或随机选择进行分支），这创造了一种真实状态并非完全可观测的情境，需要维护一个信念分布（belief distribution）而非确定性的状态。

方法论

1. 通过代码重构状态追踪（REPL Trace）

为了弥合抽象自动机基准测试与 LLM 训练目标之间的差距，作者将排列组合任务转换为 Python REPL trace。

机制： 模型不再是将排列序列 $\sigma_1, \dots, \sigma_m$ 映射为累积状态，而是处理一系列 Python 命令（变量赋值、交换）以及交织其中的、用于揭示部分状态信息的 print 语句。
目标： 模型通过下一 Token 预测进行训练，以预测 trace 中的下一个 Token，该 Token 可以是一个命令或被揭示的状态值。这迫使模型通过预测任务隐式地学习状态追踪。

2. 架构比较

本研究在这些任务上比较了三类架构：

Transformers： 标准的基于注意力的模型。
线性 RNNs (DeltaNet [0, 1])： 特征值限制在 $[0, 1]$ 范围内的线性 RNN，已知无法处理复杂的如奇偶校验（parity）或排列组合等状态追踪任务。
线性 RNNs (DeltaNet [-1, 1])： 具有扩展特征值范围 $[-1, 1]$ 的线性 RNN，理论上能够解决确定性的排列组合问题。

3. 理论框架：PFSA-SR

为了分析现实代码中状态追踪的难度，作者引入了带有状态揭示的概率有限状态自动机（PFSA-SR）。

定义： 一个元组 $(Q, \Sigma, \delta, \rho, q_0)$ ，其中 $\delta$ 是概率转移核， $\rho$ 是将输入映射到状态子集的揭示函数。
动力学： 系统通过概率转移进行演化，但观察者接收到的是“硬”约束（揭示），这些约束会剪枝可能的状态集合（支持集剪枝），而非提供“软”似然度。
信念更新： 追踪需要维护一个状态上的信念分布 $b_t$ 。更新过程包括一个揭示步骤（将不一致的状态归零并重新归一化）和一个转换步骤。至关重要的是，重新归一化步骤 $f(x) = x/\|x\|_1$ 是非线性的。

4. 表示分析

论文研究了线性 RNN 中信念状态的两种表示方法：

联合表示（Joint Representation）： 显式枚举所有 $n!$ 个排列。作者认为这对于线性 RNN 来说在数值上是不稳定的，因为输入注入项 $B(x_t)$ 无法在部分揭示后自适应地补充生存配置的概率质量，从而导致指数级的范数衰减。
边缘表示（Marginal Representation）： 独立追踪每个位置的边缘概率（一个 $n \times n$ 的双随机矩阵）。虽然这允许 $B(x_t)$ 直接为被揭示的条目注入确定性，但作者表明，在对抗性揭示序列下，即使是这种表示也会受到影响——在这种序列中，特定的未揭示条目从未被直接补充。

关键结果

1. 监督密度与架构

Transformers： 需要密集的状态监督来学习状态追踪。随着 print 语句（揭示）之间间隔的增加，Transformer 的性能迅速下降。即使在密集监督下，它们也无法泛化到全排列。
线性 RNNs (DeltaNet [-1, 1])： 在这种设定下表现出色。它们即使在稀疏监督下也能可靠地学习追踪状态，并展示出鲁棒的长度外推能力（能够泛化到比训练序列更长的序列）。
线性 RNNs (DeltaNet [0, 1])： 无法解决该任务，证实了扩展特征值范围 $[-1, 1]$ 对于复杂的状态追踪是必要的。

2. 状态追踪的可解释性

对训练好的 DeltaNet [-1, 1] 的分析显示，状态追踪被局部化到了特定的组件中：

特征值扩展： 至少有一个注意力头在特定层（例如第 12 层，第 3 个头）学习将其衰减参数 $\beta$ 设置为 2，从而产生一个为 $-1$ 的特征值。
子空间结构： 这个特定的头在一个显著低维的线性子空间中运行，与其他头相比，这与追踪排列的理论要求一致。
干预： 将该特定头的 $\beta$ 范围缩回到 $[0, 1]$ 会导致性能立即崩溃，而缩放其他头则没有影响，这证实了该头是专门的态追踪机制。

3. 线性 RNN 在概率设定下的局限性

论文证明，虽然线性 RNN 可以解决确定性任务，但在概率设定（PFSA-SR）中面临内在障碍：

指数衰减： 在联合表示下，部分揭示会导致未归一化的状态范数呈指数级衰减（ $1/2 \to 1/4 \to 1/8 \dots$ ），因为线性递归无法执行必要的非线性重新归一化。
对抗脆弱性： 即使使用边缘表示，对手也可以构建一个包含概率交换和选择性揭示的序列，使得未揭示条目的概率质量呈指数级衰减，直到在有限精度下消失。
结论： 在具有部分可观测性的概率设定中，实现稳定的信念维护对于线性 RNN 而言似乎是不可行的，除非引入非线性重新归一化或频繁的全状态重置。

意义与主张

该论文声称提供了一个更真实的测试平台，用于测试与语言模型下一 Token 预测目标相一致的状态追踪，超越了传统的合成序列到序列基准测试。

弥合差距： 它证明了能够进行状态追踪的线性 RNN（特别是具有扩展特征值的 DeltaNet）可以通过在类代码 trace 上的下一 Token 预测来学习这些能力，而 Transformer 则不能。
理论障碍： 它识别了线性 RNN 在现实代码场景中的一个根本局限。现实代码通常涉及概率性转换和部分可观测性（例如隐藏的分支、随机选择）。作者表明，在这些环境中追踪状态需要非线性重新归一化，而线性 RNN 内部无法执行此操作。这表明，在现实任务中观察到的“复杂度差距”可能源于代码的概率性质，而非仅仅是序列长度或复杂度。
基准测试的影响： 作者指出，仅在确定性有限状态自动机（DFA）任务上对架构进行基准测试，会掩盖非线性 RNN 的潜力以及线性 RNN 的局限性。提出的 PFSA-SR 框架及其相关的失效模式（指数范数衰减）为为什么标准线性递归模型在需要不确定性推理时表现挣扎提供了理论解释。

论文最后进行了谦逊的总结，指出其实验侧重于合成排列群，未来的工作需要扩展到现实世界的执行 trace，并研究交替使用线性递归与周期性非线性归一化的混合架构。