Research map / Architectures and efficiency / Attention and transformers
Looped Transformers are Better at Learning Learning Algorithms
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.
Proposes looped (weight-tied, input-injected) transformers trained from scratch to in-context learn regression function classes, matching a 12-layer GPT-2 with about 1/12 of the parameters.
Problem
Standard transformers lack an inherent iterative structure, making it hard to emulate iterative learning algorithms (e.g., gradient descent for linear regression) in-context; the paper asks whether looped transformers can do this more parameter-efficiently.
Approach
A single-layer GPT-2-style transformer applied repeatedly with input injection (Y_{t+1}=M(Y_t+P)), trained with a truncated loss over loop iterations in [b-T, b] and a schedule on b, so the fixed point stays stable beyond trained iterations.
Key results
- Looped transformer (0.79M params, 1 layer looped 20-30 times) matches the 12-layer unlooped GPT-2 (9.48M params) on in-distribution linear regression, almost matching OLS.
- On sparse linear regression, the looped transformer outperforms the standard transformer and Lasso when few in-context samples are available (fewer than 10 at 40-sample comparison).
- On OpenML datasets, the looped transformer is comparable or better on 6 of 10 test datasets (e.g., dataset 737: 0.72 vs 0.656 accuracy).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a 1-layer looped transformer instead of the 12-layer standard transformer | GPT-2 decoder in-context learner · OpenML abalone (ID 720) | test accuracy: 0.626 \pm 0.008 -> 0.662 \pm 0.008 | Helped |
| Use a 1-layer looped transformer instead of the 12-layer standard transformer | GPT-2 decoder in-context learner · OpenML bank8FM (ID 725) | test accuracy: 0.511 \pm 0.007 -> 0.504 \pm 0.008 | Hurt |
| Use a 1-layer looped transformer instead of the 12-layer standard transformer | GPT-2 decoder in-context learner · OpenML space_ga (ID 737) | test accuracy: 0.656 \pm 0.006 -> 0.72 \pm 0.01 | Helped |
| Use a 1-layer looped transformer instead of the 12-layer standard transformer | GPT-2 decoder in-context learner · OpenML delta_ailerons (ID 803) | test accuracy: 0.394 \pm 0.01 -> 0.40 \pm 0.01 | No clear effect |
| Use a 1-layer looped transformer instead of the 12-layer standard transformer | GPT-2 decoder in-context learner · OpenML kin8nm (ID 807) | test accuracy: 0.405 \pm 0.004 -> 0.416 \pm 0.005 | Helped |
| Use a 1-layer looped transformer instead of the 12-layer standard transformer | GPT-2 decoder in-context learner · OpenML puma8NH (ID 816) | test accuracy: 0.463 \pm 0.004 -> 0.462 \pm 0.004 | No clear effect |
| Use a 1-layer looped transformer instead of the 12-layer standard transformer | GPT-2 decoder in-context learner · OpenML delta_elevators (ID 819) | test accuracy: 0.483 \pm 0.005 -> 0.568 \pm 0.01 | Helped |
| Use a 1-layer looped transformer instead of the 12-layer standard transformer | GPT-2 decoder in-context learner · OpenML wind (ID 847) | test accuracy: 0.668 \pm 0.007 -> 0.757 \pm 0.006 | Helped |
| Use a 1-layer looped transformer instead of the 12-layer standard transformer | GPT-2 decoder in-context learner · OpenML pollen (ID 871) | test accuracy: 0.532 \pm 0.004 -> 0.51 \pm 0.005 | Hurt |
| Use a 1-layer looped transformer instead of the 12-layer standard transformer | GPT-2 decoder in-context learner · OpenML compas-two-years (ID 42192) | test accuracy: 0.65 \pm 0.005 -> 0.65 \pm 0.008 | No clear effect |
| Use the looped transformer rather than the standard transformer | GPT-2 decoder in-context learner · Sparse linear regression, fewer than 10 in-context samples | squared error | Helped |
| Use input injection in each loop iteration rather than weight tying without input injection (ablation) | Looped Transformer · Linear regression | test error | Helped |
Surprising findings
- The looped transformer with a 2-layer ReLU NN task required only b=12 loop iterations, fewer than for simpler linear regression, which the authors could not explain.
- Without input injection, weight-tied looped transformers diverge beyond the trained loop iterations.
Limitations
- Looped solutions have an inherent error floor and generalize poorly to out-of-distribution inputs, e.g., scaled inputs.
- Training with larger T can be unstable without regularization, and the choice of b and T requires hyperparameter tuning.
Experimental setup
Datasets: synthetic linear regression prompts sparse linear prompts decision tree prompts 2-layer ReLU NN prompts OpenML datasets (10)
Benchmarks: Synthetic linear regression Sparse linear regression Decision tree 2-layer ReLU NN OpenML classification datasets
Metrics: normalized squared error test accuracy
Baselines: Unlooped GPT-2 transformer (12 layers) Ordinary least squares Lasso Weight-tied looped transformer without input injection
Builds on: Garg et al. in-context function learning GPT-2 truncated backpropagation through time deep equilibrium models
Closest papers by meaning
- Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? (ICML 2024 · similarity 0.88)
- Looped Transformers for Length Generalization (ICLR 2025 · similarity 0.86)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.86)
- Stability and Generalization in Looped Transformers (NeurIPS 2026 · similarity 0.85)
- Transformers are almost optimal metalearners for linear classification (NeurIPS 2025 · similarity 0.85)
- Linear Transformers are Versatile In-Context Learners (NeurIPS 2024 · similarity 0.85)
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.