Research map / Optimization / Implicit bias and generalization
CausalLM is not optimal for in-context 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.
Theoretically analyzes multi-layer linear self-attention transformers for in-context linear regression, showing prefixLM converges to the least-squares solution while causalLM follows non-decaying online gradient descent and may not become optimal, supported by synthetic and T5-based experiments.
Problem
Empirical evidence shows prefixLM outperforms causalLM for in-context learning, but there is no theoretical explanation for this gap in transformer ICL.
Approach
Models multi-layer linear self-attention (LSA) under a fixed construction that implements gradient descent per layer, then shows that full (prefix) attention yields a single GD trajectory converging to the least-squares solution, whereas causal attention yields position-wise weights that follow online GD with constant step size and need not reach the optimum.
Key results
- Multi-layer LSA prefixLM query error decays linearly to the least-squares solution, while causalLM query error stalls above 1e-1 for the LSA construction.
- On 5-shot MMLU, T5X DecoderOnly prefixLM Base/Large/XL reach 28.8/32.0/39.5 vs causalLM 28.0/26.9/30.5.
- Standard 24-layer transformers: prefixLM-SL achieves 8.6e-3 query MSE on linear regression vs 1.9e-1 for causalLM-SL.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace causalLM attention with prefixLM attention in a shared-layer transformer | softmax transformer ICL · Linear regression (LR) | test query MSE: 1.9e-1 -> 8.6e-3 | Helped |
| replace causalLM attention with prefixLM attention in a shared-layer transformer | softmax transformer ICL · Non-linear regression (N-LR) | test query MSE: 2.7e-3 -> 1.5e-4 | Helped |
| replace causalLM attention with prefixLM attention in a shared-layer transformer | softmax transformer ICL · Multiclass classification (MC) | classification error: 27.0 -> 24.1 | Helped |
| replace causalLM attention with prefixLM attention in an unshared-layer transformer | softmax transformer ICL · Linear regression (LR) | test query MSE: 1.6e-2 -> 2.5e-3 | Helped |
| replace causalLM attention with prefixLM attention in an unshared-layer transformer | softmax transformer ICL · Non-linear regression (N-LR) | test query MSE: 2.9e-3 -> 9.0e-5 | Helped |
| replace causalLM attention with prefixLM attention in an unshared-layer transformer | softmax transformer ICL · Multiclass classification (MC) | classification error: 32.1 -> 27.6 | Helped |
| replace causalLM attention with prefixLM attention | T5 DecoderOnly Base · 5-shot MMLU (57 tasks), Base | accuracy: 28.0 -> 28.8 | Helped |
| replace causalLM attention with prefixLM attention | T5 DecoderOnly Large · 5-shot MMLU (57 tasks), Large | accuracy: 26.9 -> 32.0 | Helped |
| replace causalLM attention with prefixLM attention | T5 DecoderOnly XL · 5-shot MMLU (57 tasks), XL | accuracy: 30.5 -> 39.5 | Helped |
| replace causalLM attention with prefixLM attention | T5 DecoderOnly Base · 3-shot BBH (23 tasks), Base | accuracy: 24.8 -> 27.4 | Helped |
| replace causalLM attention with prefixLM attention | T5 DecoderOnly Large · 3-shot BBH (23 tasks), Large | accuracy: 29.8 -> 32.2 | Helped |
| replace causalLM attention with prefixLM attention | T5 DecoderOnly XL · 3-shot BBH (23 tasks), XL | accuracy: 32.0 -> 35.8 | Helped |
Surprising findings
- Causal ICL can eventually converge to the optimal solution in the noiseless zero-mean setting but needs more than 100 examples to reach below 1e-2 query error, and convergence worsens with input mean shift.
Limitations
- Theory applies only to a structured one-layer LSA construction with fixed parameters, not general softmax transformers.
- Causal vs prefix comparison in LLMs relies on custom-pretrained small T5 DecoderOnly checkpoints, not state-of-the-art models.
- The argument that the gap extends to general transformers rests on experiments rather than proof.
Experimental setup
Datasets: C4 FLAN synthetic linear regression, nonlinear regression, and multiclass classification sequences
Benchmarks: MMLU BBH
Metrics: mean squared error classification error accuracy
Baselines: causalLM prefixLM PaLM2 PaLI-X
Builds on: Von Oswald et al. linear self-attention gradient descent construction Ahn et al. gradient flow LSA prefixLM (UL2/T5 style) FLAN
Closest papers by meaning
- Position: Do pretrained Transformers Learn In-Context by Gradient Descent? (ICML 2024 · similarity 0.89)
- Technical Debt in In-Context Learning: Diminishing Efficiency in Long Context (NeurIPS 2025 · similarity 0.88)
- In-Context Learning through the Bayesian Prism (ICLR 2024 · similarity 0.88)
- Sparse Autoencoders Reveal Temporal Difference Learning in Large Language Models (ICLR 2025 · similarity 0.87)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.87)
- Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions (ICLR 2024 · 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.