Research map / Architectures and efficiency / Attention and transformers
Towards Understanding How Transformers Learn In-context Through a Representation Learning Lens
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.
Shows that one softmax attention layer's ICL inference equals one gradient step on a kernel "dual model" (via positive random features), derives a generalization bound in demonstration count, and proposes contrastive-inspired attention modifications tested on synthetic regression.
Problem
Understanding the mechanism of in-context learning in softmax-attention Transformers without restrictive assumptions on weights or input token structure (prior gradient-descent links were restricted to linear attention and linear regression).
Approach
Use kernel feature maps of the softmax kernel to define a dual linear model f(x)=W phi(x); the ICL output of one softmax attention layer equals the test prediction after one SGD step on this dual model with a cosine-style loss over demonstration key/value pairs. This representation-learning view is then linked to contrastive learning without negatives, and used to motivate regularized, augmented,
Key results
- On synthetic linear regression with N=15 demonstrations, the dual model's test prediction after training matches the softmax attention ICL output (reported as exact equivalence).
- Regularization with alpha>0 converged to poorer results, while alpha<0 converged faster and reached results comparable to the unregularized model.
- Augmented models with a two-layer MLP key augmentation eventually converged faster and to a better solution; negative models with k=3, beta=0.1 showed slightly faster convergence.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use the dual model's one-step gradient descent process to represent one softmax attention layer | One softmax attention layer · Synthetic linear regression | equivalence | Helped |
| add positive regularization (α > 0) (ablation) | Attention layer · Synthetic linear regression | final performance | Hurt |
| use negative regularization (α < 0) (ablation) | Attention layer · Synthetic linear regression | convergence speed | Helped |
| apply a single-layer MLP augmentation to keys only (ablation) | Attention layer · Synthetic linear regression | convergence speed | Helped |
| use a two-layer MLP for key augmentation (ablation) | Attention layer · Synthetic linear regression | convergence speed | Mixed |
| subtract negative-sample values using k=3 and β = 0.1 (ablation) | Attention layer · Synthetic linear regression | convergence speed | Helped |
Surprising findings
- Negative regularization (alpha<0) converged faster, contrary to the contrastive-loss motivation for positive alpha.
- Key-only augmentation converged slightly faster than augmenting both key and value.
Limitations
- Analysis covers a simplified Transformer; LayerNorm, residual connections and other structures are not analyzed.
- Proposed attention modifications are only validated on synthetic tasks and need task-specific design of augmentations and negatives.
- Equivalence relies on one SGD step with a specific loss and on a kernel approximation (random features) of softmax.
Experimental setup
Datasets: synthetic linear regression tasks
Benchmarks: synthetic linear regression
Metrics: L2 distance between dual-model test prediction and ICL outpu MSE training loss
Baselines: normal (unregularized) attention model ICL as implicit fine-tuning (linear attention)
Builds on: Performer (positive random features) linear attention dual form (Irie et al.) Dai et al. ICL as implicit fine-tuning von Oswald et al. transformers learn in-context by gradient descent
Closest papers by meaning
- How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations (ICLR 2024 · similarity 0.94)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.93)
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.93)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.92)
- Understanding the Generalization of In-Context Learning in Transformers: An Empirical Study (ICLR 2025 · similarity 0.90)
- Unlabeled Data Can Provably Enhance In-Context Learning of Transformers (NeurIPS 2025 · similarity 0.90)
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.