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

TL;DR

The paper provides a complete first-order gradient analysis of a single-head transformer under cross-entropy, revealing an advantage-based routing law and responsibility-weighted value updates that jointly drive specialization. It presents an EM-like two-timescale interpretation in which attention acts as soft responsibilities (E-step) and value updates function as prototype updates (M-step), yielding a frame–precision dissociation that aligns with Bayesian manifolds observed in wind-tunnel and scaling studies. Through toy and Markov-chain experiments, the authors show that EM-like updates can converge faster than standard SGD while preserving final performance, and they connect these dynamics to the emergence of low-dimensional Bayesian manifolds in representation space. The work argues that optimization drives geometry, which in turn enables in-context probabilistic reasoning, offering practical diagnostics and architectural guidance for training and interpreting transformers. Together with companion papers, it posits a coherent picture where gradient flow constructs Bayesian geometry that underpins probabilistic inference in context.

Abstract

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).

Figures (9)

• Figure 1: Geometric illustration of coupled gradient dynamics. Value vector $v_j$ (blue) is updated toward the attention-weighted upstream signal $\bar{u}_j$ (green), inducing a change $\Delta g_i = \alpha_{ij}\Delta v_j$ (purple) in the context for any query $i$ attending to $j$. This change reduces loss to first order and increases compatibility $b_{ij}$, which in turn reinforces routing through $j$.
• Figure 2: Initial attention heatmap (toy simulation)
• Figure 3: Final attention heatmap, 100 steps (toy simulation)
• Figure 4: Loss, 100 EM steps (toy simulation)
• Figure 5: PCA projection of value vectors $v_j$ (toy simulation)