Research map / Optimization / Implicit bias and generalization
On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability
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-flow training of a one-layer linear causal self-attention transformer on first-order AR data, proving it implements one step of OLS gradient descent in-context under a moment condition, but that the learned mesa-optimizer only recovers the distribution under stronger conditions.
Problem
Whether the non-convex training dynamics of autoregressively trained transformers converge to the mesa-optimizer that performs in-context gradient descent, and what capability limits such a mesa-optimizer has.
Approach
Studies gradient flow of a one-layer linear causal self-attention model on a unitary-diagonal first-order AR process without the strong diagonal-structure assumption of prior work; shows convergence to one-step GD on OLS under a moment condition on the initial token, and derives necessary and sufficient moment conditions for distribution recovery.
Key results
- Under diagonal initialization and the moment condition, the gradient flow converges to a solution with ab = kappa_1/(kappa_2 + (kappa_3/(T-2)) sum 1/(t-1)), so the trained transformer implements one step of GD for OLS in-context.
- For Gaussian initial tokens, the learned predictor converges to x_{T_te}/5 in expectation and fails to recover the true next token, while the sparse-vector example is recovered.
- Without the non-diagonal masking, the trained transformer is shown not to perform vanilla one-step GD in general (Proposition on non-diagonal gradients).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a Gaussian initial token distribution instead of a distribution satisfying the success condition | One-layer linear transformer trained on an AR process · AR process, Gaussian initial token | expected prediction-to-target ratio: 1/5 | Hurt |
| Use a sparse initial token distribution from Example (sparse vector) | One-layer linear transformer trained on an AR process · AR process, sparse initial token | prediction recovery | Helped |
| Mask non-diagonal gradients of W^KQ_32 and W^PV_12 during training (ablation) | One-layer linear transformer with fixed all-ones initial token · AR process, fixed all-ones initial token | vanilla gradient-descent mesa-optimization: ã b̃ = 1/(1 + (d-1)/(T-2) ∑_{t=2}^{T-1} 1/(t-1)) | Helped |
| Remove non-diagonal gradient masking (ablation) | One-layer linear transformer with fixed all-ones initial token · AR process, fixed all-ones initial token | vanilla gradient-descent mesa-optimization | Hurt |
| Use standard Gaussian initialization instead of diagonal initialization (ablation) | One-layer linear transformer · AR process, Gaussian initialization | test results (ratio or MSE loss) | No clear effect |
| Vary diagonal initialization (a₀,b₀) among (0.1,0.1), (0.5,1.5), and (2,2) (ablation) | One-layer linear transformer · AR process, Gaussian and sparse initial tokens | learned parameter product ab: {1}/{5σ^2}; {1}/{c^2} | No clear effect |
Surprising findings
- With a standard Gaussian initial token, the trained one-layer linear transformer fails to recover the AR process, predicting only about one-fifth of the true next token in expectation.
- The trained transformer in general does not implement vanilla gradient descent for the OLS problem once the diagonal structure assumption is dropped.
Limitations
- Analysis restricted to a one-layer linear transformer, not softmax or multi-layer models.
- General case without the diagonal/moment assumptions is not fully resolved.
- Relies on diagonal initialization for tractability, with standard Gaussian initialization only tested empirically.
Experimental setup
Datasets: synthetic first-order AR process x_{t+1}=W x_t with unitary diagonal W
Benchmarks: synthetic AR process simulations
Metrics: convergence of ab product prediction ratio to true next token MSE loss
Baselines: Zhang et al. 2024 (diagonal-structure AR transformer loss landscape)
Builds on: mesa-optimization hypothesis in-context learning for linear regression theory Ahn et al. 2023 linear transformer global minima analysis
Closest papers by meaning
- How do Transformers Perform In-Context Autoregressive Learning ? (ICML 2024 · similarity 0.89)
- Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functions (NeurIPS 2025 · similarity 0.88)
- Transformers Learn to Achieve Second-Order Convergence Rates for In-Context Linear Regression (NeurIPS 2024 · similarity 0.88)
- Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.87)
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.87)
- Does learning the right latent variables necessarily improve in-context learning? (ICML 2025 · similarity 0.87)
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.