[Paper] 使用线性 RNN 从代码学习状态跟踪

Source: arXiv - 2602.14814v1

（请提供您希望翻译的正文内容，我将按照要求将其翻译为简体中文并保留原始的格式和链接。）

Overview

论文 Learning State‑Tracking from Code Using Linear RNNs 研究了当输入是一系列类似程序的操作流，而不是干净的“输入 → 输出”对时，序列模型如何跟踪隐藏状态。通过将置换‑组合问题转化为 REPL 风格的代码跟踪（带有打印以显示当前状态），作者弥合了经典状态跟踪基准与用于训练现代语言模型的下一个 token 预测范式之间的差距。他们的实验表明，线性循环网络能够掌握此任务，而 Transformer 仍然表现不佳。

关键贡献

• 基于代码的重新表述置换组合：将操作表达为 Python 风格的语句，并交叉插入显式的状态打印，使问题兼容语言模型训练。
• 展示线性 RNN（即不含非线性激活的 RNN）能够在代码追踪设置中成功跟踪隐藏状态，在相同数据上优于 Transformer。
• 理论框架将代码中的状态追踪视为对概率有限状态自动机（PFSA）进行推理，且该自动机具有确定性状态揭示，突出任务本质上的难度。
• 实证证据表明，当底层 PFSA 具有部分可观测的动作时，非线性 RNN 相较线性 RNN 更具鲁棒性，提供了模型表达能力的细致视角。

方法论

1. 任务转换 – 经典的置换‑组合基准（给定一系列置换，输出结果置换）被重新表述为 REPL 跟踪：每个置换作用于一个变量，并在每一步后使用 print 语句显示当前变量的值。
2. 模型族 –
   • 线性 RNN：递归更新完全线性（无激活函数）。
   • 非线性 RNN：标准门控结构（例如 LSTM/GRU）。
   • Transformer：仅编码器模型，使用因果掩码进行训练。
3. 训练方案 – 所有模型都在生成的代码跟踪上进行下一个 token 预测的训练，完全等同于在源代码上训练语言模型的方式。
4. 理论分析 – 作者将动作序列建模为 PFSA（概率有限状态自动机），其中每个隐藏状态通过 print 确定性地显现。他们证明了在线性更新下无法唯一恢复 PFSA 状态的条件，而非线性动力学则可以实现。

Source: …

结果与发现

• Linear RNN 能够精确模拟底层置换组合，即使它们仅接收到下一个 token 的信号。
• 即使容量很大，Transformer 也无法可靠地从交错的打印信息中推断出隐藏的置换状态。
• 当 PFSA 包含不可完全观察的动作（即多个隐藏转移产生相同的打印状态）时，线性 RNN 的性能下降比非线性 RNN 更为明显，这验证了理论分析。

Practical Implications

• Code‑aware language models: 该研究表明，简单的线性递归在需要精确状态维护的任务（例如符号执行、静态分析或对程序轨迹进行自动推理）中可能出奇地有效。
• Model selection for tooling: 构建 IDE 助手、调试器或自动重构工具的开发者可能会倾向于使用轻量级线性 RNN 来实现确定性的状态跟踪组件，从而在速度和内存方面优于笨重的 Transformer。
• Benchmark design: 通过将算法问题转换为 REPL 风格的代码，研究人员可以在此前仅限于序列到序列设置的任务上评估下一个 token 语言模型，开启衡量推理能力的新途径。
• Security & verification: 从代码轨迹中实现准确的状态跟踪有助于在沙箱环境中检测不变式违规或恶意状态操控。

限制与未来工作

• 任务范围 – 实验聚焦于置换组合；尚未明确这些发现能在更复杂的程序分析（例如循环、条件语句或数据结构操作）中推广得多好。
• Transformer 变体 – 仅测试了标准因果 Transformer；具有显式递归或记忆的更新架构（例如 Transformer‑XL、检索增强模型）可能缩小性能差距。
• 可扩展性 – 线性 RNN 在相对短的轨迹上表现出色；其处理包含数千步的长时间运行程序的能力尚未评估。
• 概率扩展 – PFSA 框架可以扩展以建模噪声或部分损坏的打印输出，反映真实调试输出，从而进一步检验鲁棒性。

底线：通过将算法状态跟踪重新构造为代码生成问题，作者揭示了线性 RNN——常被视为“过于简单”——能够在确定性状态推断上超越现代 Transformer，为构建需要推理程序状态的下一词模型的开发者提供了全新视角。

作者

• Julien Siems
• Riccardo Grazzi
• Kirill Kalinin
• Hitesh Ballani
• Babak Rahmani

Paper Information

• arXiv ID: 2602.14814v1
• Categories: cs.LG, cs.CL
• Published: 2026年2月16日
• PDF: Download PDF