Why Does CoT Reasoning Generalize? Theory Provides First OOD Guarantee

When Chain-of-Thought Steps Outside the Training Distribution: A Neglected Theoretical Gap

Chain-of-Thought (CoT) has long been a standard technique for enhancing large models' complex reasoning capabilities, yet a fundamental question has always lacked a rigorous answer: if a model only sees single-step atomic reasoning data during training, why can it complete tasks at test time that require multi-step chaining and come from a completely different distribution?

A recent paper formalizes this problem as compositional OOD generalization. The setup is clear: neither the full reasoning trajectories nor the direct input-output pairs of the target task are visible; training provides only in-distribution (ID) single-step atomic data. The paper theoretically explains the mechanism behind CoT generalization and provides a risk upper bound. The code is open-sourced on GitHub (CoT-OOD).

Compositional OOD generalization refers to the setting where, at test time, tasks are composed of subtasks seen during training in novel ways, and the overall distribution differs from the training distribution. Unlike common covariate shift, the challenge here is not in the appearance of inputs but in the reasoning structure requiring multi-step chaining. A typical example: training only sees single-step atomic operations like "add one step" or "multiply one step," while testing requires completing multi-step mixed arithmetic. This setting demands composability: the model must reuse and concatenate transition functions learned in-distribution rather than memorize full trajectories.

Core Mechanism: Repeated Composition of a Shared Single-Step Predictor

The paper's core insight is that CoT generalizes by repeatedly invoking the same shared single-step predictor, stitching together reusable ID transitions. In other words, the model does not need to learn a separate mapping for each multi-step task; it learns a general "take one step" capability and then repeats that step T times at inference.

This "sharing" is key. It means capacity control applies only to a single function class, rather than growing linearly with task length. The subsequent generalization analysis is built entirely on this structural assumption, and it is the most insightful part of the paper: CoT's capability comes not from a larger model but from a more sensible compositional structure.

A single-step predictor is a function that, given the current state (or intermediate reasoning result), predicts the next transition. In CoT, it corresponds to the ability to "generate the next reasoning step based on the existing reasoning prefix." Sharing means all multi-step tasks reuse the same function class rather than being parameterized separately for each task length or type. This parameter sharing keeps the hypothesis space capacity from growing with the number of reasoning steps T, allowing the generalization analysis to handle tasks of different lengths uniformly—and it is also the key premise for understanding the subsequent risk upper and lower bounds.

Risk Upper Bound: Coefficients Decoupled from Step Count

The theoretical result provides a set of explicit bounds. Let T be the number of reasoning steps in the target task, n the number of training samples per subtask, and k the number of output tokens per step. Then the OOD task risk is controlled by the sum of ID subtask risks, with a scaling coefficient of order O(k) that is independent of T.

The implication of this result deserves careful reading: although the sum of subtask risks may itself grow with T, the additional penalty from composition does not worsen as the reasoning chain lengthens. That is, long-chain reasoning is not inherently harder to generalize merely because "the chain is longer"; the difficulty comes from single-step risk, not from the compositional structure.

Further analysis using Rademacher complexity shows that the estimation term scales as O(√(T/n)) rather than linearly in T. This is precisely the benefit of predictor sharing under unified capacity control; without sharing, estimation error would amplify linearly with chain length.

Rademacher complexity is a tool for measuring a function class's ability to fit random noise, commonly used to derive generalization error upper bounds. For a hypothesis class composed of T steps, a naive analysis yields complexity growing linearly with T; but if each step shares the same capacity-controlled predictor, the complexity of the composed class can be controlled to be sublinear in T (e.g., √T). This explains the origin of the estimation term O(√(T/n)) in the paper: the shared structure prevents capacity from expanding with chain length, so that sample efficiency for long-chain reasoning does not deteriorate sharply.

Autoregressive Setting and Lower Bound: Generalization Is Not Free

For autoregressive next-token prediction (Next-Token Prediction), the paper shows that the final task risk consists of the optimal single-step risk multiplied by the same O(k) coefficient, plus an estimation term of order Õ(k√(T/n)).

On the other side, the paper uses a parity-coupling construction to prove that, for empirical risk minimizers of subtasks, when T≤n, the expected final task risk is at least the optimal in-class risk plus Ω(√(T/n)). This shows that the upper bound is tight in order, and that compositional generalization has an irreducible cost that does not vanish no matter how clever the method.

The parity-coupling construction is a reduction technique commonly used in lower-bound proofs: construct a pair of subtasks that differ only in some parity, making them hard for any empirical risk minimizer to distinguish, thereby forcing an unavoidable error in the final composed task. By carefully designing the subtask distribution, one can prove a lower bound of Ω(√(T/n)) on compositional generalization error. This construction shows the upper bound is tight in order and reveals that the cost of compositional generalization stems from statistical indistinguishability among subtasks rather than algorithmic deficiencies.

Experimental Validation: Short-to-Long and Compositional Generalization

The paper validates the theory on Atom-Task and GSM8K. The results support two conclusions: CoT indeed achieves short-to-long generalization and compositional generalization; and explicit state tracking is crucial for generalization. This finding has direct value for engineering practice: to make CoT work reliably, the model needs to maintain and update intermediate states during reasoning, not merely string steps together.

From a longer-term perspective, this work advances CoT from an "empirical trick" to an analyzable object. It reminds us that the reliability of compositional generalization depends on capacity control of shared structures and on state-tracking mechanisms—these two aspects deserve more investment than simply stacking more reasoning steps.