Pith. sign in

REVIEW 3 major objections 5 minor 46 references

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

T0 review · 3 major / 5 minor · reviewed 2026-08-02 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗ Review JSON

Pith's one-line read Multi-head softmax attention specializes in two stages, learning signals strongest first.

desk verdict A solid statistical-physics analysis of softmax attention with a genuinely new Bayes-optimal normalization; the abstract oversells the sequential-specialization result, which is proven only at the saddle escape. read the letter →

arxiv 2603.03993 v2 pith:MAAZ4PBA submitted 2026-03-04 cs.LG cond-mat.dis-nn

M. Sagitova , O. Duranthon , L. Zdeborová This is my paper · ORCID
classification cs.LGcond-mat.dis-nn MSC 68T07
keywords multi-headattentionheadspecializationsoftmaxhigh-dimensionallimitorderparametersSGDdynamicsBayes-softmaxphasetransition
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper studies a stripped-down transformer task—one relevant token hidden among noise, with a signal built from a few latent directions—where attention itself is the only prediction mechanism. In the high-dimensional limit the full SGD dynamics collapses to gradient flow on a small set of order parameters, and the paper claims that training always passes through a fast unspecialized phase, in which all heads lock onto the mean signal, followed by a staged specialization phase in which heads sequentially align with the eigenvectors of the signal covariance, largest eigenvalue first. This provides a first-principles explanation for the empirically observed staged emergence and redundancy of attention heads. The paper then shows that standard softmax cannot switch off redundant heads and is generically suboptimal, and introduces Bayes-softmax attention, whose input-adaptive normalization reaches the Bayes risk and dictates the right number of heads. If correct, the framework turns questions about which patterns a transformer learns first, and how many heads are enough, into a spectral problem on the data distribution.

What carries the argument

The argument rests on two mechanisms. First, the order-parameter reduction: as D→∞, the loss factors through the alignments mhf = kh·k*_f and the residual covariance r, so SGD (with γNb = o(D^{-1})) becomes exact gradient flow on (m,r,b,v); this reduction is what makes the full training trajectory tractable. Second, the Hessian expansion around the unspecialized saddle (Lemma III.1): for m⊤Eθ=0 and small initialization, the loss is a quadratic form whose leading negative directions are the eigenvectors of Covθ, which explains why heads split in decreasing eigenvalue order; combining this with a diffusive rescaling of the magnetization yields the Θ(log D) saddle-escape time. The Bayes-softmax

What would settle it

Using the non-isotropic Gaussian prior with ν1>ν2 and Eθ=0, train the model at D=10^4 and record ‖m:,1‖ and ‖m:,2‖ over time. The theory predicts the f=1 direction (larger eigenvalue) is learned first and f=2 only after a delay of order log D; if the f=2 alignment grows first or both grow at the same rate, the Hessian-eigenvector ordering of specialization is wrong.

Watch

Extended reading notes

Core claim

The central claim is that, for D→∞, the population loss of a one-layer, multi-head softmax attention trained by SGD on the single-location task depends only on the head–feature alignments m and the residual overlap r, and the SGD trajectory converges to gradient flow on these order parameters. The dynamics then decompose into a fast unspecialized phase—where all heads share a common component aligned with the mean signal Eθ, requiring Θ(D) samples—and a slower specialization phase in which heads escape an unspecialized saddle along the eigenvectors of Covθ in descending eigenvalue order, requiring Θ(D log D) samples to begin. The paper further proves that standard softmax and softmax-v are n

Load-bearing premise

The specialization-phase proof assumes the mean signal and the initialization scale are small enough (independent of D) so the loss can be expanded around m≈0 and r≈0 and the rescaled magnetization follows a constant-diffusion SDE; if either is large, the predicted escape time, the sharp transition, and the Θ(D log D) sample complexity have no rigorous backing, only numerics.

Editorial extensions

If this is right

  • The order-parameter ODEs give an exact description of the SGD trajectory, making specialization times and head-acquisition order predictable from the data covariance before training.
  • The spectral ordering predicts that easy, high-variance patterns (large-covariance directions) are learned before rare or subtle ones, matching qualitative reports of bigram-then-induction emergence in transformers.
  • Standard softmax attention is provably suboptimal whenever the signal distribution has both θ and −θ in its support; deactivation-capable activations (softmax-1, Bayes-softmax) are needed to reach low error.
  • The Bayes-softmax result prescribes a concrete head-count: for a discrete latent prior, use one head per support point, with key vectors equal to the spike directions and biases set by the log-prior.
  • If Conjecture III.1 holds, the specialization transition is a sharp dynamical phase transition in the D→∞ limit, so sudden emergence of new attention behaviors is a concentration phenomenon rather than finite-size noise.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • The spectral mechanism is generic: it only needs anisotropy of the token-signal distribution, so the same two-phase, eigenvalue-ordered specialization should appear in real transformer layers whenever the query is input-independent; one can test this by comparing the order of head acquisition with the eigenvalue spectrum of the empirical covariance of token embeddings.
  • The Bayes-softmax estimator is effectively a nearest-prototype memory: each head stores a possible spike, and the bias stores its prior log-probability. This suggests that in trained transformers, attention heads with large biases are 'confident' prototypes, and pruning should preserve at least one head per latent mode, not merely heads with high norm.
  • The sample-complexity gap between the unspecialized phase (Θ(D)) and specialization (Θ(D log D)) implies that specialists are the bottleneck in data-limited regimes; scaling laws for transformer training might separate 'common-feature' learning from 'rare-feature' learning accordingly.
  • Because softmax learns mixture directions ±e1±e2... while Bayes-softmax learns pure directions, the choice of normalization changes the geometry of the learned representation; this is measurable in real models as the difference in cosine similarity between heads across softmax vs. softmax-1/B-softmax variants.
Share X Bluesky LinkedIn Reddit HN

Signed reviews

No signed human review yet.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 5 minor

Summary. The paper studies a minimal one-layer multi-head attention model for a synthetic single-location regression task in high dimension. It derives a low-dimensional order-parameter description of SGD dynamics (Prop. III.2), analyzes the unspecialized and specialization phases (Props. III.3–III.4), and proposes a sequential saddle-to-saddle specialization mechanism (Result III.1). In the second part it introduces a Bayes-softmax attention normalization and shows it reaches the Bayes risk for discrete signal distributions (Prop. IV.2), contrasting it with standard softmax and softmax-1. The theoretical findings are supplemented by finite-D simulations and MNIST-based experiments.

Significance. The paper contributes a solvable high-dimensional model in which attention training dynamics reduce to a closed ODE system on order parameters without fitted constants. The numerical match between the effective-dynamics ODEs and finite-D SGD (Fig. 1, 9–10) is convincing, and the B-softmax construction is an elegant way to connect attention normalization to the Bayes estimator. The paper also provides code. However, the central claim of sequential specialization (Result III.1) is only proven at early times and otherwise rests on heuristics and numerics, which substantially tempers the strength of contribution (ii). The paper is nonetheless valuable as a clearly formulated model with partial rigorous results and extensive numerical evidence.

major comments (3)
  1. [Section III, Result III.1 and contribution (ii)] The central claim — that training exhibits a hierarchy of specialization events with heads sequentially aligning with eigenvectors of Covθ — is stated in Result III.1 and in the abstract/contribution (ii) as established. However, the text immediately before Result III.1 says: 'The following discussion is based on an analytical result Lemma III.1 for the initial times of the specialization, and for the later times relies on heuristic arguments on the structure of the landscape and on a numerical integration of Prop. III.2.' Lemma III.1 is a local quadratic expansion near m≈0, r=0; Lemma A.11 extends the Hessian analysis only to small finite m and assumes Eθ=0, independent θ_f, and softmax activation. No proof is given for the global saddle-to-saddle trajectory or for the claim that softmax learns mixtures ±e1±...±eF while B-softmax learns single directions. This is a load-bearing gap for
  2. [Section III, Proposition III.4 and Lemma A.5] The specialization-phase sample-complexity statement (Prop. III.4 and Lemma A.5) is derived under the assumption that ||Eθ||² and η are small enough (independently of D) to expand the loss around m≈0 and r≈0 and to obtain a constant-diffusion SDE. The paper states that 'Prop. III.4 holds for larger values, as numerically shown by Figs. 1 to 3 and in Appendix E 1,' but no proof is provided outside this regime. Since the abstract and contribution (ii) describe a general two-stage dynamics without this qualification, the scope of the proven claim is narrower than the headline. The small-parameter condition should be stated in the main claims, or the proof extended.
  3. [Section IV, Proposition IV.2 and abstract] Contribution (iii) and the abstract state that B-softmax attention 'reaches the Bayes-risk and prescribes the optimal number of heads.' Proposition IV.2 establishes this only for distributions Pθ with discrete support {θ_h}_{h∈[H]} and with H equal to the support size. For continuous Pθ, the paper only argues approximation by discretization (Fig. 4, center-left). The exact Bayes-risk equality and the finite optimal-head-number prescription do not hold for continuous distributions. The abstract and contribution (iii) should be qualified accordingly.
minor comments (5)
  1. [Section IV, Prop. IV.2] The formula for b_h is malformed: 'bh = −||θh||2 2 logP θ(θh)/2'. It should read b_h = -||θ_h||^2/2 + log Pθ(θ_h) (or the intended expression), so that the B-softmax weights reproduce exp(-||θ||²/2) Pθ(dθ).
  2. [Section II.B] The notation ebh is introduced without definition; use e^{b_h} or define it explicitly.
  3. [Appendix E.2, Fig. 11] The figure caption says 'Cosine similarity by 2 (varying 1)' — the axes appear to be labeled with ν1, ν2, but the subscripts are missing in some labels.
  4. [Section III, Phase transition] The conjecture is introduced with 'we conjecture that ...' and then 'This conjecture relies on a heuristic argument' — this is fine, but the subsequent 'Moreover it extends to semi-realistic data' should clarify that the extension is numerical.
  5. [Section III, Sequential specialization] The phrase 'This is a key connection with what [3, 4] observe in practice' is vague; consider expanding on the specific empirical observations being referenced.

Circularity Check

0 steps flagged · score 1.0 of 10

No significant circularity: dynamics, sample complexity, and B-softmax optimality are proven in-appendix from the data model; self-citations are methodological sockets, and the partially heuristic later-stage specialization claim is a disclosed rigor gap, not a definitional reduction.

full rationale

Central derivations are self-contained. Props. III.1-III.4, Lemma III.1 and Lemmas A.1-A.11 (reparametrized loss, gamma-localizability of the order parameters, SDE limit, Hessian expansions, Bayes estimator) are proven in the appendix, so the dynamics and the Theta(D log D) weak-recovery complexity are derived from the data model rather than fitted to the simulations; the theory curves in Figs. 1/9/10 match finite-D SGD with no tuned constants, and the disclosed initialization of m and r to their empirical values at tau=0 is a state reset, not a dynamic fit. No load-bearing self-citation: the SGD-to-gradient-flow step invokes the external Theorem 2.3 of [17] atop the paper's own Lemma A.2, and self-references ([19],[21],[22],[31], shared authors) are methodological sockets ('as detailed e.g. in [31]', 'as suggested by [17]', data model 'inspired by [20,21]' with [20] external). Prop. IV.2 is transparently a construction: 'directly follows from a substitution of the given parameters into the B-softmax attention model leading to the expression of the Bayes estimator Prop. IV.1'; this is an expressivity-inclusion theorem against the derived Bayes risk, disclosed by the text ('as motivated later'), and the SGD ability to actually reach the plateau (Fig. 4) is a separate numerical claim. The flagged weakness is a rigor gap, not circularity: Result III.1's later-stage sequential hierarchy 'relies on heuristic arguments on the structure of the landscape and on a numerical integration of Prop. III.2', so the abstract's 'multi-stage specialization phase' exceeds the proof; the paper discloses this. The small-||Etheta||2/eta condition on Prop. III.4 is an explicit assumption, with numerics shown beyond it.

Assumptions & free parameters 0 free parameters · 4 assumptions · 1 invented entities

The central claims rely on the stated probabilistic data model and the high-dimensional limit; no continuous parameters are fitted to data to produce the theoretical predictions. Signal strengths, L, F, H, η, and γ are model inputs. The B-softmax is a new construction derived directly from the Bayes estimator, not a fitted entity. Some technical smallness conditions (small Eθ and η) are introduced specifically for the proofs and are the main price the reader pays.

assumptions (4)
  • standard math Random hidden spikes k*_f are i.i.d. N(0, D^{-1} I_D) and hence nearly orthogonal as D→∞; Gaussian projections χ*, ξ are independent.
    Used in Prop III.1 and Appendix A.1 to reduce the loss to order parameters (m, r, b, v).
  • standard math The SGD trajectory is γ-localizable and converges to population gradient flow in effective time τ = γt, per Theorem 2.3 of Ben Arous et al. [17].
    Basis for Prop III.2 / Lemma A.3; the proof verifies moment bounds and compact-cutoff conditions.
  • ad hoc to paper For Prop III.4 and Lemmas A.4–A.5, ||Eθ||^2 and initialization η are small enough that the loss can be expanded around m≈0 and r≈0, and the rescaled magnetization SDE has constant diffusion.
    Stated in Prop III.4 and Appendix A.2.c; beyond this regime the paper offers only numerics.
  • domain assumption For Prop IV.3, Pθ assigns bounded-away-from-zero mass to S and -S; for softmax-1 well-specifiedness, ||θ||^2 > B almost surely with B→∞.
    This is the symmetric-sign regime where softmax cannot deactivate a head; the contradiction proof relies on it.
invented entities (1)
  • Bayes-softmax (B-softmax) attention normalization
    purpose: Normalizes each head by the total logits over all heads so that unaligned heads are deactivated per input; proven to recover the Bayes estimator and attain Bayes risk when heads match the support of Pθ.
    Introduced in Sec IV (eq. (4),(13)); its optimality is verified by construction (Prop IV.2) and by in-paper simulations (Fig. 4, 5), but no external real-world benchmark provides an independent falsifiable handle.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Specialization of softmax attention heads: insights from the high-dimensional single-location model." pith.science (2026). https://pith.science/paper/MAAZ4PBA

@misc{pith2026260303993,
  author       = {Pith},
  title        = {Pith review of: Specialization of softmax attention heads: insights from the high-dimensional single-location model},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/MAAZ4PBA}},
  note         = {Machine review of arXiv:2603.03993}
}
read the original 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

Figures reproduced from arXiv: 2603.03993 by the authors.

Figure 1
Figure 1. FIG. 1: Asymptotic description of the attention trained by SGD. We compare numerical simulations at finite [PITH_FULL_IMAGE:figures/full_fig_p005_1.png] view at source ↗
Figure 2
Figure 2. FIG. 2: Evolution of the [PITH_FULL_IMAGE:figures/full_fig_p006_2.png] view at source ↗
Figure 3
Figure 3. FIG. 3: Evolution of the [PITH_FULL_IMAGE:figures/full_fig_p007_3.png] view at source ↗
Figures from the paper (12 more)
Figure 4
Figure 4. Figure 4: FIG. 4: Predicted error [PITH_FULL_IMAGE:figures/full_fig_p009_4.png]
Figure 5
Figure 5. Figure 5: FIG. 5: MNIST detection task (described in App. C). Left: learned ( [PITH_FULL_IMAGE:figures/full_fig_p009_5.png]
Figure 6
Figure 6. Figure 6: FIG. 6: Phase transition: concentration of the specialization time ˜τ [PITH_FULL_IMAGE:figures/full_fig_p027_6.png]
Figure 7
Figure 7. Figure 7: FIG. 7: Phase transition: concentration of the specialization time [PITH_FULL_IMAGE:figures/full_fig_p027_7.png]
Figure 8
Figure 8. Figure 8: FIG. 8: Histogram of the differences between magnetizations of the first head in the model trained with changed MC [PITH_FULL_IMAGE:figures/full_fig_p028_8.png]
Figure 9
Figure 9. Figure 9: FIG. 9: Asymptotic description of the attention trained by SGD. We compare numerical simulations at finite [PITH_FULL_IMAGE:figures/full_fig_p029_9.png]
Figure 10
Figure 10. Figure 10: FIG. 10: Asymptotic description of the attention trained by SGD. We compare numerical simulations at finite [PITH_FULL_IMAGE:figures/full_fig_p030_10.png]
Figure 11
Figure 11. Figure 11: FIG. 11: Cosine similarity between heads for softmax attention with [PITH_FULL_IMAGE:figures/full_fig_p031_11.png]
Figure 12
Figure 12. Figure 12: FIG. 12: Evolution of the heads in the model of attention with varying number of heads [PITH_FULL_IMAGE:figures/full_fig_p031_12.png]
Figure 13
Figure 13. Figure 13: FIG. 13: Evolution of the heads attention with varying number of heads [PITH_FULL_IMAGE:figures/full_fig_p032_13.png]
Figure 14
Figure 14. Figure 14: FIG. 14: Evolution of the [PITH_FULL_IMAGE:figures/full_fig_p033_14.png]
Figure 16
Figure 16. Figure 16: FIG. 16: Pruning head experiment. Error [PITH_FULL_IMAGE:figures/full_fig_p034_16.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

46 extracted references · 20 linked inside Pith

  1. [1]

    A. Chen, R. Shwartz-Ziv, K. Cho, M. L. Leavitt, and N. Saphra, Sudden drops in the loss: Syntax acquisition, phase transi- tions, and simplicity bias in mlms, inInternational Conference on Learning Representations (ICLR)(2024) arXiv:2309.07311

    arXiv 2024

  2. [2]

    switching off

    It corresponds to the loss of the Bayes estimator ˆyBayes, which is the optimal estimator of y in terms of population loss. We characterize it in the following proposition and derive it in Appendix A 3. Proposition IV.1(Bayes estimator).Given the spikes {k∗ f }f∈[F] ∈R F×D and a sequence X∈R L×D, the Bayes 8 estimator of the label is ˆyBayes(X, k∗) = LX ℓ...

  3. [3]

    Hoogland, G

    J. Hoogland, G. Wang, M. Farrugia-Roberts, L. Carroll, S. Wei, and D. Murfet, Loss landscape degeneracy and stagewise development in transformers, Transactions on Machine Learning Research (TMLR) (2024), arxiv:2402.02364

    arXiv 2024

  4. [4]

    Clark, U

    K. Clark, U. Khandelwal, O. Levy, and C. D. Manning, What does bert look at? an analysis of bert’s attention, arXiv preprint arXiv:1906.04341 (2019)

    arXiv 1906

  5. [5]

    Tigges, M

    C. Tigges, M. Hanna, Q. Yu, and S. Biderman, Llm circuit analyses are consistent across training and scale, inAdvances in Neural Information Processing Systems (NeurIPS)(2024) arXiv:2407.10827

    arXiv 2024

  6. [6]

    G. Wang, J. Hoogland, S. van Wingerden, Z. Furman, and D. Murfet, Differentiation and specialization of attention heads via the refined local learning coefficient, inInternational Conference on Learning Representations (ICLR)(2025) arXiv:2410.02984

    arXiv 2025

  7. [7]

    Voita, D

    E. Voita, D. Talbot, F. Moiseev, R. Sennrich, and I. Titov, Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned, inProceedings of the 57th Annual Meeting of the Association for Computational Linguistics(Association for Computational Linguistics, 2019) arXiv:1905.09418

    arXiv 2019

  8. [8]

    This sequential learning is shown for the non-isotropic Gaussian distribution, where Eθ = 0 and all the features f have different signal strengths in Figs

    for ICL of linear regression. This sequential learning is shown for the non-isotropic Gaussian distribution, where Eθ = 0 and all the features f have different signal strengths in Figs. 2b and 3, as well as in Appendix Figs. 10 and 13. In the case where two features have the same signal strength, they are learned at the same time, as shown in Fig. 13 top....

Show all 46 references
  1. [9]

    Michel, O

    P. Michel, O. Levy, and G. Neubig, Are sixteen heads really better than one?, inAdvances in Neural Information Processing Systems (NeurIPS), Vol. 32 (Curran Associates, Inc., 2019) arXiv:1905.10650

    2019 arXiv

  2. [10]

    Zhang, A

    Y. Zhang, A. K. Singh, P. E. Latham, and A. Saxe, Training dynamics of in-context learning in linear attention, in Proceedings of the 42th International Conference on Machine Learning (ICML)(2025) arXiv:2501.16265

    2025 arXiv

  3. [11]

    S. Chen, H. Sheen, T. Wang, and Z. Yang, Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality, inConference on Learning Theory (COLT)(2024) arXiv:2402.19442

    2024 arXiv

  4. [12]

    H. S. Seung, H. Sompolinsky, and N. Tishby, Statistical mechanics of learning from examples, Physical review A45, 6056 (1992)

    1992

  5. [13]

    Saad and S

    D. Saad and S. Solla, Dynamics of on-line gradient descent learning for multilayer neural networks, Advances in neural information processing systems8(1995)

    1995

  6. [14]

    D. Saad, S. A. Solla,et al., On-line learning in soft committee machines, Physical review E52, 4225 (1995)

    1995

  7. [15]

    Engel,Statistical mechanics of learning(Cambridge University Press, 2001)

    A. Engel,Statistical mechanics of learning(Cambridge University Press, 2001)

    2001

  8. [16]

    Aubin, A

    B. Aubin, A. Maillard, F. Krzakala, N. Macris, L. Zdeborov´ a,et al., The committee machine: Computational to statistical gaps in learning a two-layers neural network, Advances in Neural Information Processing Systems31(2018)

    2018

  9. [17]

    Goldt, M

    S. Goldt, M. Advani, A. M. Saxe, F. Krzakala, and L. Zdeborov´ a, Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup, Advances in neural information processing systems32(2019)

    2019

  10. [18]

    Ben Arous, R

    G. Ben Arous, R. Gheissari, and A. Jagannath, Online stochastic gradient descent on non-convex losses from high-dimensional inference, Journal of Machine Learning Research (JMLR)22, 1 (2021)

    2021

  11. [19]

    G. B. Arous, R. Gheissari, and A. Jagannath, High-dimensional limit theorems for sgd: Effective dynamics and critical scaling (2023), arXiv:2206.04030 [stat.ML]

    2023 arXiv

  12. [20]

    E. Abbe, E. B. Adser` a, and T. Misiakiewicz, Sgd learning on neural networks: leap complexity and saddle-to-saddle dynamics, inProceedings of the 36th International Conference on Machine Learning (ICML)(2023) arXiv:2302.11055. 11

    2023 arXiv

  13. [21]

    Arnaboldi, B

    L. Arnaboldi, B. Loureiro, L. Stephan, F. Krzakala, and L. Zdeborov´ a, Asymptotics of SGD in sequence-single index models and single-layer attention networks, inAdvances in Neural Information Processing Systems (NeurIPS)(2025) arXiv:2506.02651

    2025

  14. [22]

    Marion, R

    P. Marion, R. Berthier, G. Biau, and C. Boyer, Attention layers provably solve single-location regression, inThe Thirteenth International Conference on Learning Representations (ICLR)(2025) arXiv:2410.01537

    2025 arXiv

  15. [23]

    Troiani, H

    E. Troiani, H. Cui, Y. Dandi, F. Krzakala, and L. Zdeborov´ a, Fundamental limits of learning in sequence multi-index models and deep attention networks: High-dimensional asymptotics and sharp thresholds, Forty-second International Conference on Machine Learning (ICML) (2025)

    2025

  16. [24]

    Duranthon, P

    O. Duranthon, P. Marion, C. Boyer, B. Loureiro, and L. Zdeborov´ a, Statistical advantage of softmax attention: Insights from single-location regression, inThe Fourteenth International Conference on Learning Representations (ICLR)(2026) arXiv:2509.21936

    2026

  17. [25]

    Barnfield, H

    N. Barnfield, H. Cui, and Y. M. Lu, High-dimensional analysis of single-layer attention for sparse-token classification (2025), arXiv:2509.25153

    2025

  18. [26]

    Dohmatob, Understanding softmax attention layers: Exact mean-field analysis on a toy problem, inAdvances in Neural Information Processing Systems (NeurIPS)(2025)

    E. Dohmatob, Understanding softmax attention layers: Exact mean-field analysis on a toy problem, inAdvances in Neural Information Processing Systems (NeurIPS)(2025)

    2025

  19. [27]

    A. M. Saxe, J. L. McClelland, and S. Ganguli, A mathematical theory of semantic development in deep neural networks, PNAS116(2019), arxiv:1810.10531

    2019 arXiv

  20. [28]

    P. Kaul, C. Ma, I. Elezi, and J. Deng, From attention to activation: Unraveling the enigmas of large language models, in International Conference on Learning Representations (ICLR), Vol. 2025 (2025) pp. 34531–34582, arXiv:2410.17174

    2025 arXiv

  21. [29]

    Darcet, M

    T. Darcet, M. Oquab, J. Mairal, and P. Bojanowski, Vision transformers need registers, inInternational Conference on Learning Representations (ICLR)(2024) arXiv:2309.16588

    2024 arXiv

  22. [30]

    G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis, Efficient streaming language models with attention sinks, inInternational Conference on Learning Representations (ICLR)(2024) arXiv:2309.17453

    2024 arXiv

  23. [31]

    Z. Qiu, Z. Wang, B. Zheng, Z. Huang, K. Wen, S. Yang, R. Men, L. Yu, F. Huang, S. Huang, D. Liu, J. Zhou, and J. Lin, Gated attention for large language models: Non-linearity, sparsity, and attention-sink-free, inThe Thirty-ninth Annual Conference on Neural Information Process...

    2025 arXiv

  24. [32]

    Carion, F

    N. Carion, F. Massa, G. Synnaeve, N. Usunier, A. Kirillov, and S. Zagoruyko, End-to-end object detection with transformers, inComputer Vision – ECCV 2020(Springer International Publishing, 2020) pp. 213–229, arXiv:2005.12872

    2020 arXiv

  25. [33]

    Arnaboldi, Y

    L. Arnaboldi, Y. Dandi, F. Krzakala, B. Loureiro, L. Pesce, and L. Stephan, Online learning and information exponents: On the importance of batch size, and time/complexity tradeoffs, arXiv preprint arXiv:2406.02157 (2024)

    2024 arXiv

  26. [34]

    Pesme and N

    S. Pesme and N. Flammarion, Saddle-to-saddle dynamics in diagonal linear networks, inAdvances in Neural Information Processing Systems (NeurIPS)(2023) arXiv:2304.00488

    2023 arXiv

  27. [35]

    Deng, The mnist database of handwritten digit images for machine learning research, IEEE Signal Processing Magazine 29(2012)

    L. Deng, The mnist database of handwritten digit images for machine learning research, IEEE Signal Processing Magazine 29(2012)

    2012

  28. [36]

    C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. van Kerkwijk, M. Brett, A. Haldane, J. F. del R ´ ıo, M. Wiebe, P. Peterson, P. G´ erard-Marchant, K. Sheppar...

    2020

  29. [37]

    Paszke, S

    A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala, Pytorch: An imperative style, high-per...

    2019

  30. [38]

    J. D. Hunter, Matplotlib: A 2d graphics environment, Computing in Science & Engineering9(2007). Appendix A: Proofs of propositions and lemma. In this part we give the different proofs and justifications for our theoretical results of parts III and IV

    2007

  31. [39]

    Proof of Proposition III.1

    Low-dimensional characterization of SGD We start by proving the asymptotic characterization of the training dynamics in terms of the order parameters. Proof of Proposition III.1. We first derive the order parameters and the expression for the reparameterized loss. Notice that ...

  32. [40]

    2qY i Y µi # (A13) = 1 N 2q b NbX µ1 . . . NbX µk E

    Description of the dynamics a. Sample complexity of SGD In this section we formally prove the convergence of the trajectories of the order parameters under SGD to the GF in the limit and derive sample complexity of weak recovery with SGD using Theorem 2.3 from [17]. To do so, ...

  33. [41]

    Since we considerσsoftmax, to lighten the notation we writeσ(χ h) forσ(χ, b, v;h)

    We take m∈R H×F with m ¯fi = 0 for all ¯fi; it encompasses the case of a matrix where only the f -th column is not null equal tom f . Since we considerσsoftmax, to lighten the notation we writeσ(χ h) forσ(χ, b, v;h). ˜Eσ(m+ ¯m,0,0,0) =E ϵ,θ,χ∗   LX ℓ δℓ,ϵ − 1 H HX h σ(m⊤ h χ...

  34. [42]

    Proof of Proposition IV.1

    Expressiveness and performances of different activation functions In this part we derive the results about the Bayes risk and the expressivity of the softmax and softmax-1 activation functions. Proof of Proposition IV.1. The Bayes-optimal estimator is the posterior mean under ...

  35. [43]

    III.2, to the numerical simulation of SGD at finite large D

    Characterization of the training dynamics We compare our theoretical prediction, Prop. III.2, to the numerical simulation of SGD at finite large D. We consider the different activation functions, on the flipping spike or non-isotropic Gaussian. 100 101 102 103 τ 0.5 0.6 0.7 0....

  36. [44]

    The specialization of the two heads is measured as their cosine similarity and is reported in Fig

    Specialization of the heads We provide an additional figure illustrating the specialization of the heads for the flipping spike distribution, depending on the signal strength ν1 of the average direction and the signal strength ν2 of the flipping-sign direction. The specializat...

  37. [45]

    Sequential specialization We provide Figs. 13 and 14 that complement the description of the specialization dynamics of section III, showing the specialization of the heads for an anisotropic Gaussian distribution for different H or the different activation functions. The three...

  38. [46]

    irrelevant

    Comparison of the different activations after training In this section we provide further elements of comparison of the different activations after training, complementing the discussion of part IV. The attention maps of the different attentions are shown in Fig. 15 for a few ...

Pith tools

Reviewed August 2, 2026 · model on record in the stance chip above.