Research map / Architectures and efficiency / Attention and transformers
How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations
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.
Constructs synthetic in-context regression tasks where labels are a fixed MLP representation composed with a per-instance linear function, proves transformers can implement representation-then-linear-ICL, and probes trained transformers to find matching lower/upper-module mechanisms.
Problem
Existing theory and mechanistic work on in-context learning (ICL) mostly covers simple function classes like raw linear functions; it is unclear how transformers perform ICL when a fixed nonlinear representation must be learned alongside in-context linear adaptation.
Approach
ICL tasks where labels = linear function of a fixed MLP representation, with the optimal algorithm being representation transform followed by ridge regression. The paper gives decoder-only transformer constructions that predict at every token, and uses linear probes and a pasting experiment to show trained transformers split into a lower representation-computing module and an upper linear-ICL modu
Key results
- Trained transformers match the Bayes-optimal ridge predictor across settings of depth L, hidden dimension D, and noise sigma in the representation ICL setting.
- Pasting the upper layers of a trained (L=3, D=20, sigma=0.1) transformer with a one-layer transformer embedding yields nearly optimal linear ICL, while a freshly trained one-layer transformer alone achieves about trivial risk 1.01.
- Probes show representation is computed at lower layers then copied onto y-tokens (e.g., layer 5-6 for L=4), with the ridge prediction decreasing monotonically through upper layers.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train a transformer to learn representations and perform linear ICL | GPT-2-family transformer · Synthetic supervised ICL with representations | test risk | Helped |
| Train a transformer to learn representations and perform linear ICL | GPT-2-family transformer · Synthetic dynamical systems with representations | ICL risk | Helped |
| Use a mixture of multiple possible representation functions | Trained transformer · Mixture of multiple representations | ICL risk | Hurt |
| Use a one-layer transformer embedding before the upper module (ablation) | Trained transformer's upper module · Linear ICL pasting, D-dimensional, noise 0.1 | test risk | Helped |
| Use a linear-copy embedding rather than a linear embedding (ablation) | Trained transformer's upper module · Linear ICL pasting, D-dimensional, noise 0.1 | test risk | Helped |
| Use a linear embedding before the upper module (ablation) | Trained transformer's upper module · Linear ICL pasting, D-dimensional, noise 0.1 | test risk | Helped |
| Use a freshly trained single-layer transformer with separate input and label tokens (ablation) | Single-layer transformer · Linear ICL with separate input and label tokens | test risk: 1.01 | Hurt |
| Increase representation noise or hidden dimension rather than representation depth (ablation) | Trained transformer · Synthetic supervised ICL with representations | test risk | Mixed |
Surprising findings
- A freshly trained single-layer transformer performs at trivial risk (~1.01), yet the upper module with a one-layer embedding performs near-optimal linear ICL.
Limitations
- Setting uses synthetic data with idealized, fixed random-orthogonal MLP representations.
- Mixture-of-representations extension is harder and transformers approach Bayes-optimal risk less closely.
Experimental setup
Datasets: synthetic ICL regression tasks with MLP representations
Benchmarks: synthetic ICL regression with representations synthetic dynamical system ICL
Metrics: ICL test risk linear probing error pasting risk
Baselines: Bayes-optimal ridge predictor
Builds on: Garg et al. in-context learning of function classes Akyürek et al. ridge regression in-context von Oswald et al. in-context gradient descent Bai et al. algorithm selection
Closest papers by meaning
- Towards Understanding How Transformers Learn In-context Through a Representation Learning Lens (NeurIPS 2024 · similarity 0.94)
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.93)
- In-Context Learning through the Bayesian Prism (ICLR 2024 · similarity 0.93)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.92)
- Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions (ICLR 2024 · similarity 0.92)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 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.