On the Sample Complexity of Next-Token Prediction

2025-05-22

AISTATS 2025Annual Conference on Artificial Intelligence and Statistics

On the Sample Complexity of Next-Token Prediction

Oğuz Kaan Yüksel·Nicolas Flammarion
May 2025

Research context

Open in graph

Produces

Sample-complexity bounds for next-token prediction

For bounded, Lipschitz next-token predictors, empirical risk minimization attains the optimal token rate on observed contexts without a mixing-time dependence. Prediction on unseen contexts retains the same rate but incurs a factor determined by the mixing time of the data-generating process. The Markov order enters through the complexity of the hypothesis class rather than through an automatically exponential dependence.

Statement
\[\begin{aligned}\widetilde L(\widehat\theta)&\leq\widetilde{\mathcal O}\!\left(\frac{B\,C(\Theta)}{NT}\right),\\[0.35em] L(\widehat\theta)&\leq\widetilde{\mathcal O}\!\left(\frac{\tau_{\mathrm{mix}}(B+\ln d)\,C(\Theta)}{NT}\right)\end{aligned}\]
\(N\)
number of independent trajectories
\(T\)
number of predicted tokens per trajectory
\(B\)
uniform bound on the logits
\(d\)
vocabulary size
\(C(\Theta)\)
metric-entropy complexity of the hypothesis class
\(\tau_{\mathrm{mix}}\)
mixing time of the data-generating process
\(\widetilde{\mathcal O}\)
asymptotic upper-bound notation suppressing logarithmic factors in the other problem parameters and confidence level
\(\widehat\theta\)
empirical risk minimizer
\(\widetilde L\)
in-sample KL prediction error
\(L\)
out-of-sample KL prediction error