Research map / Large language models / Reasoning and chain of thought
Iteration Head: A Mechanistic Study of Chain-of-Thought
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.
Studies how chain-of-thought emerges in small two-layer transformers trained on synthetic iterative tasks (copy, parity, polynomial iteration), identifying an "iteration head" attention circuit and showing skill transfer via pretraining on another iterative task.
Problem
Why and how transformers trained with next-token prediction acquire chain-of-thought abilities for multistep iterative reasoning remains poorly understood, and single-token prediction is limited for such tasks.
Approach
Proposes an "iteration head": a two-layer attention circuit in which the first head locates the end-of-input token and the second retrieves the t-th input token by positional subtraction, letting the MLP compute the successor state from generated CoT tokens; shows it emerges in small transformers and can be transferred via data curation (pretraining on polynomial iteration speeds parity learning).
Key results
- Two-layer transformers with CoT learn the parity and polynomial iteration tasks, with attention patterns matching the iteration head (EoI-focused first layer, position-t-focused second layer).
- Next-token prediction without CoT and one-layer CoT models perform clearly worse on the polynomial iteration task across embedding dimensions and sequence lengths.
- Pretraining on polynomial iteration for 200 epochs before switching to parity reduces the total epochs needed for parity to under 300 versus 1000 epochs from scratch.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Pretrain on polynomial iteration for 200 epochs, then train on parity | Two-layer transformer · Parity | epochs to learn: less than 300 epochs vs 1000 epochs | Helped |
| Fine-tune only the second-layer MLP on parity after training on copying | Two-layer transformer with one attention head per layer · Parity | accuracy: 100% | Helped |
| Use chain-of-thought instead of next-token prediction (ablation) | Two-layer transformer · Polynomial iteration, P(X, Y) = XY + 1 in F_11 | test accuracy | Helped |
| Use a two-layer transformer instead of a one-layer transformer (ablation) | Transformer trained with CoT · Polynomial iteration, P(X, Y) = XY + 1 in F_11 | test accuracy | Helped |
| Remove CoT and use next-token prediction (ablation) | Transformer · Parity, sequence lengths up to L_max = 32 and embedding dimension d = 32 | test accuracy | Helped |
| Use a one-layer transformer with CoT instead of the iteration-head two-layer architecture (ablation) | Transformer · Parity | test accuracy | No clear effect |
| Reduce the embedding dimension (ablation) | Two-layer transformer · Polynomial iteration | test accuracy | Hurt |
| Patch in the ideal attention maps (ablation) | Transformer implementing an iteration head · Iterative tasks | accuracy | No clear effect |
| Remove the first attention head's focus on EoI (ablation) | Transformer implementing an iteration head · Iterative tasks | accuracy | Hurt |
| Remove the second attention head's focus on p_t (ablation) | Transformer implementing an iteration head · Iterative tasks | accuracy | Hurt |
| Freeze the positional embedding (ablation) | Transformer trained on parity · Parity | iteration-head learning | No clear effect |
| Train on parity alone instead of first training on polynomial iteration (ablation) | Two-layer transformer · Parity | epochs to learn: less than 300 epochs vs 1000 epochs | Hurt |
Surprising findings
- For small embedding dimensions and short-enough sequence lengths on parity, next-token prediction outperformed CoT, because position subtraction is hard while addition of up to 32 bits is easy for a two-layer model.
- Some runs found alternative circuits that subsample even positions via a previous-token copy instead of the full iteration head.
Limitations
- Results are on small synthetic tasks and a two-layer, one-head-per-layer architecture, not large LLMs.
- Attention analysis relies on a heuristic peakiness score with a single seed per configuration.
- Position subtraction becomes difficult at low embedding dimension, so the iteration head does not always appear.
Experimental setup
Datasets: synthetic binary copy synthetic parity synthetic polynomial iteration XY+1 in F_11
Benchmarks: synthetic copy, parity, and polynomial iteration test sets
Metrics: test accuracy attention peakiness score
Baselines: next-token prediction without CoT single-layer transformer with CoT
Builds on: transformer Adam
Closest papers by meaning
- SCOUT: Teaching Pre-trained Language Models to Enhance Reasoning via Flow Chain-of-Thought (NeurIPS 2025 · similarity 0.90)
- Unlocking the Capabilities of Thought: A Reasoning Boundary Framework to Quantify and Optimize Chain-of-Thought (NeurIPS 2024 · similarity 0.90)
- The CoT Encyclopedia: Analyzing, Predicting, and Controlling how a Reasoning Model will Think (ICLR 2026 · similarity 0.90)
- Understanding Chain-of-Thought in LLMs through Information Theory (ICML 2025 · similarity 0.90)
- Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis (ICLR 2025 · similarity 0.89)
- Enhancing Chain of Thought Prompting in Large Language Models via Reasoning Patterns (AAAI 2025 · similarity 0.89)
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.