Research map / Architectures and efficiency / Attention and transformers
In-context Convergence of Transformers
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.
Analyzes GD training dynamics of a one-layer softmax-attention transformer for in-context linear regression, proving convergence for balanced features (two phases) and stage-wise convergence for imbalanced features (four phases for under-represented features).
Problem
Understanding how softmax-attention transformers trained by gradient descent learn in-context, since prior theory of in-context learning dynamics only covered linear transformers; the setting includes balanced and imbalanced feature distributions.
Approach
Tracks bilinear attention weights (query-target feature weight A_k vs. off-target weights B_{k,n}) under GD to show finite-time convergence of softmax attention for in-context linear regression, and decomposes imbalanced-feature training into stage-wise phases determined by which weight dominates.
Key results
- Balanced features: GD reaches loss within epsilon of the minimum in T* = O(log(K)K^2/eta + K log(K eps^-1/2)/(eps eta)) iterations, with attention concentrating on query-matching tokens.
- Imbalanced features: the dominant feature converges in O(log(eps^-1/2)/(eta eps)) iterations, while under-represented features need a much longer run (up to O(log(K)K^2/eta + ...)) through four phases.
- The analysis shows stage-wise convergence: near-zero error for dominant-feature queries first, then for under-represented features later.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train the one-layer softmax transformer by GD on balanced-feature prompts | One-layer softmax transformer · Balanced features | population loss suboptimality: L(\theta^{(T^*)})-L^* \leq \epsilon | Helped |
| Train the one-layer softmax transformer by GD on balanced-feature prompts | One-layer softmax transformer · Balanced features | target-feature attention error: (1-\Attn_k^{(T^*)})^2 \leq O(\epsilon) | Helped |
| Train the one-layer softmax transformer by GD on imbalanced-feature prompts | One-layer softmax transformer · Imbalanced features, dominant feature v_1 | conditional prediction error: \cL_1(\theta^{(T_1)})\leq \cL_{1}^{*}+\epsilon | Helped |
| Train the one-layer softmax transformer by GD on imbalanced-feature prompts | One-layer softmax transformer · Imbalanced features, under-represented features | conditional prediction error: \cL_k(\theta^{(T_k)})\leq \cL_{k}^{*}+\epsilon | Helped |
| Train on imbalanced-feature prompts, including under-represented features | One-layer softmax transformer · Imbalanced features, under-represented features | target-feature attention error: (1-\Attn^{(T_{k})}_k)^2 \leq O(\epsilon) | Helped |
| Use imbalanced rather than balanced feature frequencies | One-layer softmax transformer trained by GD · Imbalanced features, dominant versus under-represented features | GD iterations to convergence: T_1=O(\frac{\log(\epsilon^{-\frac{1}{2}})}{\eta\epsilon}); T_k=O(\frac{\log(K)K^2}{\eta}+\frac{K \log\left(K\epsilon^{-\ | Mixed |
Surprising findings
- In-context learning remains accurate for under-represented query features despite their rare occurrence in the prompt, reaching the same near-zero error as dominant features, only more slowly.
Limitations
- Analysis restricted to a one-layer transformer with reparameterized weights (fixed value matrix nu=1, structured W^KQ), not full softmax parameterization.
- Orthonormal feature vectors and Gaussian-like linear task setting are idealized; no experiments reported in the text shown.
Experimental setup
Datasets: synthetic linear regression prompts with orthonormal feature tokens
Metrics: squared prediction loss per-feature prediction error attention score concentration
Baselines: linear self-attention training dynamics analyses
Builds on: linear transformer in-context learning dynamics (gradient flow analyses) GPT-2 in-context learning framework (Garg et al.) Ahn et al. transformers as GD learners
Closest papers by meaning
- In-context Learning on Function Classes Unveiled for Transformers (ICML 2024 · similarity 0.92)
- How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression (NeurIPS 2024 · similarity 0.91)
- In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness (NeurIPS 2024 · similarity 0.91)
- Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape (ICML 2024 · similarity 0.91)
- On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures (ICML 2025 · similarity 0.90)
- Transformers are almost optimal metalearners for linear classification (NeurIPS 2025 · similarity 0.90)
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.