Research map / Learning theory and algorithms / Sample complexity
Fundamental limits of learning in sequence multi-index models and deep attention networks: high-dimensional asymptotics and sharp thresholds
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.
Maps deep tied low-rank attention networks to sequence multi-index (SMI) models and derives Bayes-optimal and AMP/GAMP weak-recovery thresholds, showing layers are learned sequentially as sample complexity grows.
Problem
Understanding the statistical and computational limits of learning deep self-attention networks with tied low-rank weights, where existing theory covers only single attention layers or shallow multi-index models.
Approach
Shows that deep attention networks are sequence multi-index (SMI) functions of low-rank projections of the input sequence; flattening tokens turns SMI models into multi-index models with block-structured weights, enabling adaptation of Bayes-optimal and GAMP analyses (state evolution, weak-recovery thresholds) to attention models.
Key results
- For a two-layer attention model (P1=P2=1, c=1, M=2) the second layer becomes learnable at alpha≈0.14 and the first layer at alpha≈0.79, showing sequential layer-wise learning.
- GAMP achieves the same final error when initialized with strong side information, suggesting no computational-to-statistical gap in this setting.
- On the TREC task with a two-layer attention network, layer-wise sequential learning is observed, in reversed order relative to the synthetic setting.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase sample complexity past α1 ≈ 0.14 | GAMP on two-layer attention · Two-layer attention, layer 2 weak recovery | layer 2 weight overlap: α1 ≈ 0.14 | Helped |
| Increase sample complexity past α2 ≈ 0.79, after layer 2 is recovered | GAMP on two-layer attention · Two-layer attention, layer 1 weak recovery | layer 1 weight overlap: α2 ≈ 0.79 | Helped |
| Increase sample complexity beyond α1 ≈ 0.14 | GAMP on two-layer attention · Two-layer attention | prediction error: α1 ≈ 0.14 | Helped |
| Increase sample complexity beyond α2 ≈ 0.79, enabling recovery of the first layer | GAMP on two-layer attention · Two-layer attention | prediction error: α2 ≈ 0.79 | Helped |
| Fix first-layer weights to zero (ablation) | GAMP on two-layer attention · Two-layer attention | prediction error | Hurt |
| Initialize GAMP with strong side information | GAMP on two-layer attention · Two-layer attention | prediction error | No clear effect |
| Increase depth from two to three layers | GAMP on deep attention · Depth-3 attention | layer-wise weak recovery threshold | Mixed |
| Run more GAMP iterations at fixed sample complexity | GAMP on two-layer attention · Two-layer attention | layer weight cosine similarity: approximately 10 steps | Helped |
| Train for additional epochs after the second layer begins learning | SGD on two-layer attention · Two-layer attention | layer weight reconstruction: after one epoch | Helped |
| Train the practical transformer on TREC | Two-layer transformer with AdamW · TREC | layer-weight similarity: 40 epochs | Mixed |
Surprising findings
- The order in which layers are learned on TREC is reversed relative to the synthetic two-layer teacher setting.
Limitations
- SMI functions do not account for fully connected layers interleaved with attention.
- Analysis restricted to Gaussian i.i.d. inputs.
- Results are asymptotic with finite sequence length M, width P and depth L.
Experimental setup
Datasets: TREC
Benchmarks: TREC
Metrics: overlap prediction error cosine similarity
Baselines: SGD
Builds on: multi-index models generalized approximate message passing adaptive interpolation grand-staircase mechanism
Closest papers by meaning
- Bayes optimal learning of attention-indexed models (NeurIPS 2025 · similarity 0.91)
- Asymptotics of SGD in Sequence-Single Index Models and Single-Layer Attention Networks (NeurIPS 2025 · similarity 0.89)
- Degrees of Freedom for Linear Attention: Distilling Softmax Attention with Optimal Feature Efficiency (NeurIPS 2025 · similarity 0.88)
- Hierarchical Self-Attention: Generalizing Neural Attention Mechanics to Multi-Scale Problems (NeurIPS 2025 · similarity 0.86)
- High-Dimensional Learning Dynamics of Attention-Indexed Models (NeurIPS 2026 · similarity 0.86)
- Understanding the Differences in Foundation Models: Attention, State Space Models, and Recurrent Neural Networks (NeurIPS 2024 · similarity 0.86)
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.