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 .