Research map / Architectures and efficiency / Attention and transformers
Breaking through the learning plateaus of in-context learning in Transformer
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.
Studies learning plateaus in in-context learning via a weights/context component decomposition on a synthetic Shapes3D task, and proposes weights warm-up, mixed training, and extra weights loss to shorten plateaus, validated on SST-ICL and WordSelection tasks.
Problem
Transformers take a long time to acquire in-context learning because of learning plateaus with little progress; the mechanism behind these plateaus is unclear and the usual remedy (scaling the model) is computationally costly.
Approach
Conceptually decompose the Transformer's prediction into a weights component (encoding per-example factor values) and a context component (inferring the hidden factor from examples); show plateau length correlates with dysfunction of the weights component, and speed learning by explicitly improving the weights component via warm-up on a fixed-mapping task, mixed training, or an auxiliary factor-pr
Key results
- Plateau length grows with task entropy (complexity), and weights-component score after 50 epochs negatively correlates with plateau length.
- Mixed training (half D_fix, half D_rnd) lets a 6.69M model after 20 epochs match a 13.88M model trained 50 epochs on D_rnd.
- Adding an extra weights loss speeds breaking through plateaus on Shapes3D, while an extra context loss gives only marginal benefit; weights dysfunction appears on SimpleFunction when d>30.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Warm up on D_fix before training on D_rnd | Transformer · Shape3D D_rnd | accuracy | Helped |
| Warm up on D_fix before training on D_rnd | Transformer · Shape3D D_rnd | plateau length | Helped |
| Train on half D_fix and half D_rnd data | Transformer · Shape3D D_rnd | accuracy: 6.69M; 20 epochs; 13.88M; 50 epochs | Helped |
| Train on half D_fix and half D_rnd data | Transformer · Shape3D D_rnd | plateau length | Helped |
| Add extra weights loss | Transformer · Shape3D D_rnd | accuracy: λ = 0.1 | Helped |
| Add extra context loss instead of extra weights loss (ablation) | Transformer · Shape3D D_rnd | accuracy: λ = 0.1 | No clear effect |
| Add extra weights loss | Transformer · Shape3D D_rnd | weights comp. score: λ = 0.1 | Helped |
| Increase mapping-function entropy (ablation) | Transformer · Shape3D | plateau length | Hurt |
| Increase mapping-function entropy (ablation) | Transformer · Shape3D | weights comp. score | Hurt |
| Train on mixed D_fix and D_rnd data | Transformer · SST-ICL | accuracy: less than 30 epochs; after the 35 epochs | Mixed |
| Add extra weights loss | GPT2 · WordSelection | accuracy | Helped |
| Add extra weights loss | 6-layer Transformer · SimpleFunction (d > 30) | test accuracy: d > 30; 50 epochs | Helped |
Surprising findings
- The weights component can deteriorate during training on high-complexity random-mapping tasks, falling below its initialization value.
- On SST-ICL, fixed-mapping pretraining hurts performance for small epoch counts before helping later.
Limitations
- Analysis is conceptual and relies on linear probes, so the weights/context decomposition is not physically identified.
- The theoretical construction uses a relaxed Transformer (identity attention nonlinearity, ReLU MLP) and is weaker than a full learning-dynamics analysis.
- Extra weights loss is not applicable directly when no fixed-mapping setting exists, so WordSelection uses an auxiliary loss instead; the scope is limited to ICL with the stated tasks.
Experimental setup
Datasets: Shapes3D SST WordSelection (constructed from 2000 words) Pile (Pythia pretraining data, referenced) SimpleFunction (Gaussian linear sign task)
Benchmarks: Shapes3D in-context task SST-ICL WordSelection SimpleFunction
Metrics: in-context accuracy weights component score context component score MSE_p
Baselines: D_rnd training without intervention
Builds on: Shapes3D VAE
Closest papers by meaning
- Toward Understanding In-context vs. In-weight Learning (ICLR 2025 · similarity 0.90)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.89)
- Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions (ICLR 2024 · similarity 0.89)
- Emergence and Effectiveness of Task Vectors in In-Context Learning: An Encoder Decoder Perspective (ICML 2025 · similarity 0.89)
- Technical Debt in In-Context Learning: Diminishing Efficiency in Long Context (NeurIPS 2025 · similarity 0.89)
- MLPs Learn In-Context on Regression and Classification Tasks (ICLR 2025 · similarity 0.88)
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.