Research map / Architectures and efficiency / Attention and transformers
Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape
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.
Analyzes gradient dynamics of a one-MLP-plus-linear-attention Transformer for in-context learning of nonlinear features, showing via mean-field and two-timescale limits that the attention loss landscape is benign and Wasserstein gradient flow avoids saddles, with convergence rates.
Problem
How in-context learning over nonlinear representations emerges in Transformers with MLP layers trained by gradient descent; prior dynamics analyses only covered single attention layers on linear tasks.
Approach
Model the Transformer as a two-layer MLP followed by a linear attention layer, factor out attention via a two-timescale limit so the MLP's distribution-level loss is analyzed in the mean-field regime; show no spurious local minima, almost-sure saddle avoidance of Wasserstein gradient flow, and rates via birth-death and perturbation modifications.
Key results
- Shows the infinite-dimensional ICFL loss has no spurious local minima: every non-global critical point is a strict saddle or has a descent direction.
- Proves Wasserstein gradient flow almost always avoids strict saddle points (for initializations outside a measure-zero set).
- Gives convergence rates: accelerated phase with O(k^2/(gamma*eps)) time to reach loss eps with birth-death; toy experiments show modified dynamics helps for degenerate features.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the joint attention-and-MLP objective with directly minimizing the static objective | Toy Transformer on standard features · Synthetic ICL training prompts | training loss | No clear effect |
| Add birth-death and perturbation dynamics | Modified Transformer · Synthetic ICL with degenerate features | training loss | Helped |
| Learn the relevant nonlinear features during pretraining | Toy Transformer · Test task x ↦ ||h_{μ°}(x)|| | test loss | Helped |
Surprising findings
- Attention-only linear transformers can only implement linear-function learning, while the MLP-attention model learns the Barron class of nonlinear features; the static (no-attention) model behaves similarly to the attention model after two-timescale reduction.
Limitations
- Two-timescale and mean-field limits are formal; infinite-prompt and infinite-width assumptions are idealized.
- Perturbed-dynamics polynomial escape time is only conjectured, not proven.
- Experiments are small toy settings with d=20, k=5.
Experimental setup
Datasets: synthetic Gaussian inputs (N(0, I_d)) with random linear regression tasks over features
Benchmarks: synthetic ICFL toy experiments
Metrics: training loss test loss
Baselines: attention Transformer (joint optimization) static Transformer modified Transformer (birth-death and perturbation)
Builds on: mean-field neural network theory linear attention in-context learning center-stable manifold theorem
Closest papers by meaning
- In-context Convergence of Transformers (ICML 2024 · similarity 0.91)
- Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis (NeurIPS 2024 · similarity 0.90)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.89)
- Infinite Limits of Multi-head Transformer Dynamics (NeurIPS 2024 · similarity 0.89)
- Transformers Learn Faster with Semantic Focus (NeurIPS 2025 · similarity 0.89)
- From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics (NeurIPS 2025 · 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.