Research map / Architectures and efficiency / Attention and transformers
Looped Transformers for Length Generalization
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 looped decoder-only Transformers with step-dependent end-to-end supervision (no intermediate steps) so that an adaptive number of loop iterations enables length generalization on arithmetic and algorithmic tasks.
Problem
Transformers trained from scratch solve arithmetic and algorithmic tasks in-distribution but fail to length-generalize to unseen input lengths, and fixed-depth models cannot adapt computation to problem length.
Approach
Define n-RASP-L problems (a RASP-L step applied T(n) times) and train a looped decoder-only GPT-2 block with full-output prediction, supervising the output after varying numbers of loop steps using only input-output pairs and the known step count; at inference, the number of steps is set by an oracle or a maximum-confidence rule.
Key results
- Looped model trained on parity up to 20 digits generalizes to more than 40 digits near perfectly.
- On addition and copy, where vanilla NTP fails at max training length +10, the looped model remains almost perfect.
- NTP-Loop, NTP-Pause and FOP variants improve over vanilla NTP only slightly and fall behind the adaptive looped model.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use looped training with adaptive inference steps instead of NTP | GPT-2 on Parity, Copy, Addition, Binary Sum, Multiplication, and Unique Set · length generalization across studied tasks | exact match accuracy | Helped |
| use looped training with adaptive inference steps instead of NTP | GPT-2 on Parity · Parity, test length 40 | exact match accuracy | Helped |
| use looped training with adaptive inference steps instead of NTP | GPT-2 on Addition and Copy · Addition and Copy length generalization | exact match accuracy | Helped |
| use weight-tied layers with fixed loop count in NTP | GPT-2 on Addition, Copy, and Multiplication · length generalization | exact match accuracy | Helped |
| add pause tokens to NTP | GPT-2 on Binary Sum and Unique Set · length generalization | exact match accuracy | Helped |
| add input injection at each loop iteration (ablation) | FOP-Loop-Adaptive on Addition, Binary Sum, and Unique Set · length generalization | exact match accuracy | Helped |
| replace adaptive loop steps with pause tokens (ablation) | FOP on studied tasks · length generalization | exact match accuracy | Hurt |
| remove weight tying and adaptive depth (ablation) | FOP on Addition · Addition | in-distribution exact match accuracy | Hurt |
| use a fixed number of loop steps instead of variable depth (ablation) | FOP-Loop-Fixed · length generalization across studied tasks | exact match accuracy | Hurt |
| choose stopping steps per sample instead of for the full test set (ablation) | FOP-Looped-Adaptive-Instance · Addition, Copy, Multiplication, and Unique Set | exact match accuracy | No clear effect |
| choose stopping steps per sample instead of for the full test set (ablation) | FOP-Looped-Adaptive-Instance · Parity and Binary Sum | exact match accuracy | Hurt |
Surprising findings
- Fixed-depth FOP variants with pause tokens show worse in-distribution accuracy on addition, attributed to difficulty training a 20x-deep model.
Limitations
- Definition of n-RASP-L does not cover tasks requiring multiple sequential loops.
- Training with direct looping is computationally demanding for large step counts; only limited steps and lengths were trained due to compute.
- Requires ground-truth number of steps in training data, an extra assumption beyond standard end-to-end training.
Experimental setup
Datasets: synthetically generated algorithmic tasks: parity, copy, binary addition, binary sum, binary multiplication, unique set
Benchmarks: parity copy binary addition binary sum binary multiplication unique set
Metrics: exact match accuracy standard error over five seeds
Baselines: NTP NTP-Pause NTP-Loop FOP FOP-Pause FOP-Loop-Fixed PonderNet
Builds on: Universal Transformers RASP-L NoPE GPT-2
Closest papers by meaning
- What Algorithms can Transformers Learn? A Study in Length Generalization (ICLR 2024 · similarity 0.92)
- Arithmetic Transformers Can Length-Generalize in Both Operand Length and Count (ICLR 2025 · similarity 0.92)
- A Formal Framework for Understanding Length Generalization in Transformers (ICLR 2025 · similarity 0.90)
- Length Generalization for Transformers via Compression (NeurIPS 2026 · similarity 0.89)
- Universal Length Generalization with Turing Programs (ICML 2025 · similarity 0.88)
- Position Coupling: Improving Length Generalization of Arithmetic Transformers Using Task Structure (NeurIPS 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.