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

Naman Agarwal; Siddhartha R. Dalal; Vishal Misra

arxiv: 2512.22473 · v5 · pith:III6DKWDnew · submitted 2025-12-27 · 📊 stat.ML · cs.AI· cs.LG

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

Naman Agarwal , Siddhartha R. Dalal , Vishal Misra This is my paper

Pith reviewed 2026-05-21 15:59 UTC · model grok-4.3

classification 📊 stat.ML cs.AIcs.LG

keywords attention mechanismsgradient dynamicsBayesian manifoldscross-entropy lossEM algorithmtransformer optimizationin-context learning

0 comments

The pith

Cross-entropy gradients sculpt Bayesian manifolds in transformer attention heads through advantage-based routing and responsibility-weighted updates.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper shows that the first-order gradient of cross-entropy loss produces an advantage-based routing law for attention scores and a responsibility-weighted update rule for value vectors. Queries strengthen attention to values that exceed the average contribution to their error signal, while those values are drawn toward the queries that use them. This mutual specialization forms a positive feedback loop that mirrors a two-timescale EM procedure, with attention weights serving as soft responsibilities and values acting as updated prototypes. The process carves out low-dimensional manifolds that implement the Bayesian inference observed in controlled tasks and large language models.

Core claim

Our core result is an advantage-based routing law for attention scores, ∂L/∂s_ij = α_ij (b_ij − E_α_i [b]), coupled with a responsibility-weighted update for values, Δv_j = −η ∑_i α_ij u_i. These equations induce a positive feedback loop in which routing and content specialize together and behave like a two-timescale EM procedure that sculpts the low-dimensional manifolds implementing Bayesian inference.

What carries the argument

The advantage-based routing law for attention scores (∂L/∂s_ij equals attention weight times advantage of the dot product b_ij) together with responsibility-weighted value updates, which together drive coupled specialization of routing and content.

If this is right

The same gradient flow that reduces cross-entropy loss simultaneously builds the internal geometry needed for in-context probabilistic reasoning.
Attention weights perform an E-step via soft responsibilities while value vectors perform an M-step via responsibility-weighted prototype shifts.
Queries and keys jointly adjust the hypothesis frame within which the routing and updates occur.
In tasks such as the sticky Markov chain, the learned dynamics match those of a closed-form EM-style procedure.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

If the mechanism scales, it offers one route by which transformers acquire probabilistic reasoning without direct supervision on Bayesian tasks.
Attention variants could be designed that explicitly amplify the advantage signal to speed formation of the required manifolds.
Analogous first-order dynamics may operate in other transformer components such as feed-forward layers.

Load-bearing premise

The first-order gradient analysis and the identification of the resulting structures as the specific Bayesian manifolds both hold when moving from controlled simulations to the large-scale models referenced in the work.

What would settle it

Run the sticky Markov-chain simulation with the advantage term removed from the attention gradient while still minimizing cross-entropy; check whether the low-dimensional manifolds required for accurate Bayesian inference still appear.

Figures

Figures reproduced from arXiv: 2512.22473 by Naman Agarwal, Siddhartha R. Dalal, Vishal Misra.

**Figure 3.** Figure 3: Final attention heatmap, 100 steps (toy simu [PITH_FULL_IMAGE:figures/full_fig_p012_3.png] view at source ↗

**Figure 5.** Figure 5: PCA projection of value vectors 𝑣𝑗 (toy simulation) the previous symbol. We train a single-head attention layer with 𝑑𝑥 = 20, 𝑑𝑘 = 𝑑𝑞 = 10, 𝑑𝑣 = 15 for 1000 steps to predict 𝑦𝑡 from (𝑥1, . . . , 𝑥𝑡). Protocols. We compare two training schemes with causal masking and matched total parameter updates per step: (1) Standard SGD. Vanilla gradient descent on all parameters with learning rate 𝜂 = 0.01. (2) EM-lik… view at source ↗

**Figure 6.** Figure 6: Sticky Markov Chain: Loss Dynamics [PITH_FULL_IMAGE:figures/full_fig_p013_6.png] view at source ↗

**Figure 8.** Figure 8: Sticky Markov Chain: Prediction Entropy [PITH_FULL_IMAGE:figures/full_fig_p014_8.png] view at source ↗

read the original 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}} = \alpha_{ij}\bigl(b_{ij}-\mathbb{E}_{\alpha_i}[b]\bigr), \qquad b_{ij} := u_i^\top v_j, \] coupled with a \emph{responsibility-weighted update} for values, \[ \Delta v_j = -\eta\sum_i \alpha_{ij} u_i, \] where $u_i$ is the upstream gradient at position $i$ and $\alpha_{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).

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Desk Editor's Note private letter to a colleague

Derives an advantage-based routing law from cross-entropy gradients and frames the dynamics as two-timescale EM, but the link to specific Bayesian manifolds stays qualitative.

read the letter

The main point is that the authors work out how the cross-entropy gradient through softmax attention produces an explicit advantage-based routing rule for the scores together with a responsibility-weighted update for the values. This creates a feedback loop that they interpret as an EM-style process with separate timescales for routing and content updates. The derivation itself is the clearest part of the paper. Starting from the loss, they reach ∂L/∂s_ij = α_ij (b_ij − E_α[b]), where b_ij is the dot product of the upstream gradient and the value vector, and then show the value step as a weighted average pull. The sticky Markov-chain simulations then compare ordinary SGD trajectories against a closed-form EM update and find matching specialization patterns. That check is concrete and worth having on record. The EM analogy gives a useful way to think about why attention heads might separate hypothesis selection from prototype estimation during training. The paper stays focused on first-order local analysis, which keeps the claims tractable and avoids overclaiming global convergence. The soft spot is the geometry claim. The authors say the same dynamics sculpt the low-dimensional Bayesian manifolds from their companion work, yet the evidence is limited to qualitative observations of specialization and low-dimensional structure in the toy simulations. No direct metric appears that would confirm the attractors match the companion manifolds rather than some other low-dimensional set consistent with cross-entropy minimization. The extension to large language models is stated but not tested. Readers working on mechanistic interpretability or the optimization geometry of transformers will find the routing law and EM framing useful to build on. The derivations are explicit enough and the simulation is a reasonable first check, so the paper deserves a serious referee. Reviewers can ask for quantitative comparisons between the learned geometry and the claimed manifolds plus a clearer statement of what carries over from the toy task to real models. I would send it out for review after those additions.

Referee Report

2 major / 2 minor

Summary. The paper provides a first-order gradient analysis of cross-entropy loss in a transformer attention head, deriving an advantage-based routing law ∂L/∂s_ij = α_ij (b_ij − E_α_i [b]) with b_ij := u_i^T v_j and a responsibility-weighted value update Δv_j = −η ∑_i α_ij u_i. It argues that the resulting positive feedback loop between routing and content specialization behaves like a two-timescale EM procedure (attention as E-step, values as M-step) that sculpts the low-dimensional Bayesian manifolds identified in a companion work, thereby explaining how gradient training produces internal geometry for in-context probabilistic reasoning. This is illustrated via controlled simulations on a sticky Markov-chain task comparing the dynamics to a closed-form EM-style update.

Significance. If the central derivations and the quantitative link to the companion manifolds hold, the work offers a concrete mechanistic bridge between gradient flow and the emergence of Bayesian inference structures inside attention, with potential implications for understanding and designing probabilistic reasoning in transformers. The comparison of SGD dynamics to an explicit EM baseline in the simulations is a positive step toward reproducibility and falsifiability.

major comments (2)

[§4] §4 (sticky Markov-chain simulations): the central claim that the induced dynamics 'sculpt the low-dimensional manifolds identified in our companion work' rests on qualitative observation of specialization; no explicit quantitative metric (Hausdorff distance, effective dimension, or KL divergence to the companion posterior geometry) is reported to establish that the attractor is the same manifold rather than another low-dimensional structure consistent with cross-entropy minimization.
[§3] §3 (first-order analysis): while the local expressions for ∂L/∂s_ij and Δv_j follow directly from the softmax and cross-entropy assumptions, the manuscript does not supply the full intermediate derivation steps, error bounds, or verification that higher-order terms remain negligible under the controlled simulation conditions; this weakens the assertion that the dynamics generalize to the large-scale models referenced in the abstract.

minor comments (2)

Notation: the definition of the advantage term b_ij is introduced without an explicit statement of the upstream gradient u_i's dependence on the loss; a short clarifying sentence would improve readability.
Figure clarity: the simulation plots comparing SGD and EM trajectories would benefit from an additional panel showing the evolution of effective dimension or posterior alignment metric.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We thank the referee for the constructive report and the positive assessment of the work's significance. We respond point by point to the major comments and indicate the revisions we will implement.

read point-by-point responses

Referee: [§4] §4 (sticky Markov-chain simulations): the central claim that the induced dynamics 'sculpt the low-dimensional manifolds identified in our companion work' rests on qualitative observation of specialization; no explicit quantitative metric (Hausdorff distance, effective dimension, or KL divergence to the companion posterior geometry) is reported to establish that the attractor is the same manifold rather than another low-dimensional structure consistent with cross-entropy minimization.

Authors: We agree that the current presentation relies on qualitative observation and that this leaves the identification of the attractor open to alternative interpretations. In the revised manuscript we will add quantitative metrics in §4, including the Hausdorff distance between the simulated trajectories and the manifolds reported in the companion work, as well as the effective dimension of the learned attention and value representations. We will also report the KL divergence between the empirical distribution of attention weights under SGD and the posterior geometry obtained from the closed-form EM baseline. These additions will be accompanied by a brief discussion of any residual discrepancies. revision: yes
Referee: [§3] §3 (first-order analysis): while the local expressions for ∂L/∂s_ij and Δv_j follow directly from the softmax and cross-entropy assumptions, the manuscript does not supply the full intermediate derivation steps, error bounds, or verification that higher-order terms remain negligible under the controlled simulation conditions; this weakens the assertion that the dynamics generalize to the large-scale models referenced in the abstract.

Authors: We accept that the intermediate steps were omitted for conciseness. The revised version will include a dedicated appendix that presents the complete derivation of the advantage-based routing law and the responsibility-weighted value update, beginning from the cross-entropy loss and the softmax attention definition. We will also add a short verification subsection in the simulations that compares the first-order predictions against numerically computed full gradients under the same controlled conditions, thereby confirming that higher-order contributions remain small. Finally, we will revise the abstract and discussion to clarify that the analysis supplies a local mechanistic account whose direct extrapolation to very large models remains an empirical question for future work. revision: yes

Circularity Check

1 steps flagged

Manifold-sculpting claim depends on self-citation to companion work without independent geometric verification

specific steps

self citation load bearing [Abstract]
"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."

The demonstration that the derived dynamics sculpt the specific Bayesian manifolds identified in the companion work is justified solely by citation to that prior identification (by the same author group) rather than by an independent re-derivation or quantitative match performed inside the present manuscript.

full rationale

The first-order gradient derivations for the advantage-based routing law ∂L/∂s_ij = α_ij (b_ij − E_α_i [b]) and the responsibility-weighted value update Δv_j = −η ∑_i α_ij u_i follow directly from the stated cross-entropy loss and softmax attention definitions in §3, supplying independent analytic content. The paper further interprets the resulting positive-feedback specialization as a two-timescale EM procedure on the basis of those same equations. However, the central unification claim—that these dynamics sculpt the specific low-dimensional Bayesian manifolds—is supported only by qualitative simulation observations plus explicit reference to the identification performed in the companion paper by overlapping authors. No quantitative metric confirming geometric equivalence (e.g., Hausdorff distance, effective dimension, or posterior KL) is supplied here, so the most ambitious interpretation of the result reduces to that self-citation for its load-bearing step.

Axiom & Free-Parameter Ledger

0 free parameters · 1 axioms · 1 invented entities

Abstract-only access limits visibility into parameters and assumptions; entries reflect only statements explicitly present in the abstract.

axioms (1)

domain assumption The provided first-order analysis is complete for the attention head under cross-entropy.
Abstract states 'we provide a complete first-order analysis'.

invented entities (1)

Bayesian manifolds no independent evidence
purpose: Low-dimensional internal geometry implementing Bayesian inference
Referenced as identified in companion work; no independent evidence supplied in this abstract.

pith-pipeline@v0.9.0 · 5876 in / 1338 out tokens · 38218 ms · 2026-05-21T15:59:16.224011+00:00 · methodology

discussion (0)

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

IndisputableMonolith/Cost/FunctionalEquation.lean washburn_uniqueness_aczel unclear

?

unclear
Relation between the paper passage and the cited Recognition theorem.

advantage-based routing law for attention scores, ∂L/∂s_ij = α_ij (b_ij − E_α_i [b]), coupled with a responsibility-weighted update for values, Δv_j = −η ∑_i α_ij u_i. These equations induce a positive feedback loop... two-timescale EM procedure
IndisputableMonolith/Foundation/RealityFromDistinction.lean reality_from_one_distinction unclear

?

unclear
Relation between the paper passage and the cited Recognition theorem.

sculpt the low-dimensional manifolds identified in our companion work as implementing Bayesian inference

What do these tags mean?

matches: The paper's claim is directly supported by a theorem in the formal canon.
supports: The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends: The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses: The paper appears to rely on the theorem as machinery.
contradicts: The paper's claim conflicts with a theorem or certificate in the canon.
unclear: Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.

Forward citations

Cited by 2 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

The Bayesian Geometry of Transformer Attention
cs.LG 2025-12 unverdicted novelty 7.0

Small transformers reproduce known Bayesian posteriors with 10^{-3} to 10^{-4} bit accuracy in verifiable wind-tunnel tasks via residual belief states, FFN updates, and attention routing, while MLPs do not.
Geometric Scaling of Bayesian Inference in LLMs
cs.LG 2025-12 unverdicted novelty 5.0

Large language models preserve a geometric substrate in value representations that correlates with uncertainty and matches patterns from small models performing exact Bayesian inference.

Reference graph

Works this paper leans on

7 extracted references · 7 canonical work pages · cited by 2 Pith papers · 1 internal anchor

[1]

The Bayesian Geometry of Transformer Attention

Naman Agarwal, Siddhartha R. Dalal, and Vishal Misra. 2025. The Bayesian Geometry of Transformer Attention. arXiv:2512.22471 [cs.LG] https://arxiv.org/abs/2512.22471 Paper I of the Bayesian Attention Trilogy

work page internal anchor Pith review Pith/arXiv arXiv 2025
[2]

Jimmy Ba, Murat A Erdogdu, Marzyeh Ghassemi, Taiji Suzuki, Denny Wu, and Tianzong Zhang. 2022. High-dimensional asymptotics of feature learning: How one gradient step improves the representation.Advances in Neural Information Processing Systems35 (2022)

work page 2022
[3]

Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. 2...

work page 2021
[4]

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, et al. 2022. In-Context Learning and Induction Heads. Transformer Circuits Thread, Anthropic. https://transformer- circuits.pub/2022/in-contex...

work page 2022
[5]

Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. 2018. The implicit bias of gradient descent on separable data.Journal of Machine Learning Research19, 70 (2018), 1–57

work page 2018
[6]

Grewe, Bernhard Schölkopf, Claudia Clopath, and Johanni Brea

Johannes von Oswald, Christian Henning, Adrià Garriga-Alonso, Massimo Caccia, Frederik Träuble, Benjamin F. Grewe, Bernhard Schölkopf, Claudia Clopath, and Johanni Brea. 2023. Transformers as Meta-Learners for Bayesian Inference. arXiv preprint arXiv:2305.14034(2023). , Vol. 1, No. 1, Article . Publication date: January . 20 Naman Agarwal, Siddhartha R. D...

work page arXiv 2023
[7]

Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. 2022. An Explanation of In-Context Learning as Implicit Bayesian Inference. InInternational Conference on Learning Representations. , Vol. 1, No. 1, Article . Publication date: January

work page 2022

[1] [1]

The Bayesian Geometry of Transformer Attention

Naman Agarwal, Siddhartha R. Dalal, and Vishal Misra. 2025. The Bayesian Geometry of Transformer Attention. arXiv:2512.22471 [cs.LG] https://arxiv.org/abs/2512.22471 Paper I of the Bayesian Attention Trilogy

work page internal anchor Pith review Pith/arXiv arXiv 2025

[2] [2]

Jimmy Ba, Murat A Erdogdu, Marzyeh Ghassemi, Taiji Suzuki, Denny Wu, and Tianzong Zhang. 2022. High-dimensional asymptotics of feature learning: How one gradient step improves the representation.Advances in Neural Information Processing Systems35 (2022)

work page 2022

[3] [3]

Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. 2...

work page 2021

[4] [4]

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, et al. 2022. In-Context Learning and Induction Heads. Transformer Circuits Thread, Anthropic. https://transformer- circuits.pub/2022/in-contex...

work page 2022

[5] [5]

Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. 2018. The implicit bias of gradient descent on separable data.Journal of Machine Learning Research19, 70 (2018), 1–57

work page 2018

[6] [6]

Grewe, Bernhard Schölkopf, Claudia Clopath, and Johanni Brea

Johannes von Oswald, Christian Henning, Adrià Garriga-Alonso, Massimo Caccia, Frederik Träuble, Benjamin F. Grewe, Bernhard Schölkopf, Claudia Clopath, and Johanni Brea. 2023. Transformers as Meta-Learners for Bayesian Inference. arXiv preprint arXiv:2305.14034(2023). , Vol. 1, No. 1, Article . Publication date: January . 20 Naman Agarwal, Siddhartha R. D...

work page arXiv 2023

[7] [7]

Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. 2022. An Explanation of In-Context Learning as Implicit Bayesian Inference. InInternational Conference on Learning Representations. , Vol. 1, No. 1, Article . Publication date: January

work page 2022