Asymptotic Signal Subspace Recovery in Softmax Attention Models

How does the attention mechanism in a Transformer architecture identify relevant information within a sea of noisy tokens? While empirical success is well-documented, the underlying mathematical mechanism that allows attention to "discover" latent signals during training remains theoretically opaque. This paper investigates whether attention can be interpreted as a formal signal extraction procedure.

The author models the attention mechanism as a query vector being learned via stochastic gradient ascent. The input consists of a mixture of informative tokens (aligned with a hidden signal direction) and nuisance tokens (pure Gaussian noise). By exploiting the rotational symmetry of this model, the author reduces the complex high-dimensional optimization problem to a one-dimensional dynamical system governed by an alignment parameter. Using tools from stochastic approximation and dynamical systems theory, the paper establishes that the stochastic learning process tracks a deterministic ordinary differential equation (ODE).

The study demonstrates that the attention mechanism functions as a positive-feedback loop: as the query vector aligns with the latent signal, the attention weights for informative tokens increase, which in turn further reinforces the alignment. The main result proves that, under standard step-size conditions and high-dimensional scaling, the learned query vector converges almost surely to the signal subspace spanned by the latent direction. This confirms that attention-based optimization is a robust method for recovering latent structure in high-dimensional, noisy environments, effectively separating signal from noise up to an intrinsic sign ambiguity.

This work provides a rigorous theoretical foundation for understanding attention as a statistical signal-recovery tool rather than just an information-aggregation operator. By bridging Transformer theory with stochastic approximation and high-dimensional probability, the paper offers a clear, dynamical-systems perspective on how neural networks learn to focus on relevant data, providing a blueprint for future theoretical analyses of more complex attention-based architectures.