Asymptotic Signal Subspace Recovery in Softmax Attention Models
This work offers a foundational theoretical understanding of how attention mechanisms extract relevant signals from noisy high-dimensional data, addressing a key gap in the theory of attention models.
The paper provides a rigorous theoretical analysis showing that, under high-dimensional scaling assumptions, a softmax-attention model trained via stochastic gradient ascent almost surely converges to the one-dimensional signal subspace, recovering the latent informative direction up to sign ambiguity.
Attention mechanisms have demonstrated remarkable empirical success in identifying relevant information from large collections of tokens, yet the theoretical principles underlying this behavior remain poorly understood. We study a stylized softmax-attention model in which a query vector is learned by stochastic gradient ascent from a collection of informative and nuisance tokens. Exploiting the symmetry of the model, we derive a population objective and characterize the limiting ordinary differential equation governing the learning dynamics. Using tools from stochastic approximation and dynamical systems theory, we establish a rigorous connection between the stochastic learning algorithm and its deterministic limit. Our main result shows that, under suitable high-dimensional scaling assumptions and standard step-size conditions, the learned query converges almost surely to the one-dimensional signal subspace spanned by the latent informative direction. Equivalently, the query asymptotically recovers the latent signal up to the intrinsic sign ambiguity. These results provide a rigorous theoretical foundation for understanding attention mechanisms as signal extraction procedures in high-dimensional noisy environments and offer a dynamical-systems perspective on how attention discovers relevant information in the presence of substantial noise.