受约束的信念更新解释 Transformer 表示中的几何结构

next-token 的 optimal prediction 需要递归贝叶斯信念更新，注意力却只能并行加权；作者要解释这一约束如何生成中间几何。

• 出现场景：作者称 transformer 擅长 next-token prediction，但 optimal prediction 要维持对未来的信念，这是递归过程；架构却强制并行注意力（Fig. A1）。作者预期同一张力也会出现在自然语言上的大语言模型中，本文只在可解析的 HMM 上检验。
• 已有结果未解释的部分：Shai et al. (2024) 报告残差流中有信念状态几何。作者称那对应“表示了什么”；本文要解释注意力如何把贝叶斯更新弯成可并行形式，以及中间分形为何不同于最终几何。
• 核心假设：作者提出 constrained belief updating。单层注意力的每个 value 只含源位置局部信息，因此只能把各历史 token 相对平稳分布的独立修正相加，而不能做转移矩阵连乘。
• 本文的分析：在 Mess3 HMM 上训练 transformer，主要拆单层第一注意力层，用谱分解预测注意力模式、OV 向量和嵌入，并检查 MLP 与 4 层网络如何把该表示推向完整贝叶斯信念。

论文第 2.1 节把相关工作分成三组。

激活空间中的特征方向

• 线性方向：Park et al. (2024) 用激活与特征方向的对齐来看表示。Elhage et al. (2022) 用非正交向量讨论 superposition，即特征数可以多于神经元数。
• 词典学习与非线性：Cunningham et al. (2023)、Bricken et al. (2023)、Templeton et al. (2024) 用稀疏自编码器或词典学习抽取特征。Engels et al. (2024) 报告并非所有语言模型特征都是线性的；作者称本文给这类非正交几何一个机制上的 why。

从特征到电路

• 已识别的电路：Elhage et al. (2021) 给出 OV/QK 分解并讨论句法电路；Wang et al. (2022) 分析间接宾语识别；Nanda et al. (2023) 分析算术。作者称识别电路大体仍是从行为往回找的手工过程。
• 自动化方向：Conmy et al. (2023)、Marks et al. (2024) 研究自动电路发现。论文将其列为进行中的方向，未与本文做对照实验。

信念几何与计算力学

• optimal predictor 要表示什么：Shalizi & Crutchfield (2001)、Crutchfield (2012)、Riechers & Crutchfield (2018) 的计算力学，以及 Kaelbling et al. (1998) 的 POMDP、Upper (1997)，把序列预测写成对隐状态信念的维持。Crutchfield (1994)、Marzen & Crutchfield (2017) 把信念画在概率单纯形上。
• 网络中的几何：Pepper (2024) 报告 RNN 隐状态表示信念几何；Shai et al. (2024) 报告 transformer 残差流中有信念状态几何。作者称本文进一步提出约束信念更新，并检验注意力如何改变理想信念动力学。

基线与数据

• 对照：实验对照是 Eq. (2) 的完整贝叶斯信念、Eq. (5) 的约束信念，以及随机初始化的 MSE，不是另一族检测器。数据过程是 Mess3（Marzen & Crutchfield, 2017）。

作者把本文放在模型无关的贝叶斯信念与注意力并行形式的交叉处：不只核对是否编码信念，还要预测注意力模式、value 向量和残差流几何如何偏离无约束更新。

整体做法是在 Mess3 上训练 transformer 做 next-token 预测，用 constrained belief updating 解释第一层注意力的中间表示，再看 MLP 如何把它变成完整信念几何。

数据过程与两种信念

• Mess3：三个隐状态，词表为 {0,1,2}。α∈[0,1] 越大，各状态越偏向自己的发射符号；x∈(0,1/2] 控制持续性，低 x 惯性高。β=(1−α)/2，y=1−2x。附录 A 给出三个带标签转移矩阵。第 4.4.1 节另写：T 要行随机，x 须在 [0,0.5]。
• 完整更新：平稳过程的初始分布取平稳分布 π。给定整段 token 后，信念是转移矩阵连乘再归一化：

r_full=η^∅T^(z_1:d)/(η^∅T^(z_1:d)1)

作者称这是递归过程，每步依赖此前已整合的历史。

• 约束更新：注意力对源位置 value 做并行加权和，每个 value 只有该位置的局部信息。作者因此写成：未见上下文时用 π，再把每个历史 token 的独立修正加总：

r_1:d=π+∑_s=1^d(π T^|z_sT^d−s−π)

作者称，在单层注意力的限制下，这是能达到的 best possible parallel approximation；它是隐状态上的分布，因而对应整个未来而不只是下一 token。行和可为零的过程改用附录 A 的 Eq. (26)。

训练与表示分析

• 主模型：单层 transformer，dmodel=64，dff=256，学习位置嵌入，无 BOS。按 Mess3 参数用 1 或 2 个注意力头。每组 (α,x) 单独训练。
• 优化：交叉熵，Adam，学习率 10^{-4}，无 weight decay，batch size 128，约 15 million tokens（附录 B）。
• 读出：对长度不超过 10 的全部输入记录残差流，做 PCA；有时把基稍作旋转以对齐理论方向。比较嵌入、注意力后的中间激活，以及 MLP 后的最终激活。4 层网络用来检查第一层是否仍做同一计算；其宽度论文未单独写出。

注意力在单纯形上的算法

• 两条电路：残差流是 skip connection。OV 决定从源位置读出并变换什么；QK 用 query–key 内积和因果 softmax 决定权重。目的位置的更新是对 s≤d 的 value 做加权和。注意力非负，且对源求和为 1。
• Mess3 上的形态：作者在 x=0.15、α=0.6 的模型上报告，token 嵌入在 PCA 中靠近原点成三簇，OV 向量指向三角形三个顶点，并被解释为信念单纯形的顶点。注意力主要由位置距离决定，Fig. 2 图注写按 (1−3x)^{|n−1|} 衰减，且大体与具体 token 序列无关。正文另写回看 n 步的权重约等于 (1−3x)^n。
• 分形如何拼出：中间表示是三个 OV 向量的线性组合；距离越远，过去 token 的权重越小。作者称该计算可以完全看成信念单纯形上的动态过程。

谱分解给出的可检验预测

• ansatz：作者把第一层注意力的作用设为用谱机制实现约束更新。T 可对角化时，距离依赖只留在特征值的幂上。Mess3 的特征值是 1 和 ζ=1−3x，后者重数 2。脚注写：第一层头很多时，这一对应的细节会失效。
• 由此推出的关系：d>s 时 A_{d+m,s}=ζ^m A_{d,s}。同一 token 的 OV 向量互相平行，模与 A_{m+1,m} 成反比。嵌入与 OV 平行。作者称这些在实验中成立，但前两个位置的嵌入有标量偏差，附录 D 中符号与预测相反。
• 负特征值：x>1/3 时 ζ<0，模式振荡，单头不能表示，因为注意力元必须非负。两头分别承担正、负分量，OV 方向相反，并满足 A_{d+2m,s}=ζ^{2m} A_{d,s}。作者称这把数据的转移矩阵和头数联系起来。

MLP 与更深层

• 单层 MLP：作者称它在每个位置做连续非线性变形，拉伸、压缩中间分形并保持拓扑，得到最终信念几何。其数学性质的完整刻画，第 4.5 节写留待以后。
• 4 层：第一层注意力仍实现 Eq. (5)；后续层把表示推向 Eq. (2) 的完整贝叶斯信念。作者把第一层视为更深层里的基础计算。

实验设置

• 被测网络：作者训练的单层 transformer，以及同一 Mess3 数据上的 4 层 transformer。无命名预训练模型。按参数用 1 或 2 个注意力头。4 层的 dmodel 论文未单独写出。
• 数据：Mess3，3 个隐状态、3 个 token。对 (α,x) 扫描；正文与 Table 1 写出的设置包括 x=0.15、x=0.5、α=0.2、α=0.6。训练序列从对应 HMM 采样。表示分析用长度不超过 10 的全部输入，无 BOS。
• 训练：交叉熵，Adam，学习率 10^{-4}，无 weight decay，batch size 128，约 15 million tokens。单层为 dmodel=64、dff=256（附录 B）。
• 读出位置：注意力后、MLP 前的中间残差流，以及 unembedding 前的最终残差流。PCA，有时略旋转到理论方向。
• 对照与指标：Eq. (2)、Eq. (5)，以及随机初始化。指标为对理论几何回归的 MSE、预测相对 optimal 概率的 KL（Fig. A7）、PCA 累积解释方差（Table 1）。无 LLM 评审。
• 协议里没有的数：论文未报告随机种子或重复次数，正文也未给出 KL 和 Fig. A4 柱上的 MSE 数值。

主结果

定量主表是残差流 PCA 的累积解释方差。下表取 Table 1 的 component 2（从 0 编号），即前 3 个主成分合计。

条件	中间 component 2	最终 component 2
x=0.15, α=0.2	0.9673	0.9943
x=0.15, α=0.6	0.9859	0.9811
x=0.5, α=0.6	0.8913	0.8855
x=0.5, α=0.2	0.9173	0.9689

• 作者对维度的说法：作者称有效维度低，多数情况下前 3 个主成分解释超过 90% 的方差。表中 x=0.5、α=0.6 一行的中间与最终都低于 0.90，与“多数”并存。
• x=0.5：作者称该设置有效维度更高，可能与信念更新的振荡有关，并写明还需研究（附录 G）。
• 几何对比不在该表：Figure 1 图注称，最终残差流 PCA 与信念几何相符，注意力后的中间投影是另一种更曲折的结构。点按 ground-truth 信念的三个概率着成 RGB。

注意力模式与理论几何

• 测了什么：x=0.15、α=0.6 的单层模型，比较 OV、嵌入、注意力和中间表示（Fig. 2）。
• 结果：图注称 OV 在三角形顶点成三簇，嵌入靠近原点；注意力主要随位置距离按 (1−3x)^{|n−1|} 衰减，与具体序列关系不大，理论与实际注意力在 inset 中对齐。中间分形由三个 OV 向量的线性组合拼出。
• 作者的解读与例外：作者称这支持注意力在实现约束更新。附录 D 写前两个位置的嵌入与 OV 平行，但符号与预测相反，作者称尚不理解原因。α∈[0.2,0.6] 时 Eq. (5) 与中间结构相符；离开该区间偏差逐渐增大，整体结构仍相似（第 4.3 节、Fig. 4）。

负特征值与注意力头数

• 测了什么：ζ<0（x>1/3）时比较 1 头和 2 头，并在不同头数、层数下看预测与 optimal 概率的 KL（Fig. 3，附录 F，Figs. A7、A8）。
• 结果：Fig. 3 图注称两头把 token 嵌到单纯形相反两极，注意力分别充当正、负分量，合起来给出 ζ^n=(1−3x)^n=(−1)^n(3x−1)^n。Fig. A8 写 x=0.5 时单头不能得到正确信念几何，两头则与 ground truth 相符。Fig. A7 图注称单层两头在所测参数上已能得到较低 KL；正文没有 KL 数字。
• 作者的解读：作者称振荡加上注意力非负，使第二头成为预期需要，而不是只凭“头会分工”的一般观察。单头表示不完整，预测也更差（附录 F）。

训练轨迹与多层网络

• 训练中的 MSE：Fig. A3 把 MSE 除以随机初始化的初始值，Y=1 为随机基线。作者称中间激活对约束几何、MLP 后激活对完整信念的归一化 MSE 都降到 1 以下。图注称 a=20 等条件可能先短暂上升再下降。论文未说明 a 与 α 的换算。
• 相对随机初始化：Fig. A4 图注称四个条件下，训练后的绝对 MSE 低于随机初始化。作者称收敛表示与理论预测一致，MSE 比随机初始化低若干个数量级。柱上的具体 MSE 正文没有写出。
• 4 层：Fig. A6 图注写条件为 x=15、50 与 a=20、60 的组合，正文未把这些记号对应到 α、x。图注称 L0 Mid 对约束几何的拟合好于对完整信念；第一个 MLP 之后改为更符合完整信念，约束几何不再同样变好，后期可能变差。训练后 MSE 比随机基线低若干个数量级。Fig. A5 把各层激活回归到完整信念后作图；图注称 Layer 0 的 mid 仍看得出约束几何的形状。

其他消融与分析

• Table 1 component 0 的累积方差，中间四列为 0.5408、0.4648、0.4074、0.5268，最终四列为 0.9618、0.4947、0.4596、0.6503；列顺序是 (x,α)=(0.15,0.2)、(0.15,0.6)、(0.5,0.6)、(0.5,0.2)。
• 作者称 A_{2,1} 大于其余聚在一起的 A_{m+1,m}，v_1 的模相应更小（第 4.4 节）。这里的 larger 不是统计检验。
• Fig. 4 每行换一组 Mess3 参数，四列依次为 Eq. (5)、中间 PCA、Eq. (2)、MLP 后 PCA。作者称理论与观察相符；α 离开 [0.2,0.6] 时偏差增大。
• 偏差如何随 α 缩放，第 4.3 节写留待以后。MLP 变形的完整数学性质，第 4.5 节也写留待以后。
• 单头在 ζ<0 时表示不完整（附录 F）。两头时 OV 方向相反，与 Eq. (18)–(20) 的正负分解一致（Fig. 3）。
• Fig. A3 的纵轴 inset 范围为 −0.02 到 0.5；正文没有逐 checkpoint 的 MSE。

作者指出的局限与后续方向

• 网络规模：实验用 1 层和 4 层的小 transformer，分析集中在第一层注意力（第 5 节）。
• 数据族：训练数据是对所有 token 序列具有全支撑的 Mess3 HMM（第 5 节）。
• 要改到的对象：作者称技术必须改到更大的 transformer，以及更能反映真实数据复杂性的生成过程（第 5 节）。
• 自然语言：作者称该设定捕捉不到自然语言的许多方面；后续可做层次化、序列支撑稀疏的过程（第 5 节）。
• 多层多头：作者称多头注意力与更深网络的相互作用，还有单层分析刚开始触及的细节（第 5 节）。
• 优化本身：作者称已看到最后的 MLP 把部分更新细化成近似完整贝叶斯，但梯度下降为何收敛到这些电路仍待研究（第 5 节）。

实验覆盖范围

• 网络：被分析的是作者训练的单层 transformer，以及 Figs. A5、A6 的 4 层 transformer。附录 B 只写出单层的 dmodel=64、dff=256；头数为 1 或 2。
• 数据与序列：生成过程是 Mess3。Table 1 与正文例子给出 x=0.15、x=0.5、α=0.2、α=0.6；Fig. 2 为 x=0.15、α=0.6。表示分析用长度不超过 10 的全部输入，无 BOS。
• 对照：几何对照是 Eq. (2)、Eq. (5) 和随机初始化的 MSE。Fig. A7 比较不同头数和层数的 KL，正文未给出数值。
• 谱预测的适用条件：第 4.4 节的推导针对第一层注意力。脚注写，第一层头很多时该对应的细节会失效。附录 D 记录前两个嵌入符号与预测相反。
• 论文未报告的协议：未报告训练种子或重复次数。Fig. A6 图注把条件写成 x=15、50 与 a=20、60，正文未说明这些记号与 α、x 的对应。

作者在 Mess3 上把 next-token transformer 的中间几何解释成受架构约束的贝叶斯更新。

• 问题：optimal prediction 要递归更新隐状态信念，注意力只能对局部 value 做并行加权和。作者要解释因此出现的分形，以及它如何变成最终信念几何。
• 做法：主分析是单层网络。约束更新从平稳分布 π 出发，把每个历史 token 的独立修正相加。特征值 ζ=1−3x 用来预测注意力随距离的衰减；ζ<0 时用两个头分担正负分量。MLP 再把中间表示变成完整信念。
• 几何结果：作者称 Figure 1 的最终 PCA 与信念几何相符，注意力后是另一种结构。Fig. 2 在 x=0.15、α=0.6 上，注意力大体只依位置衰减。作者称 α∈[0.2,0.6] 时中间表示与 Eq. (5) 相符，离开该区间偏差增大。
• 数字与头数：Table 1 中多数设置的前 3 个主成分累积解释方差超过 90%，但 x=0.5、α=0.6 的中间与最终 component 2 为 0.8913 和 0.8855。x>1/3 时单头不够，两头可拼出振荡。训练后 MSE 被作者称为比随机初始化低若干个数量级，具体柱值正文未写出。
• 作者的结论：作者称这把模型无关的预测理论与注意力的函数形式接在一起，并预期同一张力会出现在自然语言上的大模型中；同时写明当前证据来自小网络和全支撑的 Mess3。