Research map / Large language models / Reasoning and chain of thought
Compositional Generalization from Learned Skills via CoT Training: A Theoretical and Structural Analysis for Reasoning
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.
Argues CoT training enables compositional generalization in transformers, with information-theoretic bounds decomposing error into ID/OOD KL terms and a logit-lens/causal-tracing analysis showing a two-stage circuit, validated on synthetic two-hop fact tasks.
Problem
Why and how chain-of-thought training improves generalization to both in-distribution and out-of-distribution compositional reasoning, and how this is realized internally in transformers.
Approach
Frames CoT training as teaching models to compose simpler learned skills; derives an information-theoretic generalization bound that splits into ID and OOD KL-divergence terms, and uses logit lens and causal tracing to show CoT-trained models form a two-stage circuit where intermediate results are resolved at shallower layers.
Key results
- Training with CoT reaches near-perfect ID two-hop accuracy within about 4,000 steps, whereas training without CoT shows delayed grokking on ID and no OOD generalization even after millions of steps.
- CoT-trained models resolve the bridge entity at layer 3 for ID and layer 5 for OOD, while non-CoT models resolve it at layer 5 for ID.
- Noise injected only into the second hop degrades ID and OOD accuracy, with small noise (ratio below 0.2) having little effect, while noise in both hops has a much stronger suppressive effect.
Surprising findings
- A smaller two-hop/one-hop ratio can expedite OOD generalization under CoT training.
- Scaling model size does not qualitatively change generalization; larger models only converge in fewer steps.
Limitations
- Main evidence comes from controlled synthetic fact-composition tasks with a small GPT-2-style model; real-world validation is limited in scope.
- The notion of simpler learned skills is hard to delineate in real-world tasks.
- The OOD bound relies on assumptions such as disjoint-support mixtures and linear combining rules.
Experimental setup
Datasets: synthetic atomic and two-hop fact dataset math word problems
Benchmarks: synthetic ID two-hop test synthetic OOD two-hop test math word problems
Metrics: accuracy
Baselines: training without CoT
Builds on: information-theoretic generalization bounds grokking logit lens causal tracing
Closest papers by meaning
- SCOUT: Teaching Pre-trained Language Models to Enhance Reasoning via Flow Chain-of-Thought (NeurIPS 2025 · similarity 0.91)
- Facts Don't Speak Louder than Words: The Behavioral Essence of Long-CoT Distillation (NeurIPS 2026 · similarity 0.91)
- Compositional Reasoning in Language Models under Reinforcement Learning Post-Training (NeurIPS 2026 · similarity 0.90)
- Training Nonlinear Transformers for Chain-of-Thought Inference: A Theoretical Generalization Analysis (ICLR 2025 · similarity 0.90)
- CoT Vectors: Transferring and Probing the Reasoning Mechanisms of LLMs (ICLR 2026 · similarity 0.90)
- When More is Less: Understanding Chain-of-Thought Length in 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.
