# A Mean-Field Analysis of Multi-Head Self-Attention under Cross-Entropy Training

arXiv: https://arxiv.org/abs/2606.10469 · HTML: https://mlsanity.com/paper/2606.10469/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 11 of its 17 references that are in the map.
- First posted on arXiv 9 Jun 2026.

## Builds on (most cited first, 11 of 11)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [Trainability and Accuracy of Neural Networks: An Interacting Particle System Approach](https://mlsanity.com/paper/1805.00915/) (2018)
- [A Theory of First Order Mean Field Type Control Problems and their Equations](https://mlsanity.com/paper/2305.11848/) (2023)

## Built on by (most cited first, 3 of 3)

- [A First-Order Mean Field Control Analysis of Transformer Layers under Cross-Entropy Training](https://mlsanity.com/paper/2606.23235/) (2026)
- [An Adjoint-Sensitivity Framework for Lost-in-the-Middle Phenomena in Causal Residual Transformers](https://mlsanity.com/paper/2607.17696/) (2026)
- [What Pretraining and Midtraining Make Learnable from Rewards?](https://mlsanity.com/paper/2609.38446/) (2026)

## Abstract (arXiv)

This paper develops a mean-field theory for a simplified single-layer causal multi-head self-attention model trained by cross-entropy minimization. Each attention head is treated as a particle in parameter space, and the empirical law of the heads is used as the large-head state variable. In the infinite-head limit, the averaged attention logits define a risk functional on probability measures, whose first variation generates a nonlinear Wasserstein gradient-flow equation. Unlike classical mean-field analyses of shallow networks that often focus on square-loss regression, the present model contains the softmax residual from the cross-entropy objective and the query-key-value structure of masked self-attention. We prove a static finite-head approximation bound for the optimal risk, characterize global minimizers through a variational support condition, and establish a quantitative finite-time propagation-of-chaos estimate comparing finite-head stochastic gradient descent with the limiting PDE. We then study the long-time behavior of the PDE: energy dissipation, convergence to the stationary set under compactness, convergence to a single stationary measure under topological or Kurdyka--Lojasiewicz assumptions, and explicit convergence rates under gradient-domination conditions. Finally, we prove local exponential stability under a Wasserstein strong-monotonicity condition and give verifiable stability and instability criteria for Dirac stationary measures. The results provide a rigorous baseline mean-field framework for attention-head training and clarify the additional compactness, landscape, and curvature assumptions needed to pass from stationarity to convergence and stability.