Causal Evidence of Stack Representations in Modeling Counter Languages Using Transformers

技术摘要：关于 Transformer 在建模计数语言中栈表示的因果证据

问题陈述
尽管机械解释性（mechanistic interpretability）已成功识别出在形式语言上训练的语言模型（LMs）中的潜在特征，但在建立这些特征的因果作用方面仍存在关键空白。前人的研究（如 Tiwari 等人，2025）表明，在计数语言（例如 Dyck-1, Shuffle-k）上训练的 Transformer 学习到的表示与底层的栈结构一致。然而，这些研究并未确定这些表示仅仅是与任务相关的相关性，还是模型计算所必需的因果因素。本文旨在解决以下问题：学习到的栈表示是否在 Transformer 预测下一个 token 的过程中起到了因果作用？

研究方法
本研究采用了一个包含训练、探测和因果干预的三阶段实验流程：

1. 模型训练： 一个仅编码器（encoder-only）的 Transformer 被用于 Shuffle-k 语言（其中 $k \in \{2, 4, 6, 8\}$）的下一个 token 预测任务。Shuffle-k 定义为 $k$ 个独立 Dyck-1 语言的交织，要求模型同时追踪 $k$ 个不同括号栈的深度。模型架构由单层组成，包含 4 个注意力头、32 维嵌入维度和 64 维隐藏维度。模型在第 25 个 epoch 时达到了完美的验证准确率，确保随后的所有分析均在完全收敛的模型上进行。
2. 栈深度探测： 通过线性探测器从模型的隐藏状态 ($h^{(l)}_t$) 中预测特定栈的深度 ($c_i(t)$)。为了确保探测器捕捉的是真实的结构信息而非伪相关，研究中使用了一个基于扰动标签训练的控制探测器（control probe）。指标**选择性（Selectivity）**定义为实验探测器与控制探测器之间的准确率之差。高选择性表明栈深度在隐藏空间中得到了真实的编码。
3. 因果干预（消融）：
   • 方向提取： 通过奇异值分解（SVD）对探测器的权重矩阵 $W$ 进行分解。通过识别其右上奇异向量 $w$，确定对应于栈深度表示的主成分方向。
   • 消融： 在推理过程中，使用操作 $h' = h - \alpha(w \cdot h)w$ 移除与 $w$ 对齐的隐藏状态分量，其中 $\alpha \in [0, 1]$ 代表消融强度。
   • 控制组： 同样消融随机方向作为基准。
   • 评估： 将消融后的状态直接传递给解码器头（不重新运行 Transformer 的内部层），以衡量对位置准确率（$A_{pos}$）和序列准确率（$A_{seq}$）的影响。

核心结果

• 表示的存在性： 线性探测器成功提取了栈深度信息，并表现出高选择性（例如，$k=6$, Stack 2 时为 85.8%），证实了模型学习到的表示与栈结构一致。
• 因果必要性： 对主栈方向 ($w$) 进行针对性消融会导致模型性能剧烈崩溃。随着消融强度 $\alpha$ 从 0 增加到 1，序列准确率降至接近 0%。
• 特异性： 消融随机方向对准确率没有显著影响，这证实了性能崩溃是由于移除栈表示引起的，而非向量操作产生的普遍人工痕迹。
• 指标敏感性： 研究观察到，虽然位置准确率有所下降，但序列准确率的崩溃更为剧烈。这归因于序列生成的误差累积特性，即单个错误的 token 会使整个字符串失效。

意义与主张
本文声称提供了强有力的经验证据，证明在处理计数语言的 Transformer 中，栈表示不仅是涌现的伪影，而且是因果必要的，对于模型的内部计算和下一个 token 的预测至关重要。

通过演示移除特定的、学习到的表示方向会导致模型几乎完全失效，作者认为 Transformer 明确利用了这些类栈结构来解决形式语言所需的逻辑推理任务。这使该领域从相关性分析（探测）转向了因果机械理解，验证了模型的“黑盒”行为依赖于特定的、可识别的计算原语。

未来方向
作者建议未来的工作可以探索：

• 通过在中间网络中传递消融或编辑后的状态进行干预，以观察对后续输出的影响。
• 研究不同架构在这些表示的涌现、冗余性和因果性方面的作用。
• 理解 Transformer 如何处理其他涌现出的计算原语（如计数器）的操作。