Research map / Architectures and efficiency / Attention and transformers
On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures
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.
Proves that a single-layer linear-attention transformer trained by gradient descent converges linearly to a global minimizer for in-context binary and multi-class classification of Gaussian mixtures, and bounds the ICL inference error in terms of training and test prompt lengths.
Problem
Theoretical understanding of training dynamics of transformers for in-context classification is limited; prior work mostly covers linear regression or binary classification with finite/special data patterns.
Approach
Analyze a simplified single-layer linear-attention transformer with sparse parameters on Gaussian mixture ICL tasks; show the expected loss is strongly convex and smooth on compact sets, so GD converges linearly, and the global minimizer is a constant (2Λ^{-1} or cΛ^{-1}) plus an O(1/N) error, which approximates LDA/softmax regression in context.
Key results
- Expected loss gradient descent converges linearly: ||W^t - W*||_F^2 <= exp(-t/kappa) ||W^0 - W*||_F^2 for the binary and multi-class cases.
- Global minimizer satisfies ||G||_max = O(1/N) so W* approaches 2Λ^{-1} (binary) or cΛ^{-1} (multi-class) as training prompt length grows.
- Expected total variation inference error is O(1/N + 1/sqrt(M)) for binary and O(c^2/N + c^{3/2}/sqrt(M)) for multi-class classification.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase training prompt length N and test prompt length M | 1-layer and 3-layer transformers · Gaussian mixture classification | inference error | Helped |
| Increase the number of classes c | 1-layer and 3-layer transformers · Gaussian mixture classification | inference error | Hurt |
| Use component means with equal rather than unequal Λ⁻¹-weighted norms | 1-layer and 3-layer transformers · Gaussian mixture classification; norm condition | in-context classification performance | Helped |
| Train with unequal Λ⁻¹-weighted norms | 1-layer versus 3-layer transformers · Gaussian mixture classification; norm shift | in-context classification performance | Mixed |
| Train with varying covariances rather than a fixed covariance | 1-layer versus 3-layer transformers · Gaussian mixture classification; covariance shift | in-context classification performance | Mixed |
| Use transformer models instead of classical methods | 1-layer, sparse; 1-layer, full; and 3-layer transformers · Three Gaussian mixtures | ICL performance | Helped |
Surprising findings
- Experiments show the 3-layer softmax transformer is more robust than the 1-layer model under unequal-norm and covariance shifts.
Limitations
- Analysis restricted to single-layer linear attention with sparse parameterization; multi-layer softmax not covered theoretically.
- Requires training and test covariances to match and equal-Λ^{-1}-norm class means.
- Even with infinite test prompts, error persists unless training prompt length N grows.
Experimental setup
Metrics: total variation distance inference error
Baselines: softmax regression linear discriminant analysis 3-layer softmax transformer
Builds on: Zhang et al. 2024 training dynamics of linear transformers for in-context linear Huang et al. 2023 training dynamics of single-layer transformer for in-context c
Closest papers by meaning
- Transformers are almost optimal metalearners for linear classification (NeurIPS 2025 · similarity 0.91)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.90)
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.90)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.90)
- In-Context Learning with Representations: Contextual Generalization of Trained Transformers (NeurIPS 2024 · similarity 0.89)
- Unlabeled Data Can Provably Enhance In-Context Learning of Transformers (NeurIPS 2025 · similarity 0.89)
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.