Research map / Architectures and efficiency / Attention and transformers
How Do Nonlinear Transformers Learn and Generalize in In-Context Learning?
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.
Provides the first theoretical analysis of training dynamics of a one-layer single-head Transformer with softmax attention and ReLU MLP for in-context learning on binary classification tasks, with in-domain and out-of-domain generalization and magnitude-based pruning guarantees.
Problem
How a Transformer with nonlinear self-attention and nonlinear MLP can be trained to perform in-context learning and generalize to unseen tasks and out-of-domain data, and how magnitude pruning affects ICL.
Approach
Analyzes SGD training of a one-layer softmax-attention + ReLU-MLP Transformer on structured binary classification prompts; shows attention concentrates on contexts sharing the query's relevant pattern, the MLP encodes label embeddings, and ICL generalizes to unseen and linearly-combined out-of-domain patterns, with magnitude pruning of small MLP neurons being harmless.
Key results
- Out-of-domain classification error falls below 0.01 when S1 >= 1 (the coefficient-sum condition) and deviates from 0 when S1 < 1.
- Magnitude-based pruning of low-magnitude MLP neurons does not hurt out-of-domain generalization up to about 15% pruning; pruning large-magnitude neurons increases error with pruning rate.
- Required training iterations scale as alpha^(-2/3) and training context length as alpha^(-1); in-domain generalization achieved with about M1 training tasks, a fraction (M1-1)^(-1/2) of all in-domain tasks.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase the fraction α of training contexts containing decisive patterns | One-layer Transformer training · In-domain generalization and training | required training samples: proportional to α^{-2/3} | Helped |
| Increase the fraction α of training contexts containing decisive patterns | One-layer Transformer training · In-domain generalization and training | required iterations: proportional to α^{-2/3} | Helped |
| Increase the fraction α' of testing contexts sharing the query's relevant pattern | One-layer Transformer · In- and out-of-domain generalization | required testing context length: linear in α'^{-1} | Helped |
| Increase α' for the pruned model | Magnitude-pruned one-layer Transformer · Out-of-domain generalization | generalization error | Helped |
| Prune small-magnitude neurons rather than large-magnitude neurons | One-layer Transformer · In- and out-of-domain generalization | generalization error: O(ε+M_1^{-1/2}); Ω( R+ (α' M_1)^{-1}) | Mixed |
| Prune small-magnitude neurons at a rate below around 15% | One-layer Transformer · Out-of-domain generalization | generalization error: around 15% | No clear effect |
| Increase the pruning rate when removing large-weight neurons | One-layer Transformer · Out-of-domain generalization | generalization error | Hurt |
| Set the OOD relevant-pattern coefficient sum S_1 to at least 1 | One-layer Transformer · Out-of-domain classification | classification error: <0.01 when S_1≥1 | Helped |
| Reduce the fraction α' of testing contexts sharing the query's relevant pattern | One-layer Transformer · In- and out-of-domain generalization | required testing context length | Hurt |
| Use ICL instead of the compared classification algorithms | ICL · Classification with α'=0.6 | sample efficiency | Helped |
| Use ICL instead of the compared classification algorithms | ICL · Classification with α'=0.8 | generalization performance | No clear effect |
| Train on fewer than M_1 tasks | One-layer Transformer · In-domain unseen-task generalization | generalization error: M_1 | Hurt |
Surprising findings
- Out-of-domain generalization holds only when the out-of-domain patterns are linear combinations of training patterns with coefficient sum at least 1; the authors show the error deviates from zero otherwise.
- ICL on this setup is not significantly better than other learners when alpha' = 0.8, advantage appears mainly at alpha' = 0.6.
Limitations
- Analysis restricted to a one-layer, single-head Transformer with ReLU MLP and a simplified orthogonal-pattern data model.
- Binary classification only; generation tasks left to future work.
- Multi-head analysis only sketched as an extension, not proven.
Experimental setup
Datasets: synthetic binary classification prompts from structured in-domain and out-of-domain patterns
Benchmarks: synthetic in-domain and out-of-domain binary classification tasks
Metrics: in-context classification error generalization error
Baselines: linear-attention ICL theory GPT-2 (3-layer, 2-head) experiments
Builds on: Garg et al. ICL framework feature learning theory for two-layer ReLU networks max-margin / attention training dynamics analyses
Closest papers by meaning
- How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations (ICLR 2024 · similarity 0.93)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.93)
- Towards Understanding How Transformers Learn In-context Through a Representation Learning Lens (NeurIPS 2024 · similarity 0.93)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.93)
- Position: Do pretrained Transformers Learn In-Context by Gradient Descent? (ICML 2024 · similarity 0.92)
- Exact Conversion of In-Context Learning to Model Weights in Linearized-Attention Transformers (ICML 2024 · similarity 0.92)
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.