Research map / Architectures and efficiency / Attention and transformers
Attention layers provably solve single-location regression
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 single-location regression, shows a simplified single-head attention predictor (erf-activated) is asymptotically Bayes optimal while linear predictors fail, and proves projected gradient descent recovers the oracle parameters from initializations on an invariant manifold.
Problem
Theoretical understanding of why attention handles token-wise sparsity (only one token determines the output, at a latent random position) and internal linear representations, which prior theory has not addressed.
Approach
Defines a synthetic regression task where the output depends on one latent-position token, proposes an erf-based simplification of a single-head softmax attention layer as the predictor, and analyzes its Bayes optimality and projected gradient training dynamics on a reduced invariant manifold.
Key results
- Oracle erf-attention predictor risk tends to the Bayes risk epsilon^2 when lambda*sqrt(d) -> infinity and lambda*sqrt(L) -> 0.
- Optimal linear predictor risk tends to the null predictor risk gamma^2 as L -> infinity with max p_j -> 0 (uniform J_0).
- PGD on the invariant manifold converges to +/-(k*, v*) for generic initialization and small step size.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use the oracle attention-based predictor instead of the optimal linear predictor | single-location regression · single-location regression, asymptotic regime L = o(d) | mean squared error: linear risk tends to ε²+γ²; oracle risk tends to ε² | Helped |
| train the attention-based predictor with projected gradient descent from a generic initialization on the invariant manifold | single-location regression predictor · single-location regression, initialization on M | parameter recovery | Helped |
| use a decreasing inverse-temperature schedule instead of a constant inverse temperature | projected gradient descent initialized off M · single-location regression | parameter recovery | Helped |
| use a constant inverse temperature instead of a decreasing schedule | projected gradient descent initialized off M · single-location regression | parameter recovery | Hurt |
| use a full Transformer layer to solve the task | Transformer · single-location regression | task performance | Helped |
Surprising findings
- The optimal linear predictor performs no better than predicting zero when the relevant token position is uniform over many tokens.
- Iteration-independent inverse temperature fails to recover the oracle parameters when initialized off the invariant manifold, requiring a decreasing schedule.
Limitations
- Fixed sequence length, independent tokens, and output depending on only one token.
- Analysis is on the theoretical risk and invariant manifold; general initialization left open.
- Convergence proof relies on center-stable manifold theorem with no quantitative rates.
Experimental setup
Datasets: synthetic single-location regression data
Benchmarks: synthetic single-location regression
Metrics: mean squared error (risk) alignment |k^T k*| and |v^T v*|
Baselines: linear regression lasso group-lasso null predictor
Builds on: Transformer softmax attention riemannian gradient descent
Closest papers by meaning
- Sparse Attention as Compact Kernel Regression (NeurIPS 2026 · similarity 0.89)
- In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention (ICML 2025 · similarity 0.88)
- Neural Attention Search (NeurIPS 2025 · similarity 0.87)
- Ordinary Least Squares as an Attention Mechanism (NeurIPS 2026 · similarity 0.87)
- How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression (NeurIPS 2024 · similarity 0.87)
- In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness (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.