Research map / Architectures and efficiency / RNNs and normalization
How Data Mixing Shapes In-Context Learning: Asymptotic Equivalence for Transformers with MLPs
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 a linear-attention Transformer with a two-layer MLP head (first layer trained by one gradient step, second by ridge) is asymptotically equivalent in ICL error to a degree-p Hermite polynomial model, and uses this to analyze how data mixing of heterogeneous sources affects ICL and feature learning.
Problem
In-context learning theory mostly relies on simplified architectures (no MLPs), isotropic linear-regression data, and single-source training, which limits relevance to realistic Transformers trained on heterogeneous mixtures of data sources.
Approach
Under high-dimensional proportional asymptotics, a linear-attention Transformer with a trained nonlinear two-layer MLP head has the same ICL error as a finite-degree Hermite polynomial surrogate, which makes the effects of multi-source data mixing, data quality (noise, covariance structure) and feature learning analyzable in closed form.
Key results
- The nonlinear MLP Transformer outperforms the linear Transformer in ICL error on synthetic nonlinear tasks and closely tracks its degree-p polynomial surrogate, even at moderate dimension (d=80).
- Increasing the mixing proportion of a structured-input source or structured-task source lowers ICL error, and lowering target noise in the second source improves performance.
- Feature learning (larger step size) helps only when task covariance is structured; with isotropic task vectors increasing eta does not improve ICL error.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a nonlinear MLP head instead of a linear Transformer | Transformer for nonlinear ICL · Synthetic nonlinear ICL | ICL error | Helped |
| Replace the nonlinear MLP with its polynomial surrogate | Transformer with nonlinear MLP · Synthetic nonlinear ICL | ICL error | No clear effect |
| Increase the proportion of data from a source with structured input covariance | Transformer with nonlinear MLP · Synthetic nonlinear ICL | ICL error | Helped |
| Increase the proportion of data from a source with structured task covariance | Transformer with nonlinear MLP · Synthetic nonlinear ICL | ICL error | Helped |
| Lower the target noise variance of the second data source | Transformer with nonlinear MLP · Synthetic nonlinear ICL | ICL error | Helped |
| Increase the first-layer gradient step size with isotropic task vectors | Transformer with nonlinear MLP · Synthetic ICL with structured input covariance | ICL error | No clear effect |
| Increase the first-layer gradient step size with structured task covariance | Transformer with nonlinear MLP · Synthetic ICL with structured task covariance | ICL error | Helped |
| Increase context length | Transformer with nonlinear MLP · Synthetic nonlinear ICL | ICL error | Helped |
| Increase the first-layer gradient step size | Transformer with nonlinear MLP · Multilingual Amazon Reviews Corpus | ICL error | Helped |
| Increase the proportion of English reviews in the training mixture | Transformer with nonlinear MLP · Multilingual Amazon Reviews Corpus | ICL error | Helped |
Surprising findings
- Increasing the step size does not help when only input covariance is structured, indicating feature learning depends on task-covariance structure.
Limitations
- Theory relies on strong high-dimensional assumptions (proportional limit, scaling of step size, low-rank covariance) that the authors acknowledge as limitations.
- Single-block linear attention and two-layer MLP only; no multi-layer or softmax attention analyzed.
- Real-world experiment is a small single-setting demonstration with PCA-reduced embeddings.
Experimental setup
Datasets: Multilingual Amazon Reviews Corpus
Benchmarks: Multilingual Amazon Reviews Corpus
Metrics: ICL error (squared loss)
Baselines: linear Transformer (linear attention without MLP) equivalent degree-p polynomial model
Builds on: linear attention reparameterization two-stage feature learning analysis of two-layer networks Gaussian equivalence
Closest papers by meaning
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.91)
- In-Context Learning through the Bayesian Prism (ICLR 2024 · similarity 0.90)
- How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations (ICLR 2024 · similarity 0.90)
- MLPs Learn In-Context on Regression and Classification Tasks (ICLR 2025 · similarity 0.89)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.89)
- Pretrained Transformer Efficiently Learns Low-Dimensional Target Functions In-Context (NeurIPS 2024 · similarity 0.89)
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.