Transformers Linearly Represent Highly Structured World Models¶

会议: NeurIPS 2026
arXiv: 2605.18847
领域: 可解释性
关键词: 机理可解释性、线性世界模型、约束代数、注意力头回路、单义神经元

一句话总结¶

通过对求解数独的 8 层 Transformer 进行机理可解释性分析，发现模型内部并未按棋盘单元格建构表征，而是自发形成了按行、列、宫 27 个子结构对齐的线性世界模型，并在输出端依赖由 91 个单义神经元组成的显式单一解决策回路。

研究背景与动机¶

探究序列模型在复杂推理任务中是否自发构建环境的内部世界模型，是机理可解释性领域的核心课题。此前在黑白棋（OthelloGPT）、国际象棋以及迷宫求解等棋盘网格任务上的经典工作表明，模型倾向于在其残差流中以线性方向编码任务状态。然而，这些先验研究中的内部表征几何形态几乎无一例外地严格对应物理表面分解——即以单个单元格（cell）为基本特征单元。这种现象促使研究者思考一个深层问题：当环境不再是简单的局部棋盘状态，而是受强全局代数约束驱动的组合约束满足问题时，模型学到的世界模型究竟会固守单元格的表面几何，还是会自发对齐问题的底层约束代数？

数独（Sudoku）提供了一个极具辨识度的测试基准：表面上是 81 个单元格构成的 $9 \times 9$ 物理网格，但任何单元格的填数合法性完全由 27 个相互重叠的子结构（9 行、9 列、9 宫）所约束。若模型按照单元格状态建模，表征难以直接反映规则的代数本质；而若直接以子结构是否包含某数字建模，不仅大幅降低了约束传播的复杂度，也使候选数的排除可以直接通过子结构状态的交集来线性表达。

本文在带有回溯机制的数独求解轨迹上训练了标准的 8 层因果 Transformer，并采用线性探针与因果干预手段剖析其内部计算回路。核心 idea：Transformer 在组合逻辑推理中构建的世界模型几何并不取决于任务的表面输入形式，而是由领域的底层约束代数所塑造，并在中间层形成子结构消除机制、在末层 MLP 由高度稀疏的单义神经元完成最终动作承诺。

方法详解¶

整体框架¶

本文针对 8 层、8 头因果 Transformer（维度 $d_{\mathrm{model}}=576$，$d_{\mathrm{mlp}}=3456$）开展端到端的逆向工程分析。模型输入包含题目提示线索（clues）与求解过程中的试探、回溯标记，词表涵盖 729 个动作 token（形如 [RrCc=d]）。分析以提示线索结束符 [clues_end] 为核心锚点，探查其残差流在不同层级的语义承载形态。

整个方法的机制流动分为三个主要阶段：首先，浅层残差流将线索聚合，并在第 4 层发生显著相变，完成从离散线索到 27 个子结构全局约束状态的线性编码；其次，中层（第 4 至 7 层）的注意力头在空间上特异性聚焦特定子结构，其 OV 回路执行对称的约束消除计算——抑制子结构中已出现数字的候选 logit、提升未出现数字的 logit；最后，第 8 层注意力输出已将正确数字排在候选首位，而最终的 MLP 块通过一组专职的显式单一解神经元识别唯一候选状态，向该格所有 token 注入均匀正向偏置，从而在 Softmax 下实现强力动作承诺。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入求解轨迹<br/>提示线索及回溯标记"] --> B["浅层网络 (L1-L3)<br/>线索聚合与初步处理"]
    B --> C["子结构世界模型 (L4 相变)<br/>243 个 (子结构, 数字) 线性方向形成"]
    C --> D["子结构对齐注意力与约束消除 (L4-L7)<br/>专用头按行/列/宫抑制已占数字"]
    D --> E["残差流前向累积 (L8 Attn 后)<br/>正确候选数字以高 margin 排名第一"]
    E --> F["显式单一解激活回路 (L8 MLP)<br/>91 个单义神经元触发单元格整体 Logit 放大"]
    F --> G["Softmax 输出预测<br/>准确输出下一填数动作"]

关键设计¶

1. 子结构世界模型：超越表面单元格的代数对齐表征

针对“模型究竟是按单元格还是按约束结构记忆棋盘状态”的疑问，作者对比了三套线性探针体系：针对 81 个单元格中填充数字的 9 分类探针、针对 729 个“数字 $d$ 是否为单元格 $(r,c)$ 有效候选”的二分类探针，以及针对 243 个“数字 $d$ 是否已存在于子结构 $S$（行、列、宫）”的二分类探针。实验显示，单元格状态探针在第 4 层后准确率停滞在 0.80 且均方误差不归零；而针对 243 个子结构的探针在第 4 至 7 层达到 1.0 的完全精确匹配（Exact Match），且 ROC-AUC 在整个求解过程中保持为 1.0。

这种表现差异源于数独约束的逻辑非对称性：候选数集合 $\mathsf{cand}(r,c) = \{1,\dots,9\} \setminus (\text{digits in } \mathsf{Row}_r \cup \mathsf{Col}_c \cup \mathsf{Box}_k)$ 完全是所属子结构存在状态的布尔代数运算，因此候选状态是子结构状态的直接线性组合。测量 729 个单元格候选探针法向量之间的余弦相似度发现，共享行和宫两个子结构的单元格探针相似度高达 $\approx 0.65$，共享一个子结构的相似度 $\approx 0.33$，而不共享子结构的探针相互正交。这从几何上证明模型并未维护 81 个独立的候选特征，而是完全由 243 个子结构基底线性组合而成。

2. 子结构对齐注意力与约束消除机制：基于 OV 回路的对称抑制

在中层残差流建立起子结构表征后，模型需要通过具体的注意力组件读取并操作这些特征。分析第 4 至 7 层的注意力头发现，多头注意力并没有在整个棋盘上均匀扩散，而是高度特异性地将注意力权重集中在特定的几何子结构上，如特定的某一列、某一行、某一个 $3 \times 3$ 宫，或者横向 band、纵向 stack。

通过直接 Logit 归因（DLA）将这些注意力头的 OV 电路输出投影到未嵌入矩阵，揭示了其优雅的对称消除机制：当数字 $d$ 已经出现在该头负责的子结构中时，该头会在该子结构的所有位置上对数字 $d$ 的候选 logit 产生显著负向抑制（并在相邻结构中轻微提升）；反之，当数字 $d$ 缺失时，该头则正向促进该数字。对专用注意力头进行均值消融（Mean Ablation）会导致目标子结构内非法数字的 logit 显著上升 0.94 至 1.69，确证了这些注意力头扮演着局域约束校验器的因果角色。

3. 显式单一解激活回路与单义神经元：两阶段决策承诺

对于数独中最基础的必填步——显式单一解（Naked Single，即某空格仅剩唯一可能填数），决策是如何最终敲定的？Logit Lens 分析显示，在经过第 8 层自注意力层后、尚未进入最终 MLP 时，残差流中正确填数的 logit 已经有 98.77% 的概率排名第一，且与第二名的差距极大。这意味着注意力层已经完成了候选数的逻辑推导，最终 MLP 的功能并非寻找答案，而是做出动作“承诺”（Commitment）。

作者通过统计神经元在目标单元格处于单候选状态与其他状态下的激活差异（Activation Gap），以 3.0 为阈值在第 8 层 MLP 中精准筛选出 91 个专职的“显式单一解神经元”。这 91 个神经元展现出极高的单义性：全盘 81 个格子中，71 个单元格恰好对应 1 个专属神经元，10 个单元格对应 2 个神经元，实现了对全盘所有单元格的完全覆盖。当某单元格出现唯一候选时，对应神经元剧烈激活，通过未嵌入层向该单元格对应的全部 9 个数字 token 均匀注入约 +3.00 的 logit 增益（标准差仅 0.10），而对其他 80 个格子的投影均值接近 0（$\mu = -0.02$）。由于正确候选在注意力输出端已确立优势，这一均匀的全局增益经 Softmax 非线性放大后，将概率质量几乎全部归集到该单元格的正确解上。

一个完整示例¶

以单元格 $(7,4)$ 出现显式单一解且正确数字为 3 为例： 1. 中层约束累积：在经过 L4 至 L7 层时，负责第 7 行、第 4 列以及对应中下宫的注意力头（如 L4H6 负责 4-6 列、L6H3 负责第 5 宫）分别读取盘面信息。由于 1、2、4、5、6、7、8、9 分别存在于 $(7,4)$ 所在的行、列或宫中，这些数字对应的 logit 被专用头的 OV 回路依次抑制，而数字 3 因在三者中均未出现而得以保留并被正向赋权。 2. 残差流状态：到达 L8 注意力层末端时，token [R7C4=3] 的 logit 达到 16 左右，显著高于该格其他 8 个数字 token（margin 分布在 +8 至 +15）。 3. MLP 动作承诺：L8 MLP 中负责单元格 $(7,4)$ 的专属单义神经元检测到残差流传来的“该格仅剩单候选”特征并点火，对该格所有 9 个 token 输出均匀的 +3.0 logit 提升，而对盘面其他 80 个格子的 720 个 token 几乎无干扰。 4. 最终输出：Softmax 接收到底部拉开的 margin 加上整体抬升的数值，将超过 90% 的概率集中在 [R7C4=3]，模型自信输出预测。

损失函数 / 训练策略¶

模型采用标准因果自回归下一 token 交叉熵损失： $$ \mathcal{L} = -\sum_{t > t_{\text{clues_end}}} \log P(x_t \mid x_{<t}) $$ 损失掩码仅应用于 [clues_end] 之后的求解与回溯步骤，模型不需要背诵初始线索。训练数据集源自 sudoku-3m 中的 270 万道题目，求解轨迹经由 Norvig 求解器随机化生成（打乱单一解发射顺序、候选探索顺序及回溯分支），截断至 250 tokens。在单张 Google TPU v6e 上使用 AdamW 优化器（学习率 $10^{-3}$，权重衰减 0.1，Batch Size 512，5M tokens 预热加余弦退火）训练 6 个 epoch，在 15 万道测试题上达到 98.4% 的逐格准确率与 97.5% 的整盘求解准确率。

实验关键数据¶

主实验¶

为了验证子结构表征方向是否真正参与因果计算，作者采用反事实激活补丁（Counterfactual Activation Patching）干预残差流，将真实有效填数对应的行、列、宫探针方向强制改写为“该数字已存在”，并观测目标 token 的 logit 塌陷程度。

层级	目标 Token Logit 下降量	补丁后目标 Logit	补丁后 Top-1 仍合法比例	Top-1 发生改变比例
L0	$13.98 \pm 0.31$	2.85	0.723	0.987
L1	$8.72 \pm 0.29$	8.10	0.690	0.953
L2	$15.05 \pm 0.30$	1.77	0.723	0.993
L3	$23.73 \pm 0.33$	-6.90	0.763	0.993
L4	$24.77 \pm 0.31$	-7.94	0.727	0.997
L5	$22.12 \pm 0.29$	-5.30	0.650	0.997
L6	$21.78 \pm 0.26$	-4.95	0.597	0.997
L7	$19.24 \pm 0.24$	-2.41	0.607	0.993

注：原正常前向传播下目标 Token 的平均 Logit 值为 16.82。

消融实验¶

1. 中层特异性注意力头均值消融实验 检验注意力头是否作为局域约束校验器独立发挥作用。在 [clues_end] 处将专用头的激活替换为测试集均值，测试目标子结构与非目标对照子结构中非法填数 token 的 logit 变化。

注意力头	负责子结构区域	目标区域非法 Logit 变化 ($\Delta \text{logit}$)	对照区域 Logit 变化 ($\Delta \text{logit}$)	说明
L4H6	第 4–6 列 (cols 4–6)	$+1.426 \pm 0.004$	$-0.170 \pm 0.002$	专职消除中三列已存在数字
L5H8	第 7–9 行 (rows 7–9)	$+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 显式单一解神经元消融实验 在包含两个同时并存的显式单一解棋盘状态上，对模型即将预测的单一解目标神经元执行消融，观察双目标的因果扰动。

评估对象	Logit 下降量	Softmax 概率下降量	说明
目标单一解填数 (Target NS)	$+11.408 \pm 0.084$	$+0.585 \pm 0.003$	专属神经元消融导致目标概率直接腰斩
同盘另一单一解 (Other NS)	$-0.655 \pm 0.008$	概率反向微增	概率质量精准转移至另一合格单一解

关键发现¶

表征相变点明确：跨层探针迁移矩阵显示在第 4 层存在强烈的相变。L1–L3 训练的探针无法迁移至后续层，而 L4–L7 训练的子结构探针在整个中层区间几乎具备 100% 的无缝互测精度。
动态求解过程表征几何不变：虽然随着棋盘填满，固定阈值的 Exact Match 探针因先验分布偏移而下滑，但 ROC-AUC 在整个求解过程中恒定为 1.0，表明底层子结构方向的线性可分几何始终稳固。
单义神经元的极高特异性：91 个显式单一解神经元对非所属格子的全局投影均值仅为 -0.02，且目标格 9 个数字的提升方差仅 $\sigma = 0.10$，证明了其在空间维度的高度局域性与数字维度的均匀性。

亮点与洞察¶

任务代数决定内隐表征几何：该工作首次在强约束组合推理任务中确证，神经网络构建的世界模型不会被输入的网格坐标（cell-by-cell）所局限，而是精准自发解构为满足约束传播所需的子结构代数（substructure algebra）。这一洞察为探索大型语言模型在逻辑、代码等高结构化领域的认知机制提供了理论支点。
“先推导排序、后统一放大”的两阶段决策解耦：发现注意力和前馈层在复杂推理输出时具有明确的分工模式——注意力头通过局域布尔过滤确定谁是正确答案，MLP 专职神经元则充当阈值开关，检测到单解状态后整体提升 Logit，机制清晰透明。
可复用的回路分析范式：将探针法向量余弦聚类分析与 DLA、跨层迁移矩阵及反事实因果补丁深度结合，形成了一套可推广至数理逻辑与规划任务的完整机理逆向工程范式。

局限与展望¶

模型规模与任务单一性：论文结论基于单一定制架构（8 层、8 头因果 Transformer）及数独单任务，尚未在更大参数规模（如数十亿参数 LLM）或其它组合问题（如 SAT、图着色）中验证该代数子结构对齐规律的普适性。
QK 电路注意力路由尚未完全解析：论文主要完整还原了注意力头的 OV 写入回路，但关于 Query-Key 电路如何精确引导注意力头锁定特定行、列、宫的寻址机制仍未完全形式化解析。
训练动力学与顿悟演变：未对子结构几何表征与 91 个专职神经元是在训练早期逐步形成还是经历类似 Grokking 的突变跃迁做深入追踪。

评分¶

新颖性: [5/5] 首次系统性论证了组合约束满足任务中世界模型与约束代数而非表面空间单元格的对应关系。
实验充分度: [5/5] 线性探针、余弦聚类、跨层跨步迁移、DLA 及多层级因果干预构成了严密的因果证据链。
写作质量: [5/5] 逻辑极为严谨清晰，问题定义锐利，图表与附录数据支撑扎实。
价值: [5/5] 为大模型内部如何表征和执行形式化符号逻辑推理提供了极具说服力的微观机理范本。

层级	目标 Token Logit 下降量	补丁后目标 Logit	补丁后 Top-1 仍合法比例	Top-1 发生改变比例
L0	\(13.98 \pm 0.31\)	2.85	0.723	0.987
L1	\(8.72 \pm 0.29\)	8.10	0.690	0.953
L2	\(15.05 \pm 0.30\)	1.77	0.723	0.993
L3	\(23.73 \pm 0.33\)	-6.90	0.763	0.993
L4	\(24.77 \pm 0.31\)	-7.94	0.727	0.997
L5	\(22.12 \pm 0.29\)	-5.30	0.650	0.997
L6	\(21.78 \pm 0.26\)	-4.95	0.597	0.997
L7	\(19.24 \pm 0.24\)	-2.41	0.607	0.993

注意力头	负责子结构区域	目标区域非法 Logit 变化 (\(\Delta \text{logit}\))	对照区域 Logit 变化 (\(\Delta \text{logit}\))	说明
L4H6	第 4–6 列 (cols 4–6)	\(+1.426 \pm 0.004\)	\(-0.170 \pm 0.002\)	专职消除中三列已存在数字
L5H8	第 7–9 行 (rows 7–9)	\(+1.690 \pm 0.004\)	\(-0.216 \pm 0.002\)	专职消除下三行已存在数字
L6H3	第 5 宫 (box 5)	\(+0.937 \pm 0.005\)	\(-0.108 \pm 0.003\)	专职消除中心宫已存在数字

评估对象	Logit 下降量	Softmax 概率下降量	说明
目标单一解填数 (Target NS)	\(+11.408 \pm 0.084\)	\(+0.585 \pm 0.003\)	专属神经元消融导致目标概率直接腰斩
同盘另一单一解 (Other NS)	\(-0.655 \pm 0.008\)	概率反向微增	概率质量精准转移至另一合格单一解