mlsanityData to 14 Oct 2026

Bayes optimal learning of attention-indexed models

arXiv 2506.01582 ↗First posted 2 Jun 2025cs.LG cond-mat.dis-nn cs.ITcited by 3

Not on the map. Why?

Abstract (arXiv)

We introduce the attention-indexed model (AIM), a theoretical framework for analyzing learning in deep attention layers. Inspired by multi-index models, AIM captures how token-level outputs emerge from layered bilinear interactions over high-dimensional embeddings. Unlike prior tractable attention models, AIM allows full-width key and query matrices, aligning more closely with practical transformers. Using tools from statistical mechanics and random matrix theory, we derive closed-form predictions for Bayes-optimal generalization error and identify sharp phase transitions as a function of sample complexity, model width, and sequence length. We propose a matching approximate message passing algorithm and show that gradient descent can reach optimal performance. AIM offers a solvable playground for understanding learning in self-attention layers, that are key components of modern architectures.

Key facts. Cited by 3 papers in the map (data through 14 Oct 2026). Builds on 39 of its 54 references that are in the map. First posted on arXiv 2 Jun 2025.

Builds on

39 of 54 references.

References · 39

  1. 2017Attention Is All You Needcited by 60,361
  2. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,399
  3. 2020Language Models are Few-Shot Learnerscited by 32,419
  4. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 30,186
  5. 2022High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representationcited by 150
  6. 2022The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networkscited by 142
  7. 2015Statistical physics of inference: Thresholds and algorithmscited by 128
  8. 2020Infinite attention: NNGP and NTK for deep attention networkscited by 121
  9. 2020Online stochastic gradient descent on non-convex losses from high-dimensional inferencecited by 119
  10. 2024How Transformers Learn Causal Structure with Gradient Descentcited by 118
  11. 2019Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setupcited by 112
  12. 2023Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformercited by 102
27 more
  1. 2023Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Modelscited by 61
  2. 2018The committee machine: Computational to statistical gaps in learning a two-layers neural networkcited by 58
  3. 2023Learning time-scales in two-layers neural networkscited by 57
  4. 2023On Learning Gaussian Multi-index Models with Gradient Flowcited by 57
  5. 2024Asymptotic theory of in-context learning by linear attentioncited by 55
  6. 2022Phase diagram of Stochastic Gradient Descent in high-dimensional two-layer neural networkscited by 50
  7. 2023A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networkscited by 49
  8. 2023The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limitcited by 47
  9. 2023From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networkscited by 43
  10. 2023Bayes-optimal Learning of Deep Random Networks of Extensive-widthcited by 41
  11. 2023What can a Single Attention Layer Learn? A Study Through the Random Features Lenscited by 40
  12. 2024Infinite Limits of Multi-head Transformer Dynamicscited by 39
  13. 2021Graph-based Approximate Message Passing Iterationscited by 35
  14. 2024A phase transition between positional and semantic learning in a solvable model of dot-product attentioncited by 30
  15. 2023Hitting the High-Dimensional Notes: An ODE for SGD learning dynamics on GLMs and multi-index modelscited by 26
  16. 2024Bayes-optimal learning of an extensive-width neural network from quadratically many samplescited by 23
  17. 2024Fundamental computational limits of weak learnability in high-dimensional multi-index modelscited by 19
  18. 2024Attention layers provably solve single-location regressioncited by 17
  19. 2024High-dimensional learning of narrow neural networkscited by 12
  20. 2025Fundamental limits of learning in sequence multi-index models and deep attention networks: High-dimensional asymptotics and sharp thresholdscited by 12
  21. 2024Local to Global: Learning Dynamics and Effect of Initialization for Transformerscited by 11
  22. 2024Unraveling the Gradient Descent Dynamics of Transformerscited by 11
  23. 2025Fundamental Limits of Matrix Sensing: Exact Asymptotics, Universality, and Applicationscited by 11
  24. 2024Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysiscited by 9
  25. 2023Fitting an ellipsoid to random points: predictions using the replica methodcited by 7
  26. 2023Exact threshold for approximate ellipsoid fitting of random pointscited by 7
  27. 2024Bilinear Sequence Regression: A Model for Learning from Long Sequences of High-dimensional Tokenscited by 3

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2025: 1; 2026: 2.

3 papers.

  1. 2025Single-Head Attention in High Dimensions: A Theory of Generalization, Weights Spectra, and Scaling Lawscited by 6
  2. 2026High-Dimensional Theory of LoRA Fine-Tuning in a Solvable Attention Modelcited by 3
  3. 2026Meta-Attention: Bayesian Per-Token Routing for Efficient Transformer Inferencecited by 0