We study a random model of deep multi-head self-attention in which the weights are resampled independently across layers and heads, as at initialization of training. Viewing depth as a time variable, the residual stream defines a discrete-time interacting particle system on the unit sphere. We prove that, under suitable joint scalings of the depth, the residual step size, and the number of heads, this dynamics admits a nontrivial homogenized limit. Depending on the scaling, the limit is either deterministic or stochastic with common noise; in the mean-field regime, the latter leads to a stochastic nonlinear Fokker--Planck equation for the conditional law of a representative token. In the Gaussian setting, the limiting drift vanishes, making the homogenized dynamics explicit enough to study representation collapse. This yields quantitative trade-offs between dimension, context length, and temperature, and identifies regimes in which clustering can be mitigated.
Figures & tables
Figure 1. Qualitative phase diagrams in the (η,α) -plane for fixed depth L (left), together with the same picture in logarithmic coordinates (right), where straight lines correspond to power-law scalings of η and α with respect to L . Different regions correspond to the rigorous regimes obtained from Theorem 1 . Specifically, “ODE I” (Corollary 2 ) yields ( 4.1 ) in the ballistic regime αηL=o(1) and η2L=o(1) ; “ODE II” (Corollary 3 ) yields ( 4.2 ) in the refined deterministic regime αηL=o(1) and η3L=o(1) ; and “SDE” (Corollary 4 ) yields the homogenized SDE ( 3.5 ) in the diffusive regime αηL=O(1) and η3L=o(1) . The region labelled “Static” corresponds to ηL=o(1) , so that no non-trivial macroscopic evolution is seen on the time scale tL=ηL , whereas “Failure of approximation” indicates scalings not covered by the present weak-approximation argument. We study the qualitative behavior of the homogenized model in the setting of centered Gaussian weights, for which only the diffusive regime yields non-stationary evolution—see ( 2.5 ).
Figure 2. Comparison between the curve γ in Theorem 2 ( left ) and that of [ 29 , Theorem 2.13] for d=128 ( right ).
Figure 3. Theorem 3 asserts that the second moment approaches the solution of the logistic equation ( 2.10 ) when d is large and \upbeta is small. Left: n=200 and \upbeta=0.05 . Right: n=5000 and \upbeta=0.001 .
Figure 4. Realizations of solutions to ( 2.14 ) in Theorem 4 .
Transformer self-attention can be interpreted as a gradient flow on the unit sphere, in which tokens evolve under softmax interaction potentials and tend to form clusters. While prior work has established clustering behavior for single-head attention, the multi-head setting remains less understood due to geometric interference between heads, which invalidates standard monotonicity arguments. In this work, we develop a theoretical framework for multi-head self-attention dynamics and resolve several open questions. We show that, under suitable conditions on the score matrices, a natural multi-head energy functional is non-decreasing along both flat and spherical dynamics. We identify the key obstruction to per-head monotonicity as radial shadow terms, which are projections of each head's output onto token directions, persisting even under orthogonality assumptions. We introduce a sufficient condition ensuring monotonicity and establish robustness to approximate orthogonality. In a simplified scalar-head regime with equiangular token configurations, we derive a closed-form expression for the critical inverse temperature governing clustering behavior, and show that heterogeneous heads exhibit super-additive clustering rates. In this regime, we also prove a separation in clustering time between ReLU and softmax attention in the linearized dynamics. Finally, we establish an entropy production identity and show that attention entropy increases monotonically toward equilibrium as clustering progresses. Our results provide a unified perspective on the dynamics of multi-head attention and clarify the mechanisms underlying clustering and stability in transformer models.
We prove pathwise convergence of the layerwise evolution of tokens in a finite-depth, finite-width transformer model with MultiLayer Perceptron (MLP) blocks to a continuous-time stochastic interacting particle system. We also identify the stochastic partial differential equation describing the evolution of the tokens' distribution in this limit and prove propagation of chaos when the number of such tokens is large. The bounds we establish are quantitative and the limits we consider commute. We further prove that the limiting stochastic model displays synchronization by noise and establish exponential dissipation of the interaction energy on average, provided that the common noise is sufficiently coercive relative to the deterministic self-attention drift. We finally characterize the activation functions satisfying the former condition.
Andrea Agazzi, Giuseppe Bruno, Eloy Mosig García +2
Department of Mathematics and Statistics, University of Bern (CH) · Department of Mathematics, University of Pisa (IT)
A theoretical understanding of multi-head self-attention is fundamental to the study of modern neural networks. Using random matrix theory, we establish Gaussian equivalence for multi-head self-attention: replacing softmax attention with rescaled scores plus Gaussian noise preserves the limiting spectral law of the centered output. This equivalence also covers value and output projections that depend on the keys. The resulting laws separate the effects of head allocation and projection widths, and distinguish spectrum-preserving across-head sharing from within-head key--value dependence.
Tomohiro Hayase, Ryo Karakida
Artificial Intelligence Research Center (AIRC), AIST · RIKEN AIP