💬 NLP

Compositional Generalization from Learned Skills via CoT Training: A Theoretical and Structural Analysis for Reasoning

This paper theoretically and structurally demonstrates that Chain-of-Thought training enhances large language models' compositional generalization by decomposing generalization bounds into in-distribution and out-of-distribution components and internalizing reasoning into a two-stage circuit that teaches models how to think by combining learned skills rather than merely memorizing answers.

Original authors: Xinhao Yao, Ruifeng Ren, Yun Liao, Lizhong Ding, Yong Liu

Published 2026-02-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xinhao Yao, Ruifeng Ren, Yun Liao, Lizhong Ding, Yong Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: Teaching a Child to Cook vs. Giving Them a Recipe

Imagine you are trying to teach a robot (or a large language model) how to solve complex problems.

The Old Way (No Chain-of-Thought):
You show the robot a picture of a finished lasagna and say, "This is what a lasagna looks like. Now, make me one."
The robot looks at the picture, memorizes the pattern, and tries to copy it. If you ask for a lasagna with slightly different ingredients (like using zucchini instead of noodles), the robot gets confused. It doesn't know how to cook; it only knows what the final dish looks like. If the situation changes even a little, it fails. This is like memorizing answers.

The New Way (Chain-of-Thought or CoT):
Instead of just showing the lasagna, you show the robot the process. You say, "First, boil the noodles. Then, chop the tomatoes. Then, layer them."
You are teaching the robot the skills (boiling, chopping, layering) and how to combine them.
Now, if you ask for a lasagna with zucchini, the robot doesn't panic. It thinks: "Okay, I know how to boil noodles. I know how to layer. I just swap the noodles for zucchini and follow the same steps." It has learned how to think, not just what to think.

This paper proves that this "teaching the steps" method (CoT training) is the secret sauce that makes AI smart enough to handle new, weird, or complex problems it has never seen before.

Part 1: The Theory (Why it Works)

The authors used math to prove why this works. They looked at two types of problems:

In-Distribution (ID): Problems that look exactly like the ones the robot practiced on.
Out-of-Distribution (OOD): Problems that are new, strange, or use different ingredients than what was practiced.

The Finding:

Without CoT: The robot is great at ID problems (it memorized the pattern) but terrible at OOD problems. It's like a student who memorized the answers to a specific math test but fails if you change one number in the question.
With CoT: The robot learns to break big problems into small, simple skills. When a new problem arrives, it grabs the right "skills" from its toolbox and combines them. It's like a chef who knows how to sauté, bake, and fry; they can cook any meal, even if they've never seen that specific recipe before.

The "Noise" Test:
The researchers also tested what happens if the training data is messy (e.g., the recipe says "boil for 5 minutes" but the robot sometimes sees "boil for 50 minutes").

Result: Even with some messy data, the robot still learns the general method! As long as the errors aren't 100% wrong, the robot can still figure out the logic. This is great news because it means we don't need perfect, human-written data for every single step; we just need mostly correct steps.

Part 2: The Internal Structure (How the Robot's Brain Changes)

The authors didn't just look at the results; they looked inside the robot's brain (the neural network layers) to see how the thinking process changed.

The "Two-Stage Circuit" Analogy:
Imagine the robot's brain is a factory assembly line with 8 stations (layers).

Without CoT: The robot tries to do everything at the very end of the line. It waits until the last station to figure out the intermediate steps. It's like trying to build a house by waiting until the roof is on to figure out where the walls go. It's messy and slow.
With CoT: The robot builds a two-stage assembly line.
- Stage 1 (Early Stations): It solves the first part of the problem (e.g., "What is the capital of France?"). It writes this answer down on a sticky note (an intermediate result) early in the process.
- Stage 2 (Later Stations): It takes that sticky note and uses it to solve the second part (e.g., "Who is the president of France?").

Why this matters:
By solving the first step early, the robot frees up the later stations to focus entirely on the next step. It's like a human solving a math problem: you write down "10 + 5 = 15" on a piece of paper so you don't have to hold that number in your head while you do the next multiplication. This makes the robot faster, more accurate, and able to handle harder problems.

Part 3: The Real-World Test

Finally, they tested this on real math problems (like those found in school textbooks).

Result: When they taught the models to show their work (CoT), the models got much better at math.
The "Messy Data" Surprise: Even when they intentionally put mistakes into the training math problems (like writing "2 + 2 = 5" in the middle of a long solution), the models still learned to solve the problems correctly, as long as the mistakes weren't overwhelming.

The Takeaway

This paper tells us that to make AI smarter, we shouldn't just feed it more data or ask it for the final answer. We need to teach it how to break things down.

Don't just give the answer: Show the steps.
Don't fear mistakes: If the steps are mostly right, the AI can learn the logic even if the data is a bit noisy.
The result: An AI that can take simple skills it learned today and combine them to solve complex, brand-new problems tomorrow.

In short: CoT training turns a robot that memorizes facts into a robot that understands logic.

1. Problem Statement

While Chain-of-Thought (CoT) prompting and training have significantly improved the reasoning capabilities of Large Language Models (LLMs), the underlying mechanisms driving this improvement remain poorly understood. Specifically, two critical questions are unresolved:

Generalization Scope: Does CoT training improve reasoning generalization in both In-Distribution (ID) and Out-of-Distribution (OOD) scenarios? If so, what theoretical principles govern this ability?
Internal Mechanism: How is this generalization capability realized within the model's internal representations and computational circuits?

Existing literature suggests that without CoT, transformers often fail to generalize systematically to OOD tasks (the "compositionality gap"), relying instead on memorized input-output mappings. This paper aims to formalize why CoT training bridges this gap.

2. Methodology

The authors employ a unified approach combining Information-Theoretic Analysis and Mechanistic Interpretability (Structural Analysis).

A. Theoretical Framework

Data Distribution: The authors define data distributions based on compositional patterns. ID data shares the same compositional patterns (relations) as training data, while OOD data introduces novel compositions of learned atomic skills.
Generalization Bounds: Using information-theoretic bounds based on KL divergence ( $D_{KL}$ $D_{K L}$ ), they decompose the expected generalization error into ID and OOD components.
- They model the data generation process as $P(Y|X) = \sum_C P(Y|X,C)P(C|X)$ , where $C$ represents the reasoning chain (skills).
- Without CoT: The model fails to learn $P(C|X)$ , treating it as a uniform prior, leading to high $D_{KL}$ on OOD tasks.
- With CoT: The model explicitly learns both the distribution of reasoning steps $P(C|X)$ and the final output given steps $P(Y|X,C)$ .
Noise Analysis: Theoretical bounds are extended to include noise injection in training data to analyze robustness.

B. Structural Analysis (Mechanistic Interpretability)

Models: Experiments utilize controlled synthetic datasets (atomic facts and multi-hop reasoning) and real-world datasets (GSM8K, MetaMathQA). Models range from 2-layer to 8-layer Transformers.
Techniques:
- Logit Lens: Projects intermediate hidden states into the vocabulary space to identify when and where specific reasoning results (e.g., bridge entities) are resolved.
- Causal Tracing: Intervenes on hidden states to determine causal contributions to the final output, mapping the "circuit" of reasoning.
Experimental Setup:
- Controlled: Synthetic relational data where entities and relations are tokenized uniquely. Training involves atomic facts (1-hop) and multi-hop facts (2-hop+).
- Real-world: Fine-tuning LLaMA-3.1-8B and Qwen2.5-3B on math word problems with and without CoT annotations, including noisy reasoning steps.

3. Key Contributions

1. Theoretical: Decomposition of Generalization Error

The paper proves that CoT training fundamentally changes the generalization landscape:

ID Generalization: Both CoT and non-CoT models can achieve near-zero error on ID tasks given sufficient training (memorization of patterns).
OOD Generalization: Non-CoT models fail on OOD tasks because they cannot compose unseen patterns. CoT-trained models succeed because they learn to systematically combine simpler learned skills ( $P(C|X)$ and $P(Y|X,C)$ ).
Robustness: Theoretical bounds show that CoT training maintains generalization even with tolerable levels of noise in reasoning steps, whereas non-CoT training is highly sensitive to distribution shifts.

2. Structural: The Two-Stage Compositional Circuit

Through causal tracing and logit lens analysis, the authors identify a distinct internal architecture in CoT-trained models:

Two-Stage Circuit: CoT training induces a sparse, two-stage computational circuit corresponding to explicit reasoning steps.
- Stage 1: The model resolves the intermediate result (bridge entity) at a shallower layer (e.g., Layer 3 in an 8-layer model for ID tasks).
- Stage 2: Deeper layers process the second hop using the intermediate result.
Efficiency: Non-CoT models require deeper layers (e.g., Layer 5) to resolve intermediate results, leaving fewer layers for subsequent reasoning, which limits OOD performance. CoT training "frees up" deeper layers for complex reasoning.
Skill Composition: The circuit reflects the internalization of "how to think" (compositional reasoning) rather than just "what to think" (fixed mappings).

3. Empirical Validation

Convergence: CoT training significantly accelerates convergence and achieves near-perfect accuracy on OOD tasks, whereas non-CoT models show delayed or non-existent OOD generalization (even after millions of steps).
Noise Robustness: CoT models remain robust to noisy reasoning steps (up to ~20% noise ratio) in training data, maintaining systematic generalization.
Real-World Verification: On GSM8K and Wikidata-based multi-hop reasoning, CoT fine-tuning drastically outperforms answer-only fine-tuning, even when training data contains incorrect intermediate steps.

4. Results Summary

Aspect	Non-CoT Training	CoT Training
ID Generalization	Achieved after extensive overtraining (Grokking).	Achieved rapidly.
OOD Generalization	Fails. Relies on memorized patterns; cannot compose new relations.	Succeeds. Systematically combines learned atomic skills.
Internal Circuit	Single-stage; intermediate results resolved at deeper layers (e.g., Layer 5).	Two-stage. Intermediate results resolved at shallower layers (e.g., Layer 3), freeing depth for subsequent steps.
Noise Sensitivity	Highly sensitive; performance degrades rapidly.	Robust. Maintains generalization with low-to-moderate noise.
Real-World (GSM8K)	Fine-tuning can degrade base model performance (e.g., 33% $\to$ 22%).	Significant improvement (e.g., 33% $\to$ 74%).

5. Significance and Implications

Redefining CoT's Role: The paper posits that CoT training is not merely a data augmentation technique but a fundamental mechanism that teaches models compositional reasoning. It shifts the learning objective from matching input-output pairs to learning the distribution of reasoning steps and their composition.
Data Quality vs. Quantity: The findings suggest that for CoT training, the clarity and decomposition of reasoning steps are more critical than the absolute perfection of the data. Models can learn robust circuits even from noisy data, provided the noise ratio is within a tolerable range.
Architectural Insight: The discovery of the "two-stage circuit" offers a new lens for understanding Transformer depth. It suggests that explicit reasoning steps allow models to utilize their depth more efficiently by resolving subtasks earlier in the network.
Practical Guidance: For designing LLM training strategies (e.g., for OpenAI's O1 or DeepSeek-R1), the paper recommends prioritizing the inclusion of explicit, decomposed reasoning chains in the training objective, even if they contain minor errors, as this fosters the "how to think" capability required for robust OOD generalization.

In conclusion, this work provides a rigorous theoretical and structural explanation for why CoT training works, demonstrating that it enables models to internalize a compositional circuit that allows for systematic generalization to novel, complex problems.