cs.LGJul 8, 2026

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

Authors: Zhiyuan Li

Organizations: Toyota Technological Institute at Chicago

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 halt\mathtt{halt} and every pointwise halt\mathtt{halt}-halting local class H\mathrm{H}, nPACε,δ(Rollhalt(H))=O((DSdim(H)+log(1/δ))/ε)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 DSdim(H)\operatorname{DSdim}(\mathrm{H}) is worst-case optimal, since one-step stopping recovers ordinary multiclass learning of H\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.

Explore similar work

May 7, 2026cs.LG

A Theory of Online Learning with Autoregressive Chain-of-Thought Reasoning

Autoregressive generation lies at the heart of the mechanism of large language models. It can be viewed as the repeated application of a next-token generator: starting from an input string (prompt), the generator is applied for MM steps, and the last generated token is taken as the final output. [Joshi et al., 2025] proposed a PAC model for studying the learnability of the input-output maps arising from this process. We develop an online analogue of this framework, focusing on the mistake bound of learning the final output induced by an unknown next-token generator. We distinguish between two forms of feedback. In the End-to-End model, after each round the learner observes only the final token produced after MM autoregressive steps. In the Chain-of-Thought model, the learner is additionally shown the entire MM-step trajectory. Our goal is to understand how the optimal mistake bound depends on the generation horizon MM, and to what extent observing intermediate tokens can reduce this dependence. Our main results show that the online theory of autoregressive learning exhibits a qualitative picture analogous to the statistical one found by [Hanneke et al., 2026], but with a different scale of dependence on the generation horizon. In the End-to-End model, we prove a taxonomy of possible mistake-bound growth rates in the generation horizon MM: essentially any rate between constant and logarithmic can arise. We further show that this logarithmic ceiling is unavoidable. In the Chain-of-Thought model, we show that access to the full generated trajectory eliminates the dependence on MM altogether. We also analyze autoregressive linear threshold classes, and prove optimal mistake bounds, as well as a new lower bound for the statistical setting. Along the way, our results resolve several questions left open by [Joshi et al., 2025].
Ilan Doron-Arad, Idan Mehalel, Elchanan Mossel
Jun 8, 2026cs.LG

Tight Sample Complexity of Transformers

We tightly characterize the VC dimension of depth-LL Transformers with a total of WW parameters, mapping an input sequence of length TT to a single output, establishing an upper bound of O(LWlog(TW))O(L W \log (T W)) and a nearly matching lower bound of Ω(LWlog(TW/L))Ω(L W \log (T W / L)). We further tightly characterize the sample complexity of chain-of-thought learning using such a Transformer, showing teacher forcing (i.e. selecting a predictor consistent with the entire chain-of-thought on training data) learns with sample complexity O(LWlog((T+T)W))O\left(L W \log \left(\left(T+T^{\prime}\right) W\right)\right) and that any learning rule that uses chain-of-thought data requires at least Ω(LWlog((T+T)W/L))Ω\left(L W \log \left(\left(T+T^{\prime}\right) W / L\right)\right) examples, where TT is the input length and TT^{\prime} is the number of autoregressive steps.
Chenxiao Yang, Nathan Srebro, Zhiyuan Li
Aug 7, 2026cs.LG

Stochastic Autoregressive Learning

Motivated by LLMs, which generate outputs by iteratively sampling from next-token distributions, we introduce a PAC-learning model for binary stochastic autoregressive learning. This generalizes the deterministic autoregressive learning framework of Joshi et al., COLT 2025. In our model, one fixed generator assigns a Bernoulli next-token distribution to every prompt string. Starting from an input prompt, a token is sampled and appended to the prompt; the same generator is then applied again to this expanded prompt; this procedure is repeated for MM steps. Three forms of supervision are considered: base one-step samples, chain-of-thought (CoT) samples that reveal full random trajectories of length MM, and end-to-end (e2e) samples that reveal only the final token of length MM trajectories. For a generator class, we study the minimum number of samples mbase(ε),mCoT(ε),me2e(ε)m_{base}(\varepsilon),m_{CoT}(\varepsilon), m_{e2e}(\varepsilon), resp., required to learn the one-step probabilities in the base model, and the final-token probability in the CoT and e2e models, under squared loss error~ε\varepsilon. We show that stochastic autoregressive learning fundamentally differs from the deterministic theory. At scale ε\varepsilon, there is no universal comparison between the three learning tasks: both mCoT/mbasem_{CoT}/m_{base} and me2e/mCoTm_{e2e}/m_{CoT} can be made simultaneously arbitrarily larger than M/εM/\varepsilon, the natural analogue for the existing deterministic results. Nevertheless, after altering scales, for every class, CoT learning at scale ε\varepsilon is upper-bounded by base learning at scale ε/M2\varepsilon/M^2, whereas e2e learning at scale ε\varepsilon is upper-bounded, up to logarithmic factors, by (M/ε)mCoT(Θ(ε))(M/\varepsilon) m_{CoT}(Θ(\varepsilon)). These dependencies and scales are essentially tight. We complement these bounds by studying dimension dd logistic functions in our model.
Ilan Doron-Arad, Idan Mehalel, Elchanan Mossel