Research map / Large language models / Reasoning and chain of thought
CoTFormer: A Chain of Thought Driven Architecture with Budget-Adaptive Computation Cost at Inference
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 CoTFormer, a weight-tied Transformer whose repeated passes let each new token-level representation attend to earlier-pass representations (mimicking CoT), plus a Mixture-of-Repeats router for budget-adaptive token-wise compute at inference.
Problem
Deeper weight-tied (recurrent) Transformers such as Universal Transformers struggle when tokens exit at different depths, since deeper layers cannot cleanly attend to halted tokens; the paper also seeks compute-adaptive inference without retraining.
Approach
Treat each repeat as creating a new token-level representation that attends to the representations from all previous repeats (as CoT thought tokens do), and train a token-wise router (Mixture of Repeats) with randomly sampled per-repeat capacities so the compute budget can be chosen at inference.
Key results
- On OpenWebText2 with 12 layers and 3 repeats, CoTFormer reaches perplexity 27.07 vs 27.47 for Block Universal Transformer and 28.39 for the standard 12-layer model.
- LN-CoTFormer with 24 layers and 5 repeats (perplexity 24.11 in the depth-embedding table, 24.48 in the main table) surpasses a standard 48-layer Transformer (24.17).
- The adaptive CoTFormer reduces compute substantially without noticeable accuracy loss, and the router beats fixed-depth reduction at equal compute.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Allow later repeats to attend to representations from previous repeats, as in CoTFormer rather than Block Universal Transformer | 12-layer repeated Transformer · OpenWebText2 | perplexity: 27.74(0.01) -> 27.55(0.02); 27.47(0.01) -> 27.07(0.01); 27.15(0.02) -> 26.64(0.04) | Helped |
| Allow later repeats to attend to representations from previous repeats, as in CoTFormer rather than Block Universal Transformer | 24-layer repeated Transformer · OpenWebText2 | perplexity: 25.47(0.01) -> 25.28(0.00); 25.19(0.03) -> 24.85(0.04); 24.95(0.01) -> 24.48(0.03) | Helped |
| Use Mixture of Repeats instead of Stick Breaking halting (ablation) | Adaptive Block Universal Transformer · Adaptive Block Universal Transformer | perplexity: 33.82 -> 33.08 | Helped |
| Use Mixture of Repeats instead of PonderNet halting (ablation) | Adaptive Block Universal Transformer · Adaptive Block Universal Transformer | perplexity: 41.37 -> 33.08 | Helped |
| Add a depth embedding (ablation) | LN-CoTFormer · Adaptive repeats | perplexity: 25.08(0.03) -> 24.94(0.01) | Helped |
| Add a depth embedding (ablation) | LN-CoTFormer · Fixed 5 repeats | perplexity: 24.11(0.03) -> 24.17(0.08) | No clear effect |
| Reserve beginning and final layers outside the repeated middle layers (ablation) | CoTFormer · OpenWebText2 | accuracy | No clear effect |
| Add a layer norm after each repeated pass (ablation) | CoTFormer · OpenWebText2 | perplexity | Helped |
| Use learned router weights to allocate repeats rather than fixing every token to a smaller depth | Adaptive CoTFormer · Compute-accuracy trade-off | accuracy | Helped |
| Reduce the inference compute budget using the learned router | Adaptive CoTFormer · Inference at reduced compute | accuracy | No clear effect |
| Reduce compute beyond the range with negligible accuracy impact | Adaptive CoTFormer · Inference at very low compute | accuracy | Hurt |
| Train for 60k rather than 40k steps (ablation) | Adaptive CoTFormer · Adaptive CoTFormer, final repeat | router weight | Helped |
Surprising findings
- At the lowest compute budget (no repeats), the adaptive Block Universal Transformer outperforms the adaptive CoTFormer.
Limitations
- Adaptive CoTFormer lags a non-adaptive CoTFormer trained with fixed 5 repeats at equal compute, possibly due to fewer gradients reaching deeper repeats.
- Matching a standard deeper Transformer requires more repeats than the depth ratio; CoTFormer does not fully close the gap.
- Implementation uses a simple non-causal Flash Attention variant and is not optimized for speed.
Experimental setup
Datasets: OpenWebText2
Benchmarks: OpenWebText2
Metrics: perplexity multiply-accumulate operations
Baselines: Standard Transformer Block Universal Transformer Universal Transformer PonderNet Stick Breaking Mixture-of-Depths Fixed Depth
Builds on: Universal Transformer PonderNet Mixture-of-Depths Pre-LayerNorm Transformer Adaptive Computation Time
Closest papers by meaning
- SIM-CoT: Supervised Implicit Chain-of-Thought (ICLR 2026 · similarity 0.90)
- Markovian Transformers for Informative Language Modeling (ICLR 2026 · similarity 0.89)
- Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis (ICLR 2025 · similarity 0.89)
- Token Signature: Predicting Chain-of-Thought Gains with Token Decoding Feature in Large Language Models (ICML 2025 · similarity 0.89)
- Activation Control for Efficiently Eliciting Long Chain-of-thought Ability of Language Models (NeurIPS 2025 · similarity 0.89)
- CoT Vectors: Transferring and Probing the Reasoning Mechanisms of LLMs (ICLR 2026 · 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.