Research map / Learning theory and algorithms / Kernels and regression theory
Bayes optimal learning of attention-indexed models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces the attention-indexed model (AIM), a solvable high-dimensional teacher-student setting for deep attention with extensive-width key/query matrices, deriving Bayes-optimal error predictions, an AMP algorithm matching them, and showing averaged gradient descent approaches the optimum.
Problem
Theoretical understanding of learning in attention layers with extensive-width key and query matrices, which prior sequence multi-index models could not analyze.
Approach
Define attention indices as bilinear forms x_a^T S x_b with rotationally invariant S of extensive rank, show deep attention reduces exactly to this model, and analyze it via replica theory and AMP.
Key results
- For T=2 hardmax, the Bayes-optimal estimation error decays as a power law in alpha with no finite strong-recovery threshold, while softmax reaches zero error at a finite alpha.
- Small-width hardmax weak-recovery threshold is approximately 0.563 in the rescaled sample complexity.
- AMP experiments at d=100 match the Bayes-optimal predictions, and averaged GD (AGD) approaches the Bayes-optimal error while plain GD is worse.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase sample complexity | Bayes-optimal estimation of hardmax tied attention · Hardmax tied attention, T=2 | estimation error: error approaches zero as α grows with rate compatible with O(1/α) | Helped |
| Use any positive sample complexity | Bayes-optimal estimation of hardmax tied attention · Hardmax tied attention, T=2 | estimation error: error smaller than 1 | Helped |
| Use softmax instead of hardmax | Tied attention · Single-layer attention | Bayes-optimal estimation error | Mixed |
| Vary inverse temperature β while keeping it finite | Bayes-optimal softmax tied attention · Softmax tied attention, T≥2 | estimation error | No clear effect |
| Increase token count T | Bayes-optimal softmax tied attention · Softmax tied attention, T≥2 | sample complexity for a fixed estimation error: factor 1 + 2/(T(T+1)) more samples | Hurt |
| Use AMP | Bayes-optimal estimation of hardmax tied attention · Hardmax tied attention, T=2 | estimation error | No clear effect |
| Use AMP | Bayes-optimal estimation of softmax tied attention · Softmax tied attention, T=2,3 | estimation error | No clear effect |
| Average gradient-descent estimates across initializations | Gradient descent for softmax tied attention · Softmax tied attention | estimation error | Helped |
| Use averaged gradient descent instead of standard gradient descent | Gradient descent for softmax tied attention · Softmax tied attention | estimation error | Helped |
| Increase sample complexity to the Bayes-optimal recovery threshold | Gradient descent and averaged gradient descent · Softmax tied attention | recovery: perfect recovery at the BO threshold | Helped |
Surprising findings
- The Bayes-optimal softmax error is independent of the inverse temperature beta, although the AMP output function depends on beta.
- The softmax strong-recovery threshold departs from the naive counting argument because of near-invertibility under row shifts.
Limitations
- Analysis is restricted to tied key and query (Q=K) and symmetric S; the untied case is left for future work.
- Replica results are heuristic and not rigorously proven; the multi-layer matrix denoiser is an open problem.
- Experiments are at modest size (d=100, 200) under a Gaussian data assumption.
Experimental setup
Datasets: synthetic Gaussian token embeddings with attention-indexed teacher outputs
Benchmarks: synthetic attention-indexed teacher-student tasks
Metrics: estimation error generalization error
Baselines: Adam gradient descent averaged gradient descent (AGD)
Builds on: sequence multi-index model multi-index models AMP for matrix denoising ellipsoid fitting
Closest papers by meaning
- Fundamental limits of learning in sequence multi-index models and deep attention networks: high-dimensional asymptotics and sharp thresholds (ICML 2025 · similarity 0.91)
- High-Dimensional Learning Dynamics of Attention-Indexed Models (NeurIPS 2026 · similarity 0.90)
- Dissecting the Interplay of Attention Paths in a Statistical Mechanics Theory of Transformers (NeurIPS 2024 · similarity 0.88)
- Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networks (NeurIPS 2025 · similarity 0.87)
- A Statistical Theory of Gated Attention through the Lens of Hierarchical Mixture of Experts (NeurIPS 2026 · similarity 0.87)
- In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention (ICML 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.