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

In autoregressive Chain-of-Thought (CoT) learning, a model generates a sequence of reasoning steps before producing a final answer. Because the model is trained under an exact-trace loss—where a single incorrect token invalidates the entire sequence—it is natural to wonder if the sample complexity scales poorly with the length of the reasoning chain. This paper investigates whether the autoregressive structure allows the learner to achieve the same optimal sample complexity as the local next-token rule that generates the traces.

The author proves that the sample complexity is governed by the Daniely–Shalev-Shwartz (DS) dimension of the local next-token class. To achieve this, the paper introduces a new complexity measure called "parity dimension." Unlike the standard DS dimension, which can increase during autoregressive rollouts, parity dimension is rollout-stable. The proof utilizes a low-coordinate spanning theorem on finite restrictions to bound the one-inclusion density, effectively showing that the complexity of the full trace is no greater than the complexity of the local rule.

The study establishes that the sample complexity for exact-trace learning is O((DSdim(H) + log(1/delta))/epsilon). This bound is worst-case optimal, as it recovers the standard multiclass learning rate when the rollout length is restricted to one step. The author also provides a counterexample demonstrating that the standard DS dimension can increase under rollout, necessitating the use of the parity dimension as a key invariant for the proof.

This result provides a theoretical foundation for CoT supervision, proving that "thinking" before answering does not inherently make a model harder to learn from a statistical perspective. It resolves the open question of whether CoT learning can achieve the sharp local PAC rate, outperforming previous bounds that depended on the rollout length or larger, less efficient complexity parameters.