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

• The work targets the unresolved questions of whether CoT improves ID and OOD reasoning generalization and how that capability appears in internal representations.Source: §1 · p. 2
• The paper proposes that models generalize to novel, complex problems by systematically combining simpler skills learned during CoT training.Source: §1 · p. 2
• With sufficient training, ID error approaches zero regardless of CoT, while non-CoT models fail on OOD patterns and CoT-trained models achieve near-perfect OOD performance using learned simpler skills.Source: §1 · p. 2
• CoT training accelerates convergence and improves generalization across ID and OOD settings; real-world validation finds learned skills and compositional patterns applicable to complex math word problems.Source: §1 · p. 2
• CoT training forms a two-stage compositional circuit, with each explicit reasoning step mapped to a stage and intermediate results available at a lower layer than in non-CoT models.The cited examples identify layer 3 for CoT and layer 5 for non-CoT on ID examples.Source: §1 · p. 2

2.1 PRELIMINARIES

• A CoT is a sequence of reasoning states; test compositions shared with training are ID, while novel unseen compositions are OOD.Source: §2.1 · p. 3
• Expected generalization error is the difference between population risk and empirical risk, which reflects performance on training data.Source: §2.1 · p. 3

2.2 INFORMATION-THEORETIC GENERALIZATION BOUNDS

• Theorem 1 bounds expected generalization error using a weighted combination of KL divergences between training and ID/OOD test distributions, with α denoting the OOD mixture weight.Source: §2.2 · p. 3
• For matched ID compositions, error approaches zero with sufficient training whether or not CoT is used; for novel OOD compositions, non-CoT models struggle while CoT-trained models can generalize near-perfectly by applying learned simpler skills.Source: §2.2 · p. 4
• Adding noise to training data raises the generalization-error upper bounds for both ID and OOD, and larger noise ratios produce higher bounds.Source: §2.2 · p. 4

3.1 GENERAL SETUP

• Atomic facts represent simpler learned skills, while multi-hop facts combine them into complex problems; novel relation compositions define OOD evaluation.The setup uses relation patterns to simulate compositional patterns.Source: §3.1 · pp. 4–5

3.2 COT TRAINING BOOSTS BOTH ID AND OOD GENERALIZATION

• CoT training accelerates convergence and reaches near-perfect OOD accuracy, while training without CoT generalizes only to ID examples after extensive overtraining and shows no OOD generalization even after millions of steps.The comparison uses two-hop reasoning with λ = 7.2.Source: Figure 1 · p. 6, §3.2 · p. 6
• A smaller λ can speed OOD generalization, while larger models converge in fewer steps without qualitatively changing generalization.The paper also notes that overly small λ may impair learning relevant reasoning patterns.Source: §3.2 · p. 6

3.3 MECHANISMS OF COT TRAINING: INTERNALIZING STEP-WISE REASONING

• CoT training produces a two-stage circuit matching the explicit reasoning steps: it stores the bridge entity e2 at an intermediate layer, then uses e2 to derive e3.Without CoT, the circuit emerges only for ID generalization and resolves e2 at layer 5; with CoT, the intermediate layer is 3 for ID and 5 for OOD.Source: §3.3 · pp. 7–8, Figure 2 · p. 7

3.4 ROBUST SYSTEMATIC GENERALIZATION THROUGH COT TRAINING

• With noise in the second hop, CoT still supports systematic generalization, but final ID and OOD performance decline as noise increases; below ξ = 0.2, performance remains largely unaffected.Noise in both hops suppresses generalization more strongly than noise in only the second hop.Source: §3.4 · pp. 8–9, Figure 3 · p. 8

4 MORE DISCUSSION

• CoT training optimizes predictions of both reasoning steps and final answers, whereas non-CoT training optimizes only final-answer predictions.The paper frames this as a potentially finer-grained objective when snowballing errors remain acceptable.Source: §4 · p. 9
• With CoT, e2 is resolved at layer 3 for ID and layer 5 for OOD; without CoT, the OOD circuit fails to emerge and ID resolution requires layer 5 or later.The authors suggest earlier resolution leaves more layers available for the second hop, while noting that traditional transformers cannot process subtasks in parallel.Source: §4 · p. 9
• With twice the layers and weight reuse, a CoT-trained model correctly predicts e3 from the modified input format “e1, r1, r2, [mask]”.The setup is intended to emulate multiple forward passes used in next-token prediction.Source: §4 · p. 9

5 MORE RELATED WORK

• Prior analyses connect CoT to increased transformer expressivity and effective depth, while CoT training’s specific advantages remain an open question addressed by this study.The cited effective-depth account describes generated outputs being fed back into the input.Source: §5 · p. 10
• Studies of latent reasoning include recoverable intermediate variables and sequential pathways, but focus on CoT prompting rather than CoT training.Other cited work internalizes CoT through knowledge distillation or curricula that gradually shorten CoT, without fully analyzing CoT-training mechanisms.Source: §5 · p. 10
• The paper positions its contribution as analyzing the mechanisms of explicit CoT training, beyond prior work centered on training methods or algorithms.This distinction is made in the discussion of loop transformers and related approaches.Source: §5 · p. 10

6 CONCLUSION AND LIMITATION

• CoT training fosters models’ compositional generalization by combining previously acquired skills, with staged circuits reflecting how models learn “how to think.”The conclusion argues that CoT data quality depends on clarity and decomposition of reasoning steps, not only answer correctness.Source: §6 · p. 10

A MORE RELATED WORK

• Prior work studies latent reasoning pathways and internalized CoT, while loop transformers [24] [8] [15] share similarities with the paper’s two-stage circuit.Related work also notes that autoregressive generation can make human-like planning and search difficult on more complex problems [27] [113] [121] [28] [20] [88] [49].Source: A MORE RELATED WORK · p. 20

B.1 LEMMA 1 AND PROOF

• The KL-divergence decomposition separates ID/OOD uncertainty, represented by H(α), from deviations of those distributions from the training distribution.The proof uses the fact that the test distribution is a mixture whose expectation can be split into two parts.Source: §B.1 · p. 21

B.2 EXPECTED GENERALIZATION ERROR

• The expected-generalization-error framework models learning as an algorithm mapping a training sample to a hypothesis evaluated through population and empirical risk.It uses information-theoretic bounds that depend on both the data distribution and the learning algorithm.Source: §B.2 · p. 21

B.3 PROOF OF THEOREM 1

• The proof applies Lemma 3 to the generalization error and uses the assumption that l(w, Z) is R-subgaussian for all w to justify the bound through subgaussianity of g(Z).Source: §B.3 · p. 22

B.4 EXTENSIONS OF THEOREM 1

• With sufficient ID training, the expected generalization error is dominated by the OOD component, which scales with the OOD mixture coefficient and the ID-to-training KL divergence.The text states that ID patterns match training patterns and the ID divergence approaches zero.Source: §B.4 · p. 23
• Without CoT, unseen OOD compositional patterns remain difficult because the model does not explicitly learn the reasoning-step distribution P(C|X).The text characterizes P(C|X) as a uniform prior in this setting and says the OOD generalization upper bound is not improved.Source: §B.4 · p. 23
• With CoT, the model explicitly learns P(C|X) and P(Y|X, C), supporting near-optimal OOD generalization when trained subtasks and test reasoning steps align.The text assumes sufficient training and says OOD performance can approach optimal when intermediate reasoning outcomes align with inference needs.Source: §B.4 · p. 23
• Length generalization is strongest when training and test reasoning-step counts match; it can remain good for shorter prefix tests, while unseen longer subtasks are harder.The paper describes these conditions in terms of training step count m1 and test step count m2.Source: §B.4 · p. 23

B.5 PROOF OF THEOREM 2

• Theorem 2’s proof upper-bounds divergence between composed distributions by the reasoning-step divergence plus the expected conditional-output divergence.The bound is obtained using the log-sum inequality to simplify the exact KL expression.Source: §B.5 · p. 24

B.6 THEOREM 3 AND PROOF

• Training output noise is modeled as randomly replacing original outputs with probability ξ.The resulting distribution is denoted Pnew(Y|X).Source: §B.6 · p. 25
• The proof uses convexity of KL divergence in its second distribution to bound divergence for a mixture of distributions.The argument substitutes the noise-corrupted distribution into the earlier bound and applies the mixture inequality.Source: §B.6 · p. 25
• Theorem 3 gives an expected-generalization-error bound under the assumption that the loss is R-sub-Gaussian.The noise ratio ξ and training-set size N are among the bound’s described quantities.Source: §B.6 · p. 25

C EVALUATION AND REPRESENTATIVE EXAMPLES

• ID evaluation tests unseen two-hop facts with training-seen patterns, while OOD evaluation tests whether learned patterns transfer across knowledge distributions.Strong OOD performance is presented as evidence of learning latent patterns rather than only memorizing or fitting training patterns.Source: C EVALUATION AND REPRESENTATIVE EXAMPLES · p. 26

D.1 DATA GENERATION PROCESS

• The synthetic graph uses 200 relations, links each entity to 20 unique relations, and partitions one-hop facts into 95% ID and 5% OOD sets.Two-hop facts are inferred by composing pairs of one-hop relations.Source: §D.1 · p. 26

D.2 TRAINING DETAILS OF CONTROLLABLE DATA

• The experiments use a decoder-only GPT-2-style transformer, with two-layer models used for Figures 3 and 4 and unique entity and relation tokens by default.The appendix also reports AdamW optimization with learning rate 10^-4 and batch size 512.Source: §D.2 · pp. 26–27

D.3 TWO-LAYER MODEL CIRCUIT

• A two-layer model can learn a two-stage compositional circuit for two-hop facts, with the first stage’s output autoregressively reused for the second.Logit lens interprets individual states, while causal tracing measures connection strength between states.Source: Figure 4 · p. 27

D.4 TWO-HOP TO MUTI-HOP

• When three-hop facts are absent, splitting a three-hop fact into two two-hop test facts can still be solved if both predictions are correct.The paper describes this as reproducing reasoning patterns that appear in training, consistent with [8].Source: §D.4 · p. 28
• Training only on two-hop facts does not directly generalize to three-hop queries, while adding three-hop facts yields over 0.9 accuracy on both ID and OOD tests.The three-hop training set uses a three-hop-to-one-hop ratio of 3.6; the analysis also reports a three-stage mechanism.Source: §D.4 · p. 28
• The three-stage circuit activates e1 and r1, then r2 and e2, then r3 and e3; aligned intermediate reasoning patterns are associated with stronger generalization.The authors connect this alignment insight to the small amount of long CoT data used in DeepSeek-R1’s Cold Start phase [25].Source: §D.4 · p. 28

D.5 CAUSALITY ANALYSIS DETAILS

• Causal tracing captures normal and perturbed activations, swaps a selected state into the normal run, and prunes nodes whose intervention leaves the target token unchanged.Perturbations replace an input token with a random alternative of the same type; target states are identified with the logit lens.Source: §D.5 · p. 28

D.6 COT TRAINING WITH NOISE

• At noise ratio 0.2, Table 1 S(2) IDtest accuracy is 95.1% and OOD accuracy is 77.6%, while Table 2 S(2) OOD accuracies are 11.5% and 8.4%.The comparison concerns second-hop-only noise in Table 1 versus noise in both hops in Table 2.Source: Table 2 · p. 29, Table 1 · p. 29, §D.6 · p. 30
• Adding noise to both hops suppresses generalization more than adding it only to the second hop.Figure 5 compares training and testing accuracy across ξ values 0.05, 0.1, 0.2, and 0.4 when both hops are noisy.Source: Figure 5 · p. 29, Table 2 · p. 29, Table 1 · p. 29
• CoT still enables systematic generalization with noisy training data when the noise stays within a certain range, especially at small noise ratios.The authors conclude that data quality matters more than the method itself and that some errors are acceptable.Source: §D.6 · p. 30

D.7 NOISE DISCUSSION

• As ξ increases with only the first hop noisy, performance generally declines, and S(2)_OOD falls to 0.3%.The model nevertheless retains some generalization capability under low-noise conditions.Source: §D.7 · p. 30, Table 3 · p. 30

D.8 DETAILS OF REAL-WORLD DATA VERIFICATION

• Real-world verification uses 1,319 GSM8K test samples and 100,000 MetaMathQA training samples for one epoch, with noise from skipped steps or incorrect step results.The training data is converted to Alpaca format, with reasoning steps alongside answers in the CoT format.Source: §D.8 · p. 31

D.9 LAYER-WISE GRADIENTS CHANGE

• At checkpoint 4 (4,000 optimization steps), gradients are larger across layers in the ID phase (steps 0–4) than the OOD phase (steps 5–8).The passage relates this pattern to [49], which it says indicates that slow thinking with CoT produces smaller gradients that distinguish correct from irrelevant reasoning paths.Source: Figure 6 · p. 32

E REALISTIC DATA VERIFICATION

• The GSM8K verification fine-tunes LLaMA3.1-8B and Qwen2.5-3B on MetaMathQA with LoRA rank 16, using final answers alone or reasoning steps plus answers as targets.Noise is injected only into mathematical-expression components of reasoning steps, using step skipping or incorrect step results.Source: E REALISTIC DATA VERIFICATION · p. 33
• At noise ratio 1, CoT fine-tuning reaches 68.83% accuracy versus 19.21% without CoT on GSM8K, with noise limited to mathematical expressions in reasoning steps.The comparison is reported for the right side of Table 4.Source: Table 4 · p. 33, E REALISTIC DATA VERIFICATION · p. 33
• LLaMA3.1-8B accuracy drops from 33.13 to 22.79 after answer-only fine-tuning, while CoT fine-tuning significantly improves accuracy over answer-only and pretrained baselines.The left side of Table 4 reports GSM8K performance with and without CoT fine-tuning.Source: E REALISTIC DATA VERIFICATION · p. 33, Table 4 · p. 33

E.1 MORE REALISTIC DATA RESULTS

• The two-hop evaluation uses 6,442 cases where Llama2-7B answers the full query and first hop for CoT training, and 2,320 cases where it answers each hop alone but fails the full query for testing.The dataset contains 82,020 Wikidata-based two-hop queries; filtering prompts are used to remove shortcut cases.Source: §E.1 · p. 34
• After 100 epochs of LoRA fine-tuning, the model produces “Washington.” first for a held-out query about the capital of Jenna Jameson’s country of citizenship.The authors describe this output as indicating some generalization ability from CoT training.Source: §E.1 · p. 34