Beyond In-Distribution Success: Scaling Curves of CoT Granularity for Language Model Generalization

1. Problem Statement

Transformer-based Large Language Models (LLMs) excel at in-distribution (IID) tasks but often fail to generalize to Out-of-Distribution (OOD) scenarios, particularly in compound tasks that require dynamic planning, multi-step reasoning, and adaptation to evolving contexts.

• The Core Issue: Models trained on direct Question-Answer (Q-A) pairs achieve high accuracy on training data but suffer catastrophic performance degradation when faced with distribution shifts (e.g., novel task structures or longer sequences).
• The Gap: While Chain-of-Thought (CoT) prompting is known to improve reasoning, the specific impact of CoT granularity (the level of detail in reasoning steps) and data scaling on OOD generalization remains under-explored. It is unclear if simply increasing data volume for Q-A training can overcome these generalization limits, or how the structure of CoT data influences robustness.

2. Methodology

The authors propose a controlled experimental framework using synthetic compound tasks to systematically isolate the effects of data collection strategies on generalization.

A. Task Formalization

The paper defines compound tasks as dynamic state-transition systems where subtask dependencies evolve based on previous states (unlike static sequential processing).

• Tasks Used:
  1. Longest Increasing Subsequence (LIS): Testing generalization across sequence lengths.
  2. Multi-Step Path Counting (MPC): Testing generalization under varying combinatorial constraints.
  3. Equation Restoration with Variable Computation (ERVC): Testing recovery of linear equations and variable computation under distribution shifts.
• Data Paradigms:
  • Q-A (0% CoT): Direct input-output pairs.
  • Q-CoT (Variable Granularity): Input paired with step-by-step reasoning. The authors varied the granularity by applying probabilistic dropout to CoT steps (30%, 50%, 70%, 100% retention).

B. Theoretical Framework

The authors provide theoretical proofs to explain why CoT aids generalization:

1. Recap Conditions: They establish that due to finite attention precision (RoPE), tokens beyond a certain distance become unrecallable. CoT forces the model to "recap" (copy) relevant dependencies within the attention window, ensuring information flow.
2. Prefix-Substructure Theorem: CoT sequences exhibit a property where the reasoning steps for a shorter problem are a prefix of the reasoning steps for a longer problem. This allows the model to treat OOD problems (e.g., longer sequences) as in-distribution sub-problems.
3. KL Divergence Reduction: They prove that CoT supervision reduces the KL divergence between training and evaluation distributions. By covering intermediate reasoning states, CoT transforms OOD problems into IID-like sub-problems, whereas Q-A training leaves a massive distributional gap.

C. Experimental Setup

• Models: 6-layer Transformers trained from scratch (for LIS/MPC) and fine-tuned Phi-3.5-mini (for ERVC).
• Scaling: Experiments varied dataset sizes from 1k to 10,000k (10M) samples.
• Evaluation: Measured accuracy on both IID and OOD splits, specifically analyzing the "Generalization Gap" (IID accuracy minus OOD accuracy).

3. Key Contributions

1. Identification of the Generalization Gap: Demonstrated that even with massive datasets (10M+ examples), Q-A trained models achieve near-perfect IID accuracy but fail catastrophically on OOD tasks (dropping to ~10-30% accuracy).
2. Granularity-Generalization Trade-off: Established a strong positive correlation between CoT granularity and OOD performance. Finer-grained CoT (100% step coverage) leads to significantly better generalization than coarse or partial CoT.
3. Sample Efficiency: Showed that fine-grained CoT training achieves comparable OOD generalization to Q-A training using significantly less data (up to 80% reduction in required samples).
4. Theoretical Proof of Mechanism: Proved that CoT forces the internalization of valid dependency structures and that transformer positional embeddings amplify this by emphasizing subtask condition recurrence.

4. Key Results

• OOD Performance:
  • Q-A Models: High IID accuracy (80-90%) but very low OOD accuracy (10-20%). Increasing data size to 10M did not improve OOD performance; it plateaued.
  • CoT Models (100%): Achieved high OOD accuracy (0.75–0.95) across all tasks.
  • Partial CoT: Models with 30-50% CoT retention showed degraded OOD performance compared to 100% CoT, confirming that incomplete reasoning chains weaken generalization.
• Scaling Curves:
  • CoT models showed continuous improvement in OOD accuracy as data scale increased.
  • Q-A models showed diminishing returns on OOD tasks regardless of data scale.
• Recap Condition Ablation: Removing the "recap" mechanism (copying relevant tokens to maintain attention window coverage) caused a drastic drop in both training and OOD accuracy (e.g., ERVC OOD accuracy dropped from 0.95 to 0.12), validating the theoretical necessity of explicit dependency tracking.

5. Significance and Implications

• Data Curation Strategy: The paper argues that for deploying LMs on novel, complex tasks, data quality (granularity) is more critical than data quantity. Collecting high-quality, fine-grained CoT data is a more effective strategy than scaling up Q-A datasets.
• Robustness to Distribution Shift: CoT is identified as a crucial training paradigm for enabling LMs to handle structural distributional shifts, moving beyond surface-level pattern matching to learning underlying causal dynamics.
• Theoretical Insight: The work bridges the gap between empirical CoT success and theoretical understanding, showing that CoT effectively mitigates the "shortcut learning" problem by forcing models to learn the true state-transition dynamics of a task.
• Practical Application: Provides a roadmap for industrial deployment, suggesting that even small amounts of high-quality CoT data can significantly enhance model robustness, reducing the cost and risk of deploying LMs in dynamic real-world environments (e.g., GUI automation, medical diagnosis).

Conclusion

The paper concludes that Chain-of-Thought is not just a prompting technique but a fundamental training requirement for generalization in compound tasks. By enforcing the internalization of valid dependency structures and leveraging the prefix-substructure property, fine-grained CoT training allows models to generalize to unseen scenarios where traditional Q-A training fails, regardless of dataset size.