High-Dimensional Learning Dynamics of Attention-Indexed Models

arXiv:2609.03858v1 Announce Type: new 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 SinR^{dimes d} can remain trapped in an uninformative state. Tied attention (S=WW^op) induces an automatic symmetry-breaking mechanism and yields weak recovery in Theta(d^2log d) samples. For untied attention, S=UV^op, 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 Theta(d^2log d) scale occurs when the state selected by the fast dynamics breaks the initial symmetry.

Related

• Reproducible macroscopic dynamics in a closed-loop human-AI learning system
• Spectral phase transitions and trainability in neural network learning dynamics
• When Both Layers Learn: Training Dynamics of Representing Linear Models via ReLU Networks

Source: arXiv cs.LG | 2026-09-04