揭示了交叉熵训练如何通过梯度动态塑造注意力的贝叶斯几何结构。
Gradient Dynamics of Attention: How Cross-Entropy Sculpts Bayesian Manifolds
- 提出注意力分数的增量路由法则,实现误差信号驱动的自适应聚焦。
- 值向量按责任加权更新,形成与注意力耦合的特性专精机制。
- 该机制类似两时间尺度的EM算法,解释了上下文推理的几何基础。
Transformer在精心设计的‘贝叶斯风洞’和大规模语言模型中表现出精确的概率推理能力,但梯度学习如何构建所需的内部几何结构仍不清晰。本文对交叉熵训练下注意力头中的注意力分数与值向量的梯度动态进行了完整的首阶分析。核心结果是一个基于优势的路由法则:∂L/∂s_ij = α_ij(b_ij - E_{α_i}[b]),其中 b_ij = u_i^T v_j;以及一个责任加权的值更新规则:Δv_j = -η∑_i α_ij u_i,u_i为位置i的上游梯度,α_ij为注意力权重。这些方程引发正反馈循环:查询更强烈地路由至高于平均误差信号的值,而这些值则被拉向使用它们的查询。我们证明该耦合专精行为类似于两时间尺度的EM过程:注意力权重执行E步(软责任),值向量执行M步(责任加权原型更新),查询与键调整假设框架。通过控制模拟,包括粘滞马尔可夫链任务中将闭式EM更新与标准SGD对比,我们表明最小化交叉熵的同一梯度动态也塑造了先前研究中识别出的实现贝叶斯推理的低维流形。这提供了一个统一视角:优化(梯度流)产生几何(贝叶斯流形),后者又支持功能(上下文概率推理)。
原文摘要 · Abstract (English)
Transformers empirically perform precise probabilistic reasoning in carefully constructed ``Bayesian wind tunnels'' and in large-scale language models, yet the mechanisms by which gradient-based learning creates the required internal geometry remain opaque. We provide a complete first-order analysis of how cross-entropy training reshapes attention scores and value vectors in a transformer attention head. Our core result is an \emph{advantage-based routing law} for attention scores, \[ \frac{\partial L}{\partial s_{ij}} = α_{ij}\bigl(b_{ij}-\mathbb{E}_{α_i}[b]\bigr), \qquad b_{ij} := u_i^\top v_j, \] coupled with a \emph{responsibility-weighted update} for values, \[ Δv_j = -η\sum_i α_{ij} u_i, \] where $u_i$ is the upstream gradient at position $i$ and $α_{ij}$ are attention weights. These equations induce a positive feedback loop in which routing and content specialize together: queries route more strongly to values that are above-average for their error signal, and those values are pulled toward the queries that use them. We show that this coupled specialization behaves like a two-timescale EM procedure: attention weights implement an E-step (soft responsibilities), while values implement an M-step (responsibility-weighted prototype updates), with queries and keys adjusting the hypothesis frame. Through controlled simulations, including a sticky Markov-chain task where we compare a closed-form EM-style update to standard SGD, we demonstrate that the same gradient dynamics that minimize cross-entropy also sculpt the low-dimensional manifolds identified in our companion work as implementing Bayesian inference. This yields a unified picture in which optimization (gradient flow) gives rise to geometry (Bayesian manifolds), which in turn supports function (in-context probabilistic reasoning).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。