The Optimal Sample Complexity of Learning Autoregressive Chain-of-Thought

2026-07-08

AI Summary

This work addresses the high sample complexity of autoregressive chain-of-thought learning, where a single error can invalidate an entire trajectory. Within the realizable PAC framework, the paper introduces the “parity dimension”—a rollout-stable refinement of the Daniely–Shalev-Shwartz (DS) dimension—that effectively mitigates the dimensional blow-up inherent in traditional DS dimension when applied to autoregressive unrolling. By integrating PAC theory, parity pseudo-cubes, and combinatorial analysis, the authors establish a sample complexity upper bound of \(O((\text{DSdim}(\mathcal{H}) + \log(1/\delta))/\varepsilon)\). This bound depends only on the complexity of local next-token prediction and is independent of the reasoning horizon; moreover, it is shown to be unimprovable in the worst case.

Abstract

We prove that, in the realizable PAC setting, the sample complexity of exact-trace learning for full autoregressive Chain-of-Thought traces is upper bounded by the standard multiclass rate of the local next-token class, where this rate is governed by the Daniely--Shalev-Shwartz dimension. Under exact-trace loss, one wrong action makes the whole trace incorrect; nevertheless, for every stopping rule $\mathtt{halt}$ and every pointwise $\mathtt{halt}$-halting local class $\mathrm{H}$, $n_{\mathrm{PAC}}^{\varepsilon,δ}(\operatorname{Roll}_{\mathtt{halt}}(\mathrm{H}))=O((\operatorname{DSdim}(\mathrm{H})+\log(1/δ))/\varepsilon)$, with no dependence on rollout length. The dependence on $\operatorname{DSdim}(\mathrm{H})$ is worst-case optimal, since one-step stopping recovers ordinary multiclass learning of $\mathrm{H}$. The proof introduces parity dimension, a rollout-stable refinement of DS dimension based on even pseudo-cubes. It controls one-inclusion density via a low-coordinate spanning theorem on finite restrictions and, unlike DS dimension itself, does not increase under autoregressive rollout. We also show why this detour is necessary: DS dimension can increase under rollout.

Problem

• sample complexity
• autoregressive Chain-of-Thought
• exact-trace learning
• PAC learning
• Daniely--Shalev-Shwartz dimension

Innovation

• sample complexity
• autoregressive Chain-of-Thought
• Daniely–Shalev-Shwartz dimension
• parity dimension
• PAC learning