# Bayes optimal learning of attention-indexed models

arXiv: https://arxiv.org/abs/2506.01582 · HTML: https://mlsanity.com/paper/2506.01582/

## 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 (most cited first, 30 of 39)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation](https://mlsanity.com/paper/2205.01445/) (2022)
- [The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks](https://mlsanity.com/paper/2202.08658/) (2022)
- [Statistical physics of inference: Thresholds and algorithms](https://mlsanity.com/paper/1511.02476/) (2015)
- [Infinite attention: NNGP and NTK for deep attention networks](https://mlsanity.com/paper/2006.10540/) (2020)
- [Online stochastic gradient descent on non-convex losses from high-dimensional inference](https://mlsanity.com/paper/2003.10409/) (2020)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [Dynamics of stochastic gradient descent for two-layer neural networks in the teacher-student setup](https://mlsanity.com/paper/1906.08632/) (2019)
- [Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer](https://mlsanity.com/paper/2305.16380/) (2023)
- [Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models](https://mlsanity.com/paper/2305.10633/) (2023)
- [The committee machine: Computational to statistical gaps in learning a two-layers neural network](https://mlsanity.com/paper/1806.05451/) (2018)
- [Learning time-scales in two-layers neural networks](https://mlsanity.com/paper/2303.00055/) (2023)
- [On Learning Gaussian Multi-index Models with Gradient Flow](https://mlsanity.com/paper/2310.19793/) (2023)
- [Asymptotic theory of in-context learning by linear attention](https://mlsanity.com/paper/2405.11751/) (2024)
- [Phase diagram of Stochastic Gradient Descent in high-dimensional two-layer neural networks](https://mlsanity.com/paper/2202.00293/) (2022)
- [A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks](https://mlsanity.com/paper/2310.07891/) (2023)
- [The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit](https://mlsanity.com/paper/2306.17759/) (2023)
- [From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks](https://mlsanity.com/paper/2302.05882/) (2023)
- [Bayes-optimal Learning of Deep Random Networks of Extensive-width](https://mlsanity.com/paper/2302.00375/) (2023)
- [What can a Single Attention Layer Learn? A Study Through the Random Features Lens](https://mlsanity.com/paper/2307.11353/) (2023)
- [Infinite Limits of Multi-head Transformer Dynamics](https://mlsanity.com/paper/2405.15712/) (2024)
- [Graph-based Approximate Message Passing Iterations](https://mlsanity.com/paper/2109.11905/) (2021)
- [A phase transition between positional and semantic learning in a solvable model of dot-product attention](https://mlsanity.com/paper/2402.03902/) (2024)
- [Hitting the High-Dimensional Notes: An ODE for SGD learning dynamics on GLMs and multi-index models](https://mlsanity.com/paper/2308.08977/) (2023)
- [Bayes-optimal learning of an extensive-width neural network from quadratically many samples](https://mlsanity.com/paper/2408.03733/) (2024)
- [Fundamental computational limits of weak learnability in high-dimensional multi-index models](https://mlsanity.com/paper/2405.15480/) (2024)
- [Attention layers provably solve single-location regression](https://mlsanity.com/paper/2410.01537/) (2024)

## Built on by (most cited first, 3 of 3)

- [Single-Head Attention in High Dimensions: A Theory of Generalization, Weights Spectra, and Scaling Laws](https://mlsanity.com/paper/2509.24914/) (2025)
- [High-Dimensional Theory of LoRA Fine-Tuning in a Solvable Attention Model](https://mlsanity.com/paper/2606.05899/) (2026)
- [Meta-Attention: Bayesian Per-Token Routing for Efficient Transformer Inference](https://mlsanity.com/paper/2605.28384/) (2026)

## 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.