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

TL;DR

The paper investigates how Chain-of-Thought (CoT) data granularity affects language model generalization under distribution shift for compound tasks. It combines controlled theoretical analysis with empirical experiments across LIS, MPC, and ERVC tasks, showing that QA-trained models generalize poorly to OOD data despite strong IID performance, while fine-grained CoT data substantially improves OOD generalization and sample efficiency. The authors formalize compound tasks as dynamic state-transition problems, prove that CoT mitigates shortcut learning, and demonstrate that transformer positional embeddings amplify this effect via subtask recurrence. Practically, the findings guide data collection toward granular CoT annotations to achieve robust generalization in real-world deployment where distribution shifts are common, with implications for designing scalable, reliable LLM systems. The work also provides a constructive framework for analyzing and quantifying distribution shift through prefix coverage and KL-divergence bounds between training and evaluation distributions.

Abstract

Generalization to novel compound tasks under distribution shift is important for deploying transformer-based language models (LMs). This work investigates Chain-of-Thought (CoT) reasoning as a means to enhance OOD generalization. Through controlled experiments across several compound tasks, we reveal three key insights: (1) While QA-trained models achieve near-perfect in-distribution accuracy, their OOD performance degrades catastrophically, even with 10000k+ training examples; (2) the granularity of CoT data strongly correlates with generalization performance; finer-grained CoT data leads to better generalization; (3) CoT exhibits remarkable sample efficiency, matching QA performance with much less (even 80%) data. Theoretically, we demonstrate that compound tasks inherently permit shortcuts in Q-A data that misalign with true reasoning principles, while CoT forces internalization of valid dependency structures, and thus can achieve better generalization. Further, we show that transformer positional embeddings can amplify generalization by emphasizing subtask condition recurrence in long CoT sequences. Our combined theoretical and empirical analysis provides compelling evidence for CoT reasoning as a crucial training paradigm for enabling LM generalization under real-world distributional shifts for compound tasks.

Figures (9)

• Figure 1: The illustration of the impact of the granularity of Chain-of-Thought on In-Distribution (IID) and Out-of-Distribution (OOD) performance. Left: IID performance. Right: OOD performance. Results are averaged over four compound tasks. While models trained without CoT achieve high IID accuracy ( 80%), they exhibit a substantially poor generalization performance ( 10%) on OOD data.
• Figure 2: The illustration of shortcut learning, wherein the model exploits spurious correlations between questions and answers (e.g., answer length) rather than genuinely understanding the underlying reasoning, ultimately results in poor generalization ability.
• Figure 3: Chain-of-thought alleviates distribution shift by breaking down complex problems into simpler, familiar sub-problems.
• Figure 4: Outside Window Recap Condition: Recapping a token from outside the attention window, s5 depends on s1
• Figure 5: Inside Window Recap Condition: Recapping tokens to preserve causal order and skip the irrelevant token, s6 depends on s1 and s3

Theorems & Definitions (20)

• Definition 3.1
• Definition 4.1: Chain of Thought
• Proposition 4.2
• Lemma 4.3: Prefix Substructure
• Theorem 4.4: KL Divergence Reduction with Training Size
• Remark 4.5: Relationship with Training Size
• Theorem C.1: Outside Token Recap Condition under RoPE
• proof
• Theorem C.2: Inside window recap condition under Rope
• proof