Research map / Architectures and efficiency / Attention and transformers
How do Transformers Perform In-Context Autoregressive Learning ?
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.
Trains linear Transformers on next-token prediction for first-order AR processes s_{t+1}=W s_t with orthogonal/unitary context matrices, and characterizes the minima showing in-context learning as gradient descent (augmented tokens) or trigonometric geometric relations with orthogonal heads (non-augmented).
Problem
Understanding how Transformers perform in-context learning of autoregressive processes, i.e., how a trained causal Transformer estimates the context-dependent transition matrix W in-context and then predicts the next token.
Approach
Frames in-context autoregressive learning as a two-step procedure (in-context mapping estimating W, then a prediction mapping) and analytically characterizes the global minima of a one-layer linear attention model's next-token loss under commutativity assumptions, showing a gradient-descent step (augmented tokens) or an orthogonal multi-head solution with trigonometric structure (non-augmented tok
Key results
- In the augmented setting, linear attention layers beat L gradient-descent steps only for L<3; GD outperforms the linear Transformer by orders of magnitude for L>=3.
- Under commutativity, the optimal one-layer augmented model implements one gradient-descent step with step size asymptotically 3/(2 T_max).
- Trained non-augmented positional-encoding models exhibit orthogonality between heads, and the orthogonal-case solution implements the identity 2cos(theta)R_theta - I = R_{2theta}.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use softmax attention instead of linear attention (ablation) | Transformer on commuting context matrices · Autoregressive next-token prediction | in-context mapping fit | Hurt |
| Increase the number of attention heads (ablation) | Non-augmented linear Transformer · Autoregressive next-token prediction | MSE | No clear effect |
| Duplicate token representations across dimensions (ablation) | Non-augmented linear Transformer · Autoregressive next-token prediction | MSE | Helped |
| Use multiple linear Transformer layers instead of the same number of inner-loss gradient descent steps (ablation) | Augmented-token linear Transformer · Autoregressive next-token prediction | MSE | Mixed |
| Add softmax attention and component-wise MLP layers (ablation) | Vanilla Transformer · Augmented-token autoregressive next-token prediction | MSE | Helped |
| Shuffle tokens across sequences instead of preserving their order (ablation) | Autoregressive-process fit · Moby Dick token representations | sequences with fitting error above 10^-12: 4 times more | Hurt |
| Increase μ, narrowing the context-matrix distribution (ablation) | Positional encoding-only model · Autoregressive process with λ = e^{i θ / μ} | positional encoding smoothness | Helped |
Surprising findings
- Linear attention fails to compete with gradient descent beyond two layers, while a full softmax+MLP Transformer outperforms gradient descent.
- Softmax attention prevents finding the global minimum of the training loss even in dimension one.
Limitations
- Analysis restricted to commuting (diagonalizable) context matrices; non-commuting case only studied experimentally.
- One-layer linear attention with a fixed tau=T-1 prediction target; the tau=T case yields no zero-loss solution.
- Simple AR process is an overly simplistic model for real sentences.
Experimental setup
Datasets: Moby Dick text (GPT-2 token embeddings) for validating AR process choice synthetic orthogonal AR process sequences
Metrics: mean squared error population loss
Baselines: gradient descent on inner loss softmax attention Transformer cosine absolute positional encoding
Builds on: Transformers learn in-context by gradient descent One-step gradient descent is optimal in-context learner for linear self-attentio
Closest papers by meaning
- Compositional Capabilities of Autoregressive Transformers: A Study on Synthetic, Interpretable Tasks (ICML 2024 · similarity 0.89)
- In-Context Learning with Representations: Contextual Generalization of Trained Transformers (NeurIPS 2024 · similarity 0.89)
- On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability (NeurIPS 2024 · similarity 0.89)
- Towards Understanding the Universality of Transformers for Next-Token Prediction (ICLR 2025 · similarity 0.88)
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.88)
- In-context Learning on Function Classes Unveiled for Transformers (ICML 2024 · similarity 0.88)
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.