Specialization of softmax attention heads: insights from the high-dimensional single-location model

M. Sagitova; O. Duranthon; L. Zdeborová

cs.LGMar 4, 2026

Specialization of softmax attention heads: insights from the high-dimensional single-location model

Authors: M. Sagitova, O. Duranthon, L. Zdeborová

Organizations: Statistical physics of computation laboratory, École Polytechnique Fédérale de Lausanne, Switzerland

Abstract

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=10^{4}

(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

\sigma

softmax attention,

F=2

features,

\theta

drawn from the flipping spike distribution, with signal strengths

\nu_{1}=\nu_{2}=2

. Initialization

\eta=1

Panel of Figure 2 — Figure 2: Evolution of the $H$ heads at $F=2$ , softmax activation, $L=4$ , and $\eta=1$ . Left: for the flipping spike distribution, $\nu_{1}=\nu_{2}=2$ , and right for the non-isotropic Gaussian distribution, $\nu_{1}=8,\nu_{2}=2$ .

Panel of Figure 3 — Figure 3: Evolution of the $H=8$ heads for the non-isotropic Gaussian distribution at $F=3$ , according to Prop. III.2 . Left: $\sigma$ softmax; right: $\sigma$ B-softmax. $L=5$ , $\nu_{1}=20$ , $\nu_{2}=1$ and $\eta=1$ .

Figure 4: Predicted error

\mathcal{E}^{\infty}_{\sigma}

of the different activation functions.

L=5

\eta=1

. Left:

F=4

\nu=10

, varying

H

; right:

F=2

H=4

, varying signal strength. We performed 5 independent runs.

Panel of Figure 5 — Figure 5: MNIST detection task (described in App. C ). Left: learned $(k_{h})_{h\in[H]}$ at $H=10$ ; right: achieved error $\mathcal{E}^{\infty}_{\sigma}$ ; 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

\tilde{\tau}^{\mathrm{s}}

. Numerical simulations of SGD. We consider sequence length

L=5

H=2

heads with

\sigma

softmax attention,

F=2

features,

\theta

drawn from the flipping spike distribution, with signal strengths

\nu_{1}=\nu_{2}=2

. The threshold to determine the specialization is

\delta=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

\tilde{\tau}_{f}\tau^{\mathrm{s}}

. Numerical simulations of SGD on the MNIST detection task App. C .

H=10

heads with

\sigma

softmax attention. The threshold to determine the specialization time

\tau_{f}^{\mathrm{s}}

in each direction

f

||m_{:,f}-H^{-1}\sum_{h}^{H}m_{h,f}\mathds{1}_{H}||_{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,\nu_{1}=\nu_{2}=2,L=8

Figure 9: Asymptotic description of the attention trained by SGD. We compare numerical simulations at finite

D=10^{4}

(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

\theta

distributed according to the flipping spike distribution, with signal strengths

\nu_{1}=\nu_{2}=2

. Initialization

\eta=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=10^{4}

(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

\theta

distributed according to the non-isotropic Gaussian distribution, with signal strengths

\nu_{1}=8,\nu_{2}=2

. Initialization

\eta=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\in\{2,4,5,8\}

, and flipping spike distribution with

F=2

L=4

, with signal strengths

\nu_{1}=\nu_{2}=2

\eta=1

Figure 13: Evolution of the heads attention with varying number of heads $H\in\{2,5,8\}$ , and non-isotropic Gaussian distribution with $F=2$ , $L=4$ , with signal strengths $\nu_{1}=\nu_{2}=2$ (top) and $\nu_{1}=8,\nu_{2}=2$ (bottom). $\eta=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.

Softmax Attention

Softmax attention weights token representations by query–key similarity and underpins Transformer retrieval, sequence modeling, and long-context computation.

See all 75 papers

Incremental Learning of Sparse Attention Patterns in Transformers
Feb 22, 2026Oğuz Kaan Yüksel, Rodrigo Alvarez Lucendo, Nicolas FlammarionSelf-Attention Transformer Training Dynamics
TML Lab, EPFL
Multi-Head Self Attention is a Parameter Identification Mechanism
Sep 1, 2026W. Ross MorrowSoftmax Attention Parameter Identifiability
Gradient Flow Structure and Quantitative Dynamics of Multi-Head Self-Attention
May 5, 2026Ayan PendharkarSoftmax Attention Attention Head Analysis

Feb 22, 2026J/K move · Enter open · S save

Incremental Learning of Sparse Attention Patterns in Transformers

Oğuz Kaan Yüksel, Rodrigo Alvarez Lucendo, Nicolas Flammarion

TML Lab, EPFL

Self-Attention Transformer Training Dynamics Transformer Attention Sparse Attention Incremental Learning Attention Mechanisms+3

Open paper AskPDF arXiv

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.

Abstract

Softmax Attention

Explore similar work

Incremental Learning of Sparse Attention Patterns in Transformers

Multi-Head Self Attention is a Parameter Identification Mechanism

Gradient Flow Structure and Quantitative Dynamics of Multi-Head Self-Attention