Gradient Dynamics of Attention: How Cross-Entropy Sculpts Bayesian Manifolds

Transformer 模型在精心构建的“贝叶斯风洞”以及大规模语言模型中，能够实证地执行精确的概率推理，然而基于梯度的学习机制如何形成所需内部几何结构仍不明确。我们对交叉熵训练如何重塑 Transformer 注意力头中的注意力得分和值向量进行了完整的一阶分析。我们的核心结果是一个关于注意力得分的“优势导向路由定律”： \[ \frac{\partial L}{\partial s_{ij}} = α_{ij}\bigl(b_{ij}-\mathbb{E}_{α_i}[b]\bigr), \qquad b_{ij} := u_i^\top v_j, \] 以及一个关于值向量的“责任加权更新”规则： \[ Δv_j = -η\sum_i α_{ij} u_i, \] 其中 \(u_i\) 是位置 \(i\) 处的上游梯度，\(α_{ij}\) 是注意力权重。这些方程共同引发了一个正反馈循环，使得路由机制与内容表征协同专业化：查询会更强地将注意力路由到那些对其误差信号而言高于平均水平的值，而这些值则被使用它们的查询拉向自身方向。我们证明，这种耦合的专业化过程类似于一种双时间尺度的 EM（期望最大化）算法：注意力权重实现 E 步骤（即软性责任分配），而值向量实现 M 步骤（即基于责任加权的原型更新），同时查询和键则调整假设空间的框架。通过受控模拟实验——包括一个粘性马尔可夫链任务，在其中我们将闭式解形式的类 EM 更新与标准 SGD 进行比较——我们展示了最小化交叉熵的相同梯度动态，也会塑造出我们在配套研究中识别出的低维流形，而这些流形正是实现贝叶斯推断的结构基础。这提供了一个统一的图景：优化过程（梯度流）催生了几何结构（贝叶斯流形），而该几何结构反过来支撑了功能实现（上下文内的概率推理）。

论文试图揭示Transformer模型在通过梯度学习进行交叉熵训练时，如何在其内部构建实现精确概率推理的几何结构。具体而言，尽管Transformer在‘贝叶斯风洞’和大规模语言模型中表现出强大的概率推理能力，但其注意力机制和值向量是如何被优化过程塑造以支持这种功能的，仍然不明确。这是一个相对较新的问题，聚焦于深度学习模型中优化动态与内部表示几何之间的根本联系。

论文提出了一个一阶分析框架，揭示了注意力得分遵循‘优势路由定律’（advantage-based routing law），而值向量则通过‘责任加权更新’进行调整。这两个机制形成正反馈循环：查询将更强地路由到对其误差信号而言表现高于平均水平的值，同时这些值又被拉向使用它们的查询。该过程类似于两时间尺度的EM算法——注意力权重执行E步（软分配责任），值向量执行M步（基于责任更新原型），从而在优化过程中自动生成支持贝叶斯推理的低维流形结构。这一思路新颖之处在于首次将标准SGD训练动态直接解释为实现统计推断的隐式EM过程。

作者通过受控模拟验证理论，包括设计一个粘性马尔可夫链任务，并比较闭式EM风格更新与标准SGD的表现，证明两者动态一致。实验展示了梯度下降如何自发形成用于上下文内概率推理的低维表示流形。虽然未明确提及开源代码，但工作为理解Transformer内部机制提供了可建模、可仿真的分析工具。值得深入的方向包括扩展至多头注意力、探索不同损失函数下的几何演化，以及将此EM类比应用于模型编辑或可解释性研究。

1. “In-context Learning and Induction Heads as Emerging Statistical Estimators” 2. “Emergent Properties of the Local Geometry in Transformer Language Models” 3. “The Role of Attention in Bayesian Inference with Transformers” 4. “Gradient-Induced Manifolds in Deep Sequence Models” 5. “From Optimization to Inference: Geometric Perspectives on Transformers”