🤖 machine learning

Task Structure Reverses Layerwise State Encoding in Sequence Models

本文表明，序列模型中状态编码的层级分布并非固定的架构特征，而是会根据任务的计算结构发生反转，这揭示了诸如交换律之类的代数属性，其对机制特征的预测能力不如前缀更新与栈操作等底层计算需求。

原作者： Yuhang Jiang

发布于 2026-06-02

📖 1 分钟阅读☕ 轻松阅读

原作者： Yuhang Jiang

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

想象一下，你正试图理解不同类型的“思考机器”（AI 模型）在阅读故事时是如何记录信息的。长期以来，研究人员认为机器的“思考风格”是固定的，就像一个人的性格一样。他们曾相信：

Transformer（当前的 AI 标准）会将他们的思想均匀地分布在整个大脑中。
循环模型（旧式的、基于循环的 AI）会将他们的核心思想藏在脑部的最后一部分。

这篇论文说：“别急，没那么简单。”

作者们发现，这些机器并没有固定的“人格”。相反，它们会根据正在执行的任务来改变如何组织自己的思想。这就像一位厨师会根据是切菜还是片鱼而使用不同的刀具。工具的变化取决于任务本身，而不仅仅是厨师的身份。

以下是利用简单的类比对他们发现的拆解：

1. 两种类型的任务

研究人员给机器出了三类不同的谜题：

“累加总数”任务（Parity & S3）： 想象你在计算抛硬币的次数。你只需要记住当前的计数（0 或 1）并随着进行不断更新即可。这是一个简单的、线性的更新过程。
“堆栈”任务（Dyck）： 想象在检查一个句子是否具有平衡的括号，例如 (( ))。你不能仅仅计数；你需要记住哪个左括号对应哪个右括号。你必须建立一个心理上的“堆栈”（就像一叠盘子）来追踪嵌套关系。

2. 大反转

论文发现，机器会根据任务的不同而切换策略：

在“累加总数”任务中：
- 循环模型（如 Mamba）表现得像个秘密守护者。他们在早期做了大量工作，但最终答案直到大脑的最末层才会变得清晰。这就像一位魔术师在阴影中完成了所有的手法，只在最后时刻才揭晓戏法。
- Transformer 表现得像一场团队集会。他们从头到尾，通过每一层逐渐构建出答案。每个人都在过程中贡献了一点力量。
在“堆栈”任务中：
- 角色发生了互换！
- Transformer 突然变成了秘密守护者。他们在最初的几层就几乎立刻得出了答案，然后只是将其保留下来。
- 循环模型 变成了团队集会。他们必须通过一层又一层地缓慢构建，才能得到正确的答案。

核心结论： 你不能说“Transformer 总是分散信息”或“Mamba 总是隐藏信息”。策略完全取决于任务是简单的更新还是复杂的堆栈。

3. “非交换性”测试（S3 谜题）

为了证明这不仅仅是数学规则的问题（比如加法是否满足交换律），他们加入了第三个更难的谜题，叫做 S3。这是一个“累加总数”任务，但其中的顺序非常重要（就像穿袜子要在穿鞋之前，顺序颠倒了结果就不同）。

预测： 如果差异仅仅在于数学规则，那么这个棘手的任务应该看起来像“堆栈”任务。
现实： 机器处理这个棘手任务的方式与简单的“累加总数”任务完全一样。它们使用了相同的“秘密守护者”或“团队集会”策略。
结论： 机器并不是在对数学规则做出反应；它们是在对计算结构（是简单的更新还是复杂的堆栈？）做出反应。

4. “可读性”与“重要性”的陷阱

研究人员还观察了信息存储的位置。他们发现了一个令人惊讶的脱节，尤其是在大型预训练模型中：

“可读层”： 这是你可以轻松从机器大脑中“读出”答案的地方（就像在某个地方找到了写着答案的便签纸）。
“重要层”： 这是大脑中如果被破坏，就会导致机器失效的部分。

发现：

在小型、简单的模型中，“可读层”和“重要层”通常是同一个。
在大型预训练模型中，它们往往并不匹配。
- 例子： 在“堆栈”任务中，一个大模型可能在脑部中间（第 7 层）清晰地写着答案，但如果你破坏脑部末端（第 11 层），模型仍然可以正常工作。然而，如果你破坏中间部分，它就会崩溃。
- 例子： 在“累加总数”任务中，答案可能在最后清晰可见，但破坏那个特定位置并不会阻止模型运行，因为信息已经分散在各处了。

教训： 仅仅因为你能在大脑的某个特定部分“看到”答案，并不意味着那个部分是维持机器运转的唯一关键。机器可能同时在许多不同的地方持有该答案。

总结

这篇论文告诉我们，AI 架构并不是僵化的。它们是灵活的。

任务至关重要： 机器的内部策略会根据它是执行简单的更新还是复杂的堆栈而发生变化。
结构至关重要： 问题的“形状”（更新 vs 堆栈）比数学规则更能决定其策略。
可见性 $\neq$ 必要性： 在大模型中，找到答案“写在哪里”并不一定能告诉你机器最脆弱的地方在哪里。

作者得出结论：我们不能只问“这个 AI 是如何工作的？”我们必须问：“这个 AI 在执行这项特定任务时是如何工作的？”

技术摘要：任务结构反转了序列模型的层级状态编码

问题陈述

机制可解释性研究通常将层级状态编码视为固定的架构特征。普遍的假设是，循环模型（如 LSTM、GRU）和状态空间模型（SSM，如 Mamba）将可读状态集中在深层，而基于注意力机制的模型（Transformer）则将状态积累分布在不同深度。本文挑战了这种认为编码策略是稳定架构不变性的观点。本文指出，当计算结构发生变化时，相同的架构会反转其层级状态编码剖面。此外，文献中经常混淆两种潜在驱动因素：代数结构（特别是交换性）与计算结构（前缀更新 vs. 栈追踪）。作者旨在解构这些因素，并确定状态编码是架构本身的属性，还是架构与任务之间的相互作用。

研究方法

本研究通过五种架构家族进行了受控实验评估：Transformer、Mamba、Mamba-2、LSTM 和 GRU。

任务：
作者使用了三种形式化状态追踪任务和一种半真实基准测试：

Parity（奇偶校验）： 一个具有交换性的、单比特前缀更新任务。
$S_3$ 置换组合： 一个非交换的、六类前缀更新任务。这作为关键判别器：如果代数交换性驱动编码剖面， $S_3$ 的表现应类似于 Dyck；如果计算结构驱动它， $S_3$ 则应表现得像 Parity。
Dyck- $k$ ： 一个有界深度的栈追踪任务（括号匹配），最大深度为 10。
Code-Depth（代码深度）： 一个源自 CodeSearchNet 中 Python 代码的 Token 级嵌套深度任务，作为半真实验证。

实验设置：

形式化任务： 实验使用 4 层模型，且 $d_{model}=128$ 。在可能的情况下，架构进行了参数匹配（约 83 万个参数）。在序列长度 $L=40$ 时进行训练，并在 $L \in \{40, 60, 80, 100, 120\}$ 时进行评估。
预训练规模： 在 Mamba-130M（24 层）和 Pythia-160M/410M（12/24 层）上进行微调。
测量指标：
- 线性探测（Linear Probing）： 在冻结的隐藏状态上训练逻辑回归分类器，以测量每一层的状态可读性。
- 因果干预（Causal Interventions）： 零消融（投影掉探测方向）和子空间消融（移除前 PCA 或奇异方向）以测试功能必要性。
- OOD 泛化： 将在 $L=40$ 处训练的消融应用于未见长度（ $L=80, 100, 120$ ）。
- 组件归因： 在匹配的纯 PyTorch 后端下消融 Mamba 的 Conv1D 组件，以隔离其在不同任务中的贡献。
- 激活修补（Activation Patching）： 在预训练 Mamba 中，将干净输入的中间位置激活值修补到损坏的输入中，以测量 Logit 恢复情况。

核心贡献

任务依赖的反转： 本文证明了“集中于深层”与“分布式构建”的模式并非架构不变性。该模式取决于任务而发生翻转：
- 前缀更新任务（Parity、 $S_3$ ）： 循环模型和 SSM 将可读状态集中在最后一层，而 Transformer 则在各层之间逐渐构建状态。
- 有界栈追踪（Dyck）： 模式发生反转。Transformer 能够立即实现近乎完美的读取度（早期层），而循环模型和 SSM 则呈现渐进式构建。
解构代数与计算： 通过引入非交换的 $S_3$ 任务，作者表明编码剖面追踪的是计算结构（前缀更新 vs. 栈）而非代数结构（交换性）。 $S_3$ 在所有架构中都与 Parity 归为一类，从而反驳了交换性是编码剖面主要驱动因素的假设。
可读性与因果性的脱节： 研究发现，状态最线性可读的层（最佳探测层）并不总是与执行任务在因果上最必要的层重合。这种脱节在预训练模型中尤为显著，因为其表示更加冗余且分布广泛。
组件级任务依赖： 即使在单一架构（Mamba）内，特定组件的重要性也随任务而异。例如，Conv1D 组件对于 $S_3$ 和 Parity 至关重要，但对 Dyck 的性能没有显著影响，这挑战了关于架构具有单一“机制”的观念。

关键结果

行为层面： 循环基准模型在 Parity 和 $S_3$ 上表现出完美的泛化，并在 Dyck 上优于 SSM 和 Transformer。 $S_3$ 的性能排序（Recurrent > SSM > Transformer）与 Parity 一致，而非 Dyck。
探测剖面：
- Parity/ $S_3$ ： Mamba 和 LSTM 在最后一层（4 层模型中的 $L_3$ ）显示出探测准确率的剧烈跳跃。Transformer 则在 $L_1$ 到 $L_3$ 之间表现出逐渐增加的趋势。
- Dyck： Transformer 在 $L_0$ 或 $L_1$ 处达到高准确率。Mamba 和 LSTM 显示出逐渐增加的过程，仅在较晚的层才达到饱和。
- $S_3$ 确认： $S_3$ 遵循 Parity 剖面（循环模型/SSM 集中于深层，Transformer 逐渐构建），证实了计算结构假设。
因果干预：
- 在 4 层形式化模型中，消融最佳探测方向会导致准确率大幅下降，且这些方向在更长的序列长度（OOD）下仍然是必要的。
- 在预训练模型中，出现了分歧：
  - Pythia (Dyck)： 存在一个强大的中间层瓶颈（在 160M 模型中，在 $L_6$ – $L_7$ 处的消融导致准确率下降约 81%）。最佳探测层（ $L_{11}$ ）并非因果上最关键的层。
  - Mamba (Parity/Dyck/ $S_3$ )： 最后一层具有高度可读性，但单方向消融导致的下降微乎其微（<0.2%）。然而，在最后一层进行中间位置激活修补可以恢复 ~97–98% 的 Logit 差距，这表明状态是充分的，但分布在许多方向上，无法被单个线性探测所捕捉。
Conv1D 归因： 移除 Mamba 中的 Conv1D 会导致 $S_3$ （-12.64 pp）和 Party（-6.83 pp）的性能大幅下降，但在 Dyck 上没有显著影响（+4.04 pp，不显著）。这表明组件的重要性取决于任务。

意义与主张

本文认为，如果不指定任务的计算结构，询问“架构 X 如何编码状态？”是不完整的。作者主张：

机制签名是联合属性： 状态编码策略不是固定的架构特征，而是由模型的归纳偏置与任务的计算需求（前缀更新 vs. 栈）之间的相互作用产生的。
方法论上的谨慎： 从一种任务范式（如前缀更新）推断出的机制发现，若直接应用于另一种范式（如有界栈追踪），可能会误将任务特定的特征误认为架构的不变性。
可读性 vs. 因果性： 探测识别了信息在何处线性可用，但这并不总是映射到计算瓶颈所在的层，尤其是在冗余度较高的预训练大型模型中。
代数 vs. 计算： 在这些形式化任务的语境下，计算结构（顺序前缀更新 vs. 分层栈）比代数属性（如交换性）能更好地预测层级编码剖面。

作者总结道，其发现不仅适用于形式化任务，也适用于预训练规模（高达 410M 参数）和半真实代码基准，这表明序列建模中架构与任务结构之间存在稳健的相互作用。