High-Dimensional Learning Dynamics of Attention-Indexed Models
Abstract
Attention mechanisms are central to modern foundation models, yet their training dynamics remain poorly understood, especially when the attention matrices have extensive rank. In this work, we study attention-indexed models, a broad framework that can represent multi-layer and multi-head attention architectures. First, we show that, in a suitable high-dimensional limit, the population-loss landscape is characterized by a finite set of trace order parameters. In contrast, online stochastic gradient descent (SGD) is governed by an infinite hierarchy of matrix moments, which we show can be exponentially well-approximated by a finite truncated system. Second, this framework reveals that attention parameterization itself can act as an architectural implicit bias. Direct optimization of an attention matrix S∈ℝd× d can remain trapped in an uninformative state. Tied attention (S=WW⊤) induces an automatic symmetry-breaking mechanism and yields weak recovery in Θ(d2log d) samples. For untied attention, S=UV⊤, we uncover a fast-slow mechanism: the pre-activation mean first evolves on a fast timescale, while the overlaps evolve on a slower one. Weak recovery on the Θ(d2log d) scale occurs when the state selected by the fast dynamics breaks the initial symmetry.
Related
- Learning Decision-Stump Thresholds in Context: Dynamics of Softmax Attention
- Langevin dynamics for high-dimensional optimization: the case of multi-spiked tensor PCA
- Tensor-Train Weak SINDy: Identifying High-Dimensional Nonlinear Dynamics
- Characterizing High-dimensional Dynamics by Combinatorial-Topological Methods on a Latent Space
- Energetic Variational Deep Learning for Solving High-Dimensional Stationary Fokker-Planck Equations
BibTeX
@misc{xu2026highdimensional,
title = {{High-Dimensional Learning Dynamics of Attention-Indexed Models}},
author = {Yizhou Xu and Margarita Sagitova and Lenka Zdeborová and Florent Krzakala},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.03858},
eprint = {2609.03858},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.03858},
}
RIS
TY - UNPB
TI - High-Dimensional Learning Dynamics of Attention-Indexed Models
AU - Yizhou Xu
AU - Margarita Sagitova
AU - Lenka Zdeborová
AU - Florent Krzakala
PY - 2026
DA - 2026/09/03
PB - arXiv
DO - 10.48550/arxiv.2609.03858
UR - https://arxiv.org/abs/2609.03858
AB - Attention mechanisms are central to modern foundation models, yet their training dynamics remain poorly understood, especially when the attention matrices have extensive rank. In this work, we study attention-indexed models, a broad framework that can represent multi-layer and multi-head attention architectures. First, we show that, in a suitable high-dimensional limit, the population-loss landscape is characterized by a finite set of trace order parameters. In contrast, online stochastic gradient descent (SGD) is governed by an infinite hierarchy of matrix moments, which we show can be exponentially well-approximated by a finite truncated system. Second, this framework reveals that attention parameterization itself can act as an architectural implicit bias. Direct optimization of an attention matrix S∈ℝ^d× d can remain trapped in an uninformative state. Tied attention (S=WW^⊤) induces an automatic symmetry-breaking mechanism and yields weak recovery in Θ(d^2log d) samples. For untied attention, S=UV^⊤, we uncover a fast-slow mechanism: the pre-activation mean first evolves on a fast timescale, while the overlaps evolve on a slower one. Weak recovery on the Θ(d^2log d) scale occurs when the state selected by the fast dynamics breaks the initial symmetry.
ER -