注意力的梯度动力学：交叉熵如何塑造贝叶斯流形——贝叶斯注意力三部曲之Paper II

Gradient Dynamics of Attention: How Cross-Entropy SculptsBayesian Manifolds

注意力梯度动力学：交叉熵如何塑造贝叶斯流形

https://arxiv.org/pdf/2512.22473

Transformer 在精心构建的“贝叶斯风洞”和大规模语言模型中，经验上展现出精确的概率推理能力，但基于梯度的学习如何生成所需的内部几何结构仍不透明。我们对交叉熵训练如何重塑 Transformer 注意力头中的注意力得分和值向量进行了系统性的一阶分析。我们的核心结果是一个基于优势（advantage-based）的注意力得分路由法则。

其中， u i
是位置 i i 处的上游梯度， α i j
是注意力权重。这些方程引发了一个正反馈循环，在该循环中路由和内容协同专业化：查询更强烈地路由到那些对其误差信号高于平均值的值，而这些值又被拉向使用它们的查询。我们表明，这种耦合的专业化行为类似于一个双时间尺度的 EM 过程：注意力权重实现 E 步（软责任分配），而值实现 M 步（基于责任加权的原型更新），同时查询和键调整假设框架。通过受控模拟（包括一个粘性马尔可夫链任务，我们在其中将闭式 EM 风格更新与标准 SGD 进行比较），我们证明了最小化交叉熵的相同梯度动力学也会塑造我们在相关工作中识别出的低维流形，从而实现贝叶斯推断。此外，在 EM 与 SGD 的比较中，EM 收敛得更快，且两者的预测熵都接近最小贝叶斯熵。这产生了一幅统一的画面：优化（梯度流）催生几何结构（贝叶斯流形），而几何结构反过来又支持功能（上下文内概率推理）。

1 引言

Transformer 已成为序列建模的主导架构，但我们仍缺乏对梯度下降如何塑造其内部表征的机制性理解。最近的研究表明，在受控的“贝叶斯风洞”中，小型 Transformer 可以精确再现解析后验分布，其中键形成正交的假设轴，查询实现渐进式信念更新，值则沿着由后验熵参数化的一维流形展开 [1]。我们配套的缩放论文将这一图景扩展到生产规模模型，展示了 Pythia、Phi-2 和 LLaMA 中类似的几何特征。

这些发现提出了一个自然的问题：为什么单纯的交叉熵训练会产生贝叶斯推断所需的几何结构？理解这一点需要从静态几何转向梯度动力学：注意力分数、查询、键和值在优化过程中如何共同演化？

1.1 贡献

我们的主要贡献如下：

一个由责任加权的上游梯度平均值。这引发了一个正反馈循环：查询路由到有助于它们的值；这些值则向其使用者移动，从而强化路由并促成专业化。

(4) 双时间尺度 EM 解释。我们表明，这些动力学实现了一种隐式的类 EM 算法：注意力权重充当软责任分配（E 步），值作为在这些责任下更新的原型（M 步），而查询/键则作为潜在分配模型的参数。注意力通常较早稳定，而值则持续细化——这种框架-精度分离与我们在风洞实验和大型模型中的经验观察相符。

(5) 玩具实验及 EM 与 SGD 的比较。在合成任务中（包括一个粘性马尔可夫链序列），我们将由我们的梯度分析所诱导的闭式 EM 风格更新与标准 SGD 进行比较。EM 能显著更快地达到低损失、高准确率和尖锐的预测熵；SGD 收敛到类似解，但速度更慢且路由更分散。对值轨迹的 PCA 可视化揭示了新兴的低维流形。

结合 [1]，我们的结果提供了一个统一的叙述：

梯度下降 ⇒ 贝叶斯流形 ⇒ 上下文内推理。

关于“贝叶斯推断”的澄清。本文中，“贝叶斯推断”指的是对潜在任务变量的贝叶斯后验预测——而非对网络权重的后验。我们表明，交叉熵训练塑造了几何结构，该结构实现了对上下文内假设的贝叶斯滤波，而非贝叶斯权重不确定性。

2 设置与符号说明
我们分析一个作用于长度为 T 的序列的单头注意力机制。除非另有说明，下标 i , j , k 的取值范围均为 1 到 T 。

2.1向前传球

2.2 辅助量
为简洁起见，我们定义：

3 一阶梯度推导
我们现在不跳过任何步骤地推导所有相关梯度，重点关注能揭示其几何意义的形式。

3.1 输出梯度
对于每个 i i，关于 logits 的交叉熵梯度为：

3.2 关于值的梯度

3.3 关于注意力权重的梯度

定义一个优势量是很方便的，其符号的选择与梯度下降方向一致：

4 耦合动力学与专业化

我们现在展开第 3 节中梯度流的含义，重点关注路由（通过分数和注意力）与内容（通过值）之间的相互作用。

4.1 基于优势的注意力重分配

公式 (23) 表明，对于固定的查询 i i，定义一个符号与梯度下降方向一致的优势量是方便的：

4.2 值更新作为责任加权的原型
定义第 j j 列的注意力加权上游信号：

4.4 反馈循环与专业化

4.5 几何示意图

5 类 EM 的双时间尺度动力学

上述推导出的耦合动力学与期望最大化（EM）算法存在一种有用的类比，但并非对显式潜在变量似然函数的直接优化，而是梯度流与责任加权更新之间的机制性对应。注意力权重表现得如同对潜在源的责任分配，而值向量则作为在这些责任下更新的原型。与经典 EM 不同，此处的更新由上游梯度驱动，而非观测数据，并且没有独立优化关于值的似然函数。

5.1 注意力作为责任

5.3 近似 EM 与 SGD 的比较

在经典 EM 中，E 步和 M 步是分离的：首先固定参数重新计算责任，然后固定责任更新参数。而在使用 SGD 训练的 Transformer 中，这些步骤是交织且嘈杂的，但其一阶图景仍类似于 EM：

在第 7.2 节中，我们将此类 EM 风格的调度方案与标准 SGD 在粘性马尔可夫链任务上进行比较，发现两者均收敛至相似解，但 EM 风格的更新能更快地达到低损失和尖锐、聚焦的注意力。

5.5 贝叶斯视角与 EM 视角的对比

EM 是一种优化过程：它产生一个最大化（后验）似然的点估计 θ ∗ 。而完整的贝叶斯处理方式则应对 θ 进行积分，但这对于 Transformer 来说是不可行的。因此，我们的分析停留在 EM/SGD 层面。

然而，至关重要的是，我们的配套工作 [1] 表明，以这种方式学习到的点估计参数支持表征空间中的贝叶斯计算：值流形、键框架和查询轨迹在上下文中实现了贝叶斯信念更新。本文解释了为什么交叉熵和梯度下降会自然地形成这些结构。

6 从梯度流到贝叶斯流形

我们现在将上述推导出的梯度动力学与在贝叶斯风洞实验和生产模型中观察到的几何结构联系起来。

6.1 值流形展开

在风洞实验 [1] 中，我们观察到：

• 训练早期，注意力熵下降，注意力聚焦于相关假设。
• 训练后期，注意力模式看似稳定，但值表征沿一条平滑曲线展开；前几个主成分解释了大部分方差，且主轴与后验熵强相关。
• 即使注意力图在视觉上保持不变，校准误差仍持续下降。

逐渐地，会沿着残差误差景观的主方向对齐。在反复更新下，值向量最终落在由下游泛函（例如后验熵）参数化的低维流形上。

6.2 假设框架与键的正交性

如果不同的查询子集持续地发现不同的键更有帮助，相应的梯度贡献会将这些键在 k k-空间中彼此推开，从而促进不同假设轴之间的近似正交性。我们的“风洞”论文正好测量了这种正交性，并将其与竞争性假设的清晰分离联系起来。

6.3 框架-精度分离

经验观察到的“框架-精度分离”——即注意力稳定地定义一个假设框架，而校准仍在持续改进——现在很容易解释：

因此，一个训练后期的 Transformer 具有固定的贝叶斯框架（假设轴与路由），但其后验几何结构仍会继续变得更加锐利。

7 实验

我们现在通过受控模拟来说明该理论。所有实验均使用一个单头、单层注意力块，不包含残差连接或 LayerNorm，以保持动力学过程清晰透明。

7.1 玩具注意力模拟

观察结果： 在约 100 步内，我们观察到：

(1) 注意力热图逐渐变锐：每个查询的注意力质量集中于少数几个位置（见图 Figure 2、Figure 3）。

(2) 值向量在低维子空间中协同移动；它们在 PCA 投影下的轨迹显示出新兴的流形结构（见图 Figure 5）。

(3) 交叉熵损失平滑衰减（见图 Figure 4），大部分收益发生在专业化现象出现之时。

7.2 粘性马尔可夫链模拟：EM 与 SGD 对比

接下来，我们研究一个更具结构的任务，在该任务中注意力可以利用时间持久性：一个基于符号的粘性马尔可夫链。

要点： EM 风格和 SGD 训练最终都收敛到相似的定性解：专业化的值向量和聚焦的注意力。然而，EM 风格的调度方案以更少的步骤达到此状态，并实现更锐利的专业化。这符合双时间尺度的故事：责任分配（注意力）可被视为近似收敛，而闭式值更新则可利用这种稳定性加速流形的形成。

8 实用影响与诊断方法

梯度分析为训练和解释 Transformer 注意力机制提供了有用的诊断工具和设计原则。

8.1 诊断方法

8.2 正则化与稳定性

• 值上的 LayerNorm 可稳定范数，同时保持方向动力学不变。
• 注意力 Dropout 会扰乱反馈循环，限制过度专业化，并鼓励值被更均匀地使用。
• 学习率选择 调节路由与内容之间的时间尺度分离；较小的学习率使一阶图像更精确。

8.3 架构选择

• 多头注意力 允许多个专业化的路由流形共存，减少单个头内部的竞争。

• 深度 自然支持我们在风洞实验和大型模型中观察到的绑定–消除–精炼层级结构。
• 残差连接 有助于在各个头强烈专业化的同时，维持有用的中间表征。

9 相关工作

9.1 Transformer 的贝叶斯解释

一些研究认为，Transformer 通过行为方式或探针方式实现了近似的贝叶斯推断 [例如 7, 8]。我们的配套论文 [1] 在小型风洞中展示了精确的贝叶斯行为和几何特征，而一篇缩放论文则在生产级大语言模型中展示了类似的模式。本论文解释了梯度动力学如何产生这些几何结构。此外，本文还表明理论最小贝叶斯预测熵接近经验熵，这一点在配套论文中有更详细的探讨。

9.2 机制性可解释性

机制性可解释性研究旨在识别执行复制、归纳及其他算法任务的具体头和电路 [3, 5]。我们的框架通过解释专业化如何源于路由与内容的相互作用，补充了这一方向，而非将专业化头视为原始构件。

9.3 优化与隐式偏差

线性及深层网络中梯度下降的隐式偏差已被广泛研究 [2, 6]。我们将这些思想扩展至注意力机制：梯度下降隐式偏好那些路由与误差几何对齐、且值位于支持贝叶斯更新的低维流形上的表征。

此处推导的责任加权值更新让人联想到神经 EM 和槽注意力模型，其中软分配驱动原型更新。关键区别在于，在 Transformer 中，责任是通过内容可寻址的注意力计算的，而原型更新由反向传播的误差信号驱动，而非重建似然。我们的重点并非提出一种新的 EM 风格架构，而是展示标准交叉熵训练在注意力层中如何作为梯度流的结果，自然诱导出类 EM 的专业化动力学。

10 局限性与未来方向

我们的分析是有意保持简洁和受控的。

一阶近似。 我们在“一阶”框架下工作，假设学习率较小，并忽略高阶及随机效应（例如动量、Adam、小批量噪声）。将分析扩展至更现实的优化器是一个重要的下一步。

单头单层聚焦。 我们分析的是一个孤立的单头，不包含残差路径或 LayerNorm。多头、多层动力学——包括头间协调与层级专业化——仍是开放问题。

有限宽度 vs. 无限宽度。 我们并未明确将我们的分析与神经切向核或无限宽度极限联系起来。弥合这些机制可能有助于澄清 Transformer 何时在特征学习模式与懒惰训练模式下运行。

大规模经验验证。 我们的玩具模拟有意设计得规模很小。将第 8 节中的诊断工具应用于全规模大语言模型（LLM）的训练过程，跟踪优势矩阵和流形形成随时间的变化，是一个有前景的方向。

11 结论

本文聚焦于最小化设定下的“一阶”机制；配套工作已证实，相同的几何结构在更大规模上依然存在，并在受控风洞实验和大型模型中支持精确的贝叶斯推断。

我们的关键发现如下：

结合我们的风洞实验和缩放研究论文，这形成了一个连贯的三部曲：优化动力学构建贝叶斯几何，而该几何使 Transformer 能够作为上下文内的贝叶斯推理器运作。

原文：https://arxiv.org/pdf/2512.22473