Transformers Linearly Represent Highly Structured World Models¶

会议: NeurIPS 2026
arXiv: 2605.18847
OpenReview: https://openreview.net/forum?id=8RPfEpnyrk
领域: 其他
关键词: 机制可解释性、世界模型、线性表征假说、注意力头约束消除、裸单数电路

一句话总结¶

本文通过对在数独求解轨迹上训练的 8 层 Transformer 进行机制可解释性分析，发现模型并未按表面棋盘逐单元格建模，而是以行、列、宫这 27 个约束子结构为基本代数单元构建线性世界模型，并通过中层注意力头消除候选数与顶层 MLP 中的单义性“裸单数电路”实现稀疏决策。

研究背景与动机¶

在大语言模型与序列决策模型的研究中，Transformer 在自回归学习推理轨迹时是否会自发构建底层任务的内部“世界模型”，是机制可解释性领域的核心议题。此前以 OthelloGPT 为代表的网格类博弈研究表明，仅基于棋谱序列训练的自回归模型能够在残差流中线性表征棋盘的占用状态；国际象棋与迷宫导航任务上也相继验证了类似的线性表征现象。然而，这些经典任务都具有高度局域化的几何结构，模型内部表征的最小信息单元恰好契合人类最直观的表面网格划分——即每个单元格（cell）对应独立的离散特征，这引发了一个更根本的理论疑问：模型表征空间的几何结构究竟只是对输入表面形态的浅层复刻，还是能真正映射任务内在的约束代数？

数独（Sudoku）提供了一个结构截然不同的绝佳代数试验场。数独不仅是在 \(9 \times 9\) 网格上填入数字，其核心计算实质是 27 个相互重叠的约束子结构（9 行、9 列、9 个 \(3 \times 3\) 宫）之间的约束传播。任一单元格的合法性并非取决于该单元格自身，而是由其所属的行、列、宫中数字的存在性严格决定的。如果模型仅像处理黑白棋那样在单元格级别分别表征状态，将面临严重的特征冗余与跨单元格约束计算复杂度；而若模型学会直接在子结构维度组织特征，则意味着神经网络具有根据问题内在约束逻辑重构隐空间几何的能力。

本文基于带回溯的 Norvig 风格数独求解轨迹训练了一个标准因果 Transformer，并采用线性探针、激活修补（Activation Patching）、直接 Logit 归因（DLA）等可解释性工具对残差流展开端到端逆向工程。核心 idea：Transformer 在组合约束任务中构建的世界模型并非按表面单元格组织，而是直接以约束子结构为几何基底线性表征候选数，并通过中层约束消除注意力头与顶层 MLP 稀疏单义的“裸单数电路”协同完成符号推理与决策提交。

方法详解¶

整体框架¶

本文针对一个 8 层、8 头的标准因果 Decoder-only Transformer 展开机制分析。输入为包含初始线索与求解过程（包括试错与回溯）的 token 序列，模型词表包含 729 个单元格赋值 token（形如 [RrCc=d]）以及 [clues_end]、[push]、[pop]、[success] 等控制 token。分析的核心切入点位于线索结束位置 [clues_end]，该 token 汇聚了初始盘面的全部信息，且处于一切后续自主填数推理之前。

整体推理机制呈现清晰的阶段性分工：浅层（L1–L3）完成初始 token 的局部特征提取与归纳；中间层（L4–L7）发生表征相变，自发建立覆盖全部 27 个子结构（9 行、9 列、9 宫）的线性世界模型，并通过中层专用注意力头实现对称的候选数消除与提升；最终决策层（L8）则由两步级联完成——L8 注意力块在残差流中将合法候选数排至首位，紧接着 L8 MLP 中极稀疏的单义“裸单数检测神经元”被单候选状态激活，对该单元格的所有数字施加统一的 Logit 跃迁，完成决策提交。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入求解轨迹<br/>线索序列 + [clues_end]"] --> B["浅层局部特征提取<br/>L1–L3 线性可分性较低"]
    B --> C["子结构世界模型相变<br/>L4–L7 建立 243 个子结构线性表征"]
    C --> D["中层注意力头约束消除<br/>QK 定位子结构，OV 对称抑制已有数字"]
    D --> E["候选数概率前向累积<br/>L8 注意力块输出将合法解排至首位"]
    E --> F["顶层 MLP 裸单数电路<br/>91 个单义神经元触发单元格 Logit 跃迁"]
    F --> G["最终动作输出<br/>预测目标填数 token [RrCc=d]"]

关键设计¶

1. 子结构世界模型：按约束代数而非表面单元格组织状态 传统机制可解释性假设模型倾向于在最小物理单位（单元格）上线性编码状态。但在数独中，直接训练 81 个单元格多分类线性探针预测各格数字时，Top-1 准确率在第 4 层后仅停留在 80% 左右（MSE 约为 0.03），说明单元格填充状态并未被干净地线性表征。相反，当训练 729 个二分类探针预测“单元格 \((r,c)\) 是否以数字 \(d\) 为合法候选数”时，探针在中层达到了 100% 的完全匹配率（Exact Match）。深入分析这 729 个候选探针法向量的余弦相似度发现：共享 1 个子结构的单元格探针相似度约为 0.33，共享 2 个子结构（如同时在同行同宫）的相似度达到 0.65，而无共享子结构的单元格近乎正交。这直接证明单元格候选数向量并非独立存储，而是由子结构特征线性组合而成。研究团队进而训练 243 个探针直接预测全局命题“数字 \(d\) 是否存在于子结构 \(S\) 中（9行+9列+9宫各9个数字）”，发现其在中层不仅达到完美的 100% 准确率与 1.0 的 ROC-AUC，且在跨层迁移中展现出在第 4 层的剧烈相变（L4–L7 探针可零样本无损跨层泛化），在整个解题步长推进过程中表征方向保持严格恒定。

2. 子结构对齐的中层注意力头：对称消除与提升候选数 模型在残差流中建立子结构表征后，中层的多头注意力机制扮演了去中心化“约束检查器”的角色。注意力权重分析表明，大量中层注意力头在 QK 回路上高度聚焦于特定的单一行、列、宫，或是由 3 个相连宫组成的行带（band）与列堆（stack）。在 OV 回路上，这些专用注意力头执行严格对称的代数消除机制：通过直接 Logit 归因（DLA）将单个头的输出投影到反嵌入矩阵，发现当某数字 \(d\) 已经存在于其聚焦的子结构中时，该头会在残差流中显著压低该子结构内所有单元格填入 \(d\) 的 Logit；反之，若 \(d\) 在该子结构中尚未出现，则对该子结构内的候选位置施加正向提升。反事实均值消融（Mean Ablation）证实，遮蔽特定子结构专用头会导致对应区域非法填数的 Logit 显著上升 0.94–1.69，而对照区域不受干扰，验证了该约束消除操作的因果局域性。

3. 顶层 MLP 裸单数电路：单义神经元触发全局提交 数独中最基础的确定性填数法则为“裸单数”（Naked Single，即某个空格所属的行、列、宫已出现全部其他 8 个数字，仅剩唯一候选数）。Logit Lens 显示，在残差流刚穿过 L8 注意力块、尚未进入 L8 MLP 时，正确的唯一数字就已经在 98.77% 的盘面状态下位列单元格候选排序第一位，表明中层已经完成了充分的消解运算。随后，L8 MLP 并非用于重新判断哪一个是正确答案，而是负责“拍板提交”。通过计算神经元在单候选状态与多候选状态下的激活差值（阈值设为 3.0），研究团队在 L8 MLP 的 3,456 个神经元中精确隔离出了 91 个专用的裸单数检测神经元。这 91 个神经元展现出极高的单义性：每个神经元严格绑定特定的单元格（71 个单元格各自对应 1 个神经元，10 个单元格对应 2 个），仅在该格缩减为单候选时剧烈激发。DLA 投影显示，激活时该神经元对所属单元格的全部 9 个数字 token 施加约 +3.00 的均匀 Logit 提升（9 个数字间的提升标准差仅 0.10），而对其余 80 个空格几乎无影响（\(\mu=-0.02\)）。由于正确数字已在前序层建立了显著优势，这一格内广播式的均匀加分使正确答案在经 Softmax 归一化后概率压倒性集中，实现极度稀疏、可解释的决策提交。

损失函数 / 训练策略¶

模型采用标准自回归因果语言建模目标（交叉熵损失），但特别设计了掩码策略：仅对序列中 [clues_end] 之后的求解步骤 token 计算损失，对初始盘面线索 token 不予监督。这确保了网络专注于根据约束规律自发推演解答，而非死记硬背输入线索。模型在包含 270 万道题目的 sudoku-3m 数据集上训练 6 个 epoch，采用 AdamW 优化器（学习率 \(10^{-3}\)，权重衰减 0.1，Batch Size 512，5M token 线性预热配合余弦退火），在单张 Google TPU v6e 上完成训练，最终在 15 万道未见过的测试题上实现了 98.4% 的单步单元格准确率与 97.5% 的整盘求解准确率。

实验关键数据¶

主实验¶

为了验证子结构表征是否因果性地参与下一步预测，研究团队在 [clues_end] 处收集了 300 组盘面状态对 \((G1, G2)\)，其中 \(G2\) 为在单元格 \((r,c)\) 填入数字 \(d\) 的真实后续步。通过将 \(G2\) 中对应行、列、宫的三组子结构探针分量反事实修补（Patching）至 \(G1\) 的残差流中（伪造出该数字已在所属子结构出现的假象），观察目标 token 的 Logit 变化：

\[x^{\prime}_{G1} = x_{G1} + \sum_{S \in \{\mathsf{Row}_r, \mathsf{Col}_c, \mathsf{Box}_k\}} \bigl(\hat{w}_{S,d}^\top x_{G2} - \hat{w}_{S,d}^\top x_{G1}\bigr) \hat{w}_{S,d}\]

残差流层级	目标 Logit 下降值 (\(\Delta \text{logit}\))	修补后绝对 Logit	Top-1 预测改变比例 (%)	改变后仍为合法走法比例 (%)
L0	\(13.978 \pm 0.311\)	2.845	98.7%	72.3%
L1	\(8.724 \pm 0.293\)	8.099	95.3%	69.0%
L2	\(15.054 \pm 0.295\)	1.769	99.3%	72.3%
L3	\(23.727 \pm 0.332\)	-6.904	99.3%	76.3%
L4	\(24.766 \pm 0.313\)	-7.943	99.7%	72.7%
L5	\(22.124 \pm 0.285\)	-5.301	99.7%	65.0%
L6	\(21.775 \pm 0.264\)	-4.952	99.7%	59.7%
L7	\(19.235 \pm 0.243\)	-2.413	99.3%	60.7%

原始未修补状态下模型目标 token 的平均 Logit 为 16.82。由上表可见，修补在 L4 达到峰值，造成近 25 个 Logit 的剧烈压制，并在 99.7% 的情况下成功改变了模型决策，且超过 70% 的新预测依然符合盘面其他合法走法，表明修补并未摧毁隐空间的全局连贯性。

消融实验¶

为验证中层注意力头和顶层 MLP 神经元的机制特异性，论文进行了两组严格的因果消融实验：

表 1：中层专用注意力头均值消融实验 通过替换特定头在 [clues_end] 处的激活为其均值，测量其目标子结构与非目标对照子结构中非法填数 Logit 的变化（\(\mathbb{E}[\ell_{\text{ablated}} - \ell_{\text{clean}}]\)）：

专用注意力头	负责区域	目标区域非法 Logit 变化 (\(\Delta \text{logit}\))	对照区域非法 Logit 变化 (\(\Delta \text{logit}\))
L4H6	列 4–6 (Stack 2)	\(+1.426 \pm 0.004\)	\(-0.170 \pm 0.002\)
L5H8	行 7–9 (Band 3)	\(+1.690 \pm 0.004\)	\(-0.216 \pm 0.002\)
L6H3	宫 5 (Box 5)	\(+0.937 \pm 0.005\)	\(-0.108 \pm 0.003\)

表 2：顶层 MLP 裸单数专用神经元消融实验 在包含两个同时存在的裸单数盘面上，仅消融模型正要填入的那个目标单元格关联的 NS 神经元，观察目标与非目标位置的指标降幅：

评估位置	预测 Logit 降幅 (\(\text{clean} - \text{ablated}\))	Softmax 概率降幅
目标裸单数单元格	\(11.408 \pm 0.084\)	\(0.585 \pm 0.003\)
盘面另一个共存裸单数格	\(-0.655 \pm 0.008\)	不受负面影响 (概率转移承接)

关键发现¶

子结构探针与单元格探针的分野：243 个子结构探针在 L4–L7 的分类准确率与 ROC-AUC 均达到 1.0，而单元格状态探针最高仅 80%。这揭示出模型内部并不是记录“第几格填了什么”，而是记录“哪个子结构缺什么”。
第 4 层相变特性：跨层探针泛化矩阵表明，L1–L3 完全不具备后续可读的子结构特征，而在第 4 层该几何结构瞬间定型，并直接稳定保持到第 7 层。
解题全程的方向不变性：尽管随着空格减少、类别分布改变导致固定偏置的精确匹配率下降，但冻结探针的 ROC-AUC 在整个求解过程中始终为 1.0，说明世界模型的代数几何方向在序列生成中是静态恒定的。
决策机制的双阶段解耦：中层注意力负责“排除不可行项并排序”，最终层 MLP 的 91 个稀疏神经元负责“检测唯一解并执行放大提交”，二者在结构和因果上完全解耦。

亮点与洞察¶

证明了世界模型几何由任务代数决定而非表面形态：打破了过去认为“棋盘任务必然以单元格为基本表征单元”的认知惯性，有力证明了深度网络表征学习具备发现底层代数对称性的能力。
发现了极为稀疏且单义的“裸单数电路”：在现代大型前馈层（3,456 维）中定位到了功能极其明确的 91 个单义神经元，每个神经元直接对应具体物理单元格的单候选判定并提供均匀 Logit 增益。
对复杂逻辑任务实现了端到端算法闭环解析：从特征构建（L4 子结构相变）、候选消解（中层 QK/OV 对称抑制）到动作触发（L8 NS 神经元提交），完整还原了 Transformer 执行符号约束推理的微观神经回路。

局限与展望¶

模型规模与任务单一性：作者坦承实验仅在单个 8 层 Transformer 和数独这一特定谜题上验证，尚未拓展至更大参数模型或其它不同类型的逻辑约束问题（如 SAT 求解或图染色）。
注意力 QK 机制刻画尚不完全：本文详尽分析了中层注意力头的 OV 回路及其 Logit 抑制机制，但尚未完全逆向出注意力权重如何根据输入前序序列动态聚焦到特定行/列/宫的 QK 寻址算法。
探针修补协议单一：当前反事实因果证据依赖于特定的子结构探针方向替换协议，未来需要结合路径修补（Path Patching）或稀疏自动编码器（SAE）等手段进一步检验表征的抗扰动性。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 首次揭示 Transformer 世界模型由约束代数而非表面单元格决定的机制，发现了完备的裸单数单义电路。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖线性探针、余弦几何分析、跨层跨步泛化、反事实探针修补、注意力头 DLA/消融以及神经元级因果消融。
写作质量: ⭐⭐⭐⭐⭐ 逻辑极为严密，从宏观表征到微观回路环环相扣，实验设计与图表呈现清晰规范。
价值: ⭐⭐⭐⭐⭐ 为大模型复杂逻辑推理与世界模型涌现提供了极具说服力的机制可解释性范本。