Multi-head attention enables transformer models to represent multiple attention patterns simultaneously. Empirically, head specialization emerges in distinct stages during training, while many heads remain redundant and learn similar representations. We propose a theoretical model capturing this phenomenon, based on the multi-index and single-location regression frameworks. In the first part, we analyze the training dynamics of multi-head softmax attention under SGD, revealing an initial unspecialized phase followed by a multi-stage specialization phase in which different heads sequentially align with latent signal directions. In the second part, we study the impact of attention activation functions on performance. We introduce the Bayes-softmax attention, which achieves optimal prediction performance in this setting.
Figures & tables
Figure 1: Asymptotic description of the attention trained by SGD. We compare numerical simulations at finite D=104 (dots) and the theoretical description stated in proposition III.2 (continuous lines). Left: loss; center left: alignment along f=1 the constant direction; center right: alignment along f=2 the “flipping sign” direction; right: orthogonal component. We consider sequence length L=10 , H=2 heads with σ softmax attention, F=2 features, θ drawn from the flipping spike distribution, with signal strengths ν1=ν2=2 . Initialization η=1 .
Figure 2: Evolution of the H heads at F=2 , softmax activation, L=4 , and η=1 . Left: for the flipping spike distribution, ν1=ν2=2 , and right for the non-isotropic Gaussian distribution, ν1=8,ν2=2 .
Figure 3: Evolution of the H=8 heads for the non-isotropic Gaussian distribution at F=3 , according to Prop. III.2 . Left: σ softmax; right: σ B-softmax. L=5 , ν1=20 , ν2=1 and η=1 .
Figure 4: Predicted error Eσ∞ of the different activation functions. L=5 , η=1 . Left: F=4 , ν=10 , varying H ; right: F=2 , H=4 , varying signal strength. We performed 5 independent runs.
Figure 5: MNIST detection task (described in App. C ). Left: learned (kh)h∈[H] at H=10 ; right: achieved error Eσ∞ ; for the different activation functions after training.
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Phase transition: concentration of the specialization time τ~s . Numerical simulations of SGD. We consider sequence length L=5 , H=2 heads with σ softmax attention, F=2 features, θ drawn from the flipping spike distribution, with signal strengths ν1=ν2=2 . The threshold to determine the specialization is δ=0.2 . The means and variances are empirically computed over 100 independent runs for each point.
Figure 7: Phase transition: concentration of the specialization time τ~fτs . Numerical simulations of SGD on the MNIST detection task App. C . H=10 heads with σ softmax attention. The threshold to determine the specialization time τfs in each direction f is ∣∣m:,f−H−1∑hHmh,f\mathds1H∣∣2>1 . The means and variances are empirically computed over 100 independent runs for each point.
Figure 8: Histogram of the differences between magnetizations of the first head in the model trained with changed MC samples and the same MC samples at each step of the training. The data are sampled from the flipping spike distribution, H=2,F=2,ν1=ν2=2,L=8 .
Figure 9: Asymptotic description of the attention trained by SGD. We compare numerical simulations at finite D=104 (dots) and the theoretical description stated in proposition III.2 (continuous lines). We consider sequence length L=5 , H=2 heads, F=2 features and θ distributed according to the flipping spike distribution, with signal strengths ν1=ν2=2 . Initialization η=1 . Top: softmax; middle: softmax-1; bottom: B-softmax.
Figure 10: Asymptotic description of the attention trained by SGD. We compare numerical simulations at finite D=104 (dots) and the theoretical description stated in proposition III.2 (continuous lines). We consider sequence length L=5 , H=4 heads, F=2 features and θ distributed according to the non-isotropic Gaussian distribution, with signal strengths ν1=8,ν2=2 . Initialization η=1 . Top: softmax; middle: softmax-1; bottom: B-softmax.
Figure 11: Cosine similarity between heads for softmax attention with H=2 , and flipping spike distribution with F=2 , L=4 depending on the signal strengths for the constant and flipping-sign directions.
Figure 12: Evolution of the heads in the model of attention with varying number of heads H∈{2,4,5,8} , and flipping spike distribution with F=2 , L=4 , with signal strengths ν1=ν2=2 . η=1 .
Figure 13: Evolution of the heads attention with varying number of heads H∈{2,5,8} , and non-isotropic Gaussian distribution with F=2 , L=4 , with signal strengths ν1=ν2=2 (top) and ν1=8,ν2=2 (bottom). η=1 .
Figure 14: Evolution of the H=8 heads for the non-isotropic Gaussian distribution at F=3 . The 2 lines are 2 different runs with different initial conditions.
This paper studies simple transformers trained on a high-order Markov chain, where the model must incorporate information from multiple past positions, each with different statistical importance. We show that transformers learn the task incrementally, with each stage corresponding to learning how to copy information from a subset of positions via a sparse attention pattern. Notably, the learning dynamics transition from a competitive phase, where all heads focus on the statistically most important positions, to a cooperative phase, where different heads specialize in different patterns. We model these dynamics with simplified differential equations and prove stage-wise convergence of the resulting system. Functionally, these stages correspond to a sequence of increasingly expressive misspecified models, with the full model class reached only at the end. Overall, we give a theoretical account of how structured attention patterns and head specialization emerge in stages without an explicit curriculum, with implications for generalization in sequential tasks.
Oğuz Kaan Yüksel, Rodrigo Alvarez Lucendo, Nicolas Flammarion
We prove that a multi-head scaled dot product attention can be viewed as a parameter identification strategy. The ratio of unidentified parameters to the total number of parameters scales like the reciprocal of the number of heads (1/2→1/(2H)), meaning models with more heads are structurally more identified. A subtle side effect of the mathematics observation that attention can never be fully identified. Similarly we also show that some bias terms can have no effect on softmax-based attention layers in both the single- and multiple-head settings, though this is mostly a curiosity that should have a marginal effect on model size and model training/prediction efficiency. We also touch on modern improvements to transformers including RoPE and GQA from this perspective, illustrating how those as well can improve the ratio of meaningful'' parameters to all parameters. Simple numerical examples demonstrate that training can indeed involve updates that overlap model-invariant subspaces that arise from a lack of identification. As part of our experiments we use a rebalancing'' approach that can ``fix'' updates that overlap unindentified subspaces but do not try to present evidence this should actually be adopted. Instead we simply view our numerical results as exploring and confirming the theoretical results. As a whole we discuss a purely mathematical/statistical explanation, identification, for why specific architectural choices in transformers may have improved performance.
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.