Bayes optimal learning of attention-indexed models

2025/06/02 by Fabrizio Boncoraglio, Boncoraglio, Fabrizio, Emanuele Troiani +5 · 1 citation
Computer Science · #Disordered Systems and Neural Networks (cond-mat.dis-nn) #FOS: Computer and information sciences #FOS: Physical sciences #Information Theory (cs.IT) #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Neural Networks and Applications

paper · pdf · doi:10.48550/arxiv.2506.01582

openalex publication_date 2025/06/02 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

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.

Citations

• On learning Gaussian multi‐index models with gradient flow part I: General properties and two‐timescale learning
• Fundamental Limits of Matrix Sensing: Exact Asymptotics, Universality, and Applications
• Fundamental limits of learning in sequence multi-index models and deep attention networks: High-dimensional asymptotics and sharp thresholds
• Unraveling the Gradient Descent Dynamics of Transformers
• Bilinear Sequence Regression: A Model for Learning from Long Sequences of High-dimensional Tokens
• Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis
• Attention layers provably solve single-location regression
• High-dimensional learning of narrow neural networks
• Local to Global: Learning Dynamics and Effect of Initialization for Transformers
• Fundamental computational limits of weak learnability in high-dimensional multi-index models
• Infinite Limits of Multi-head Transformer Dynamics
• Asymptotic theory of in-context learning by linear attention
• How Transformers Learn Causal Structure with Gradient Descent
• A phase transition between positional and semantic learning in a solvable model of dot-product attention
• On Learning Gaussian Multi-index Models with Gradient Flow
• A Theory of Non-Linear Feature Learning with One Gradient Step in Two-Layer Neural Networks
• Exact threshold for approximate ellipsoid fitting of random points
• Fitting an ellipsoid to random points: predictions using the replica method
• Hitting the High-Dimensional Notes: An ODE for SGD learning dynamics on GLMs and multi-index models
• What can a Single Attention Layer Learn? A Study Through the Random Features Lens
• The Shaped Transformer: Attention Models in the Infinite Depth-and-Width Limit
• Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer
• Smoothing the Landscape Boosts the Signal for SGD: Optimal Sample Complexity for Learning Single Index Models
• From high-dimensional & mean-field dynamics to dimensionless ODEs: A unifying approach to SGD in two-layers networks
• High-dimensional Asymptotics of Feature Learning: How One Gradient Step Improves the Representation
• The merged-staircase property: a necessary and nearly sufficient condition for SGD learning of sparse functions on two-layer neural networks
• Graph-based Approximate Message Passing Iterations
• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Infinite attention: NNGP and NTK for deep attention networks
• Language Models are Few-Shot Learners
• State Evolution for Approximate Message Passing with Non-Separable\n Functions
• Attention Is All You Need
• Statistical physics of inference: thresholds and algorithms
• Three unfinished works on the optimal storage capacity of networks
• Exact results and critical properties of the Ising model with competing interactions
• BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding

Cited by

• Inductive Bias and Spectral Properties of Single-Head Attention in High Dimensions
• Two failure modes of deep transformers and how to avoid them: a unified theory of signal propagation at initialisation

Related

• Fundamental limits of learning in sequence multi-index models and deep attention networks: High-dimensional asymptotics and sharp thresholds
• Theoretical Insights into Fine-Tuning Attention Mechanism: Generalization and Optimization
• Inductive Bias and Spectral Properties of Single-Head Attention in High Dimensions
• A phase transition between positional and semantic learning in a solvable model of dot-product attention
• Attention: Marginal Probability is All You Need?