Research map / Architectures and efficiency / Attention and transformers
In-Context Learning with Representations: Contextual Generalization of Trained 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 gradient-descent training of a one-layer multi-head softmax-attention transformer on noisy non-linear regression ICL with partial prompts, proving linear convergence and showing the trained model implements ridge regression over basis functions.
Problem
How transformers learn contextual (template) information to predict unseen examples when prompts contain only a few noisy query-answer pairs that underdetermine the template; prior provable ICL training-dynamics work requires long prompts or restrictive assumptions.
Approach
Model ICL as regression where labels are a linear combination of m basis functions with Gaussian label noise; prove GD on a one-layer multi-head softmax transformer converges linearly to a global minimum without orthogonal-token or long-prompt assumptions, and that the learned model performs ridge regression over the basis at inference time.
Key results
- Theorem shows training loss converges linearly: L(theta_t)-inf L <= (1 - eta_w zeta_0/(2K))^t times initial gap, under assumptions on lambda distribution and initialization with H >= N.
- Inference output converges to ridge regression predictions y_hat* with error bound epsilon after T ~ log(1/epsilon) iterations, with probability 1-delta.
- Experiments with N=30, K=200, d=100, m=20, H=64 show training and inference losses descend for both 1-layer and 4-layer transformers, including on out-of-domain lambda; the 1-layer training loss stalls for H<N=30 and diverges at H=400.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train a 1-layer or 4-layer transformer with multi-head attention | Transformer on synthetic regression ICL · In-domain and OOD test sets | inference loss | Helped |
| Evaluate on OOD tasks sampled from N(1_m, 4I_m) instead of in-domain tasks | 1-layer and 4-layer transformers · OOD test set vs. in-domain test set | inference loss | Hurt |
| Set the representation dimension m closer to prompt length N | Pretrained 1-layer transformer · Synthetic regression ICL; N from 50 to 150, m=100 | performance gap to best prediction | Helped |
| Use fewer than N attention heads (ablation) | 1-layer transformer · Synthetic regression ICL; N=30 | training loss | Hurt |
| Set the number of attention heads to H=30 (ablation) | 1-layer transformer · Synthetic regression ICL; N=30 | training loss | Helped |
| Increase the number of attention heads to H=400 (ablation) | 1-layer transformer · Synthetic regression ICL | training loss | Hurt |
| Increase the number of attention heads in the 4-layer transformer (ablation) | 4-layer transformer · Synthetic regression ICL; fixed 100-second training time | training and inference loss | Hurt |
| Use H=4 rather than larger head counts in the 4-layer transformer (ablation) | 4-layer transformer · Synthetic regression ICL; fixed 100-second training time | final training loss | Helped |
Surprising findings
- Larger numbers of attention heads slow convergence and H=400 causes instability and divergence in the 1-layer model.
- In the 4-layer experiments the final training loss was smallest at H=4, so deeper transformers do not need many heads for convergence.
Limitations
- Analysis is restricted to a one-layer transformer with a simplified parameterization (identity output matrix, value and key-query structure fixed).
- Requires H >= N for the initialization assumption to hold almost surely, and larger H slows convergence.
- Assumes lambda has independent zero-mean unit-variance entries and inference lambda is bounded in norm.
Experimental setup
Datasets: synthetic Gaussian tokens and representations synthetic regression tasks with Gaussian lambda
Benchmarks: synthetic in-domain and out-of-domain regression test sets
Metrics: training loss inference loss (1/K ||y_hat - y_hat*||^2) performance gap to best ridge estimator
Baselines: Huang et al. one-layer softmax ICL dynamics Li et al. multi-task linear regression dynamics Zhang et al. linear attention dynamics
Builds on: Transformers learn in-context by gradient descent Training dynamics of one-layer softmax attention ICL (Huang et al.) Transformers as statisticians / representation-based ICL
Closest papers by meaning
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.94)
- How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations (ICLR 2024 · similarity 0.91)
- In-Context Learning through the Bayesian Prism (ICLR 2024 · similarity 0.91)
- In-context Learning in Presence of Spurious Correlations (NeurIPS 2026 · similarity 0.91)
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.91)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.91)
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.