Search papers, labs, and topics across Lattice.
This paper investigates the training dynamics of attention-indexed models, revealing that the population-loss landscape can be characterized by a finite set of trace order parameters in a high-dimensional limit. The authors demonstrate that online stochastic gradient descent (SGD) operates under an infinite hierarchy of matrix moments, which can be approximated by a finite truncated system, highlighting the role of attention parameterization as an implicit architectural bias. Notably, the study uncovers mechanisms of symmetry breaking in tied attention and a fast-slow dynamic in untied attention, both of which influence the model's recovery capabilities based on sample size.
Attention parameterization can trap models in uninformative states, but symmetry-breaking mechanisms can dramatically enhance recovery efficiency.
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\in\mathbb{R}^{d\times d}$ can remain trapped in an uninformative state. Tied attention ($S=WW^\top$) induces an automatic symmetry-breaking mechanism and yields weak recovery in $\Theta(d^2\log d)$ samples. For untied attention, $S=UV^\top$, 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^2\log d)$ scale occurs when the state selected by the fast dynamics breaks the initial symmetry.
