The Coverage Principle: How Pre-Training Enables Post-Training
Authors: Fan Chen
Audrey Huang
Noah Golowich
Sadhika Malladi
Adam Block
Jordan T. Ash
Akshay Krishnamurthy
Dylan J Foster
Venue: ICLR 2026 Oral (2026)
Keywords: language models
reinforcement learning
test-time scaling
statistical learning theory
Primary area: learning theory
TL;DR: We introduce the coverage profile, which captures the relationship between pre- and post-training performance and admits a rich statistical theory

Abstract: Language models demonstrate remarkable abilities when pre-trained on large text corpora and fine-tuned for specific tasks, but how and why pre-training shapes the success of the final model remains poorly understood. Notably, although pre-training success is often quantified by cross entropy loss, cross entropy can be poorly predictive of downstream performance. Instead, we provide a theoretical perspective on this relationship through the lens of coverage, which quantifies the probability mass the pre-trained model places on high-quality responses and which is necessary and sufficient for post-training and test-time scaling methods like Best-of-N to succeed. Our main results develop an understanding of the coverage principle, a phenomenon whereby next-token prediction implicitly optimizes toward a model with good coverage. In particular, we uncover a mechanism that explains the power of coverage in predicting downstream performance: coverage generalizes faster than cross entropy, avoiding spurious dependence on problem dependent parameters such as the sequence length. We also study practical algorithmic interventions with provable benefits for improving coverage, including (i) model/checkpoint selection procedures, (ii) gradient normalization schemes, and (iii) test-time decoding strategies.

Published as a conference paper at ICLR 2026
THE COVERAGE PRINCIPLE: How PRE- TRAINING ENABLES POST- DEST-TRAINIG
Fan Chen
Audrey Huang
Noah Golowich Microsoft Research NYC nzgêmit t. edu
MIT
UIUC andreyh5eilinniis. edu
fanchen@mit eddu
Sadhika Malladi
Adam Block
Jordan T. Ash Microsoft Research NYC ash. www.jorrddoomooofooto..oooo
Microsoft Research NYC aadiika.malladiiea...... wel.ladi@@@mail.ccom
Columbia University adam. block@columbia.ccmm edu
sadhika waalaaallamaa..mmm
Akshay Krishnamurthy Microsoft Research NYC bsshaykremicroooft.com
Dylan J. Foster
Microsoft Research NYC
salaaffostermmicroosoto com
ABSTRACT
Language models demonstrate remarkable abilities when pre-trained large text corpora and fine-tuned for specific tasks, but how and why pre-training shapes the success of the final model remains poorly understood. Notably, although pre- training success is often quantified by cross-entropy loss, cross-entropy can be a poor predictor of downstream performance. Instead, we provide theoretical perspective this relationship through the lens of coverage, which quantifies the probability mass the pre-trained model places on high-quality responses and which is necessary and sufficient for post-training and test-time scaling methods such as Best-of-N to succeed. Our main results develop an understanding of the coverage principle, phenomenon whereby next-token prediction (more gener- ally, maximum likelihood) implicitly optimizes toward model with good cover- age. In particular, we uncover a mechanism that explains the power of coverage in predicting downstream performance: coverage generalizes faster than cross- entropy, avoiding spurious dependence on problem-dependen parameters such as the sequence length. We also study practical algorithmic interventions with provable benefits for improving coverage, including (i) model/checkpoint selec- tion procedures, (ii) gradient normalization schemes, and (iii) test-time decoding strategics.
INTRODUCTION
The remarkable capabilities of language models stem from a two-stage training process: (1) large-scale pre-training via next-token prediction with the cross-entropy loss (predicting what token should follow a prefix) and (2) targeted tet-traiinng-typpcalll adapt the model specific domains and tasks. Investing more compute and data into pre-training often enables post-training to produce a stronger model, but theoretical understanding of how these stages interact is limited. Indeed, despite substantial investment into scaling pre-training (Gadre et al., 2025; Sardana et al., 2024; Hoffmann et al., 2022), several works have demonstrated that starting post-training from a better next-token predictor does not ensure stronger performance on downstream tasks (Liu et al., 2022: Zeng et al., 2025; Chen et al., 2025; Lourie et al., 2025). Motivated by this disconnect, we theoretically investigate the connection between pre-training objectives and downstream success, asking:
via reinforcement learning-to
Can we precisely characterize the relationship between the next-token prediction loss and downstream performance? What metrics are most predictive of downstream success?
Motivated by the recent interest in test-time scaling, focus attention post-training via Best-of-N (BoN) sampling or reinforcement learning with verifiable rewards. For a prompt x, Best- of-N draws N responses y from the model and returns the best response according to task-specific

---

Published as a conference paper at ICLR 2026
KL
RL
KL
0.0
0.5
1.0
1,5
0.0
0.5
1.0
1.5
0.0
0.5
1.0
1.5
1.0
1.00
1.00
1.5
Cakeo
Passig
0.98-
10
0.96
0.2
0.5
0.8-
Plss 84
Selected model
Selecced model Mode min KL.
Panage
Selected model
Pasaga2
wi
0.92
Model min.
0.0
10k
20k
30k
40k
50k
0.0
0.2
0.4
0.6
0.8
0.2
0.05 CoYS
0.i0
Iteratioo
CoV2o
Cova.3
Figure 1: The coverage profile predicts Pass better than KL divergence. We train models in a graph reasoning task and record KL divergence, coverage profile (both measured w.r.t To). and Pass@N performance; see Appendix D for details. Left: Convergence of coverage and KL divergence over truining, showing that KL improves monotonically but coverage can degrade with training. Right: Scatter plots of KI. (top axis), C/v (lower axis) and Pass@N of checkpoints. Although KI. and Cov n exhibit comparable predictive power for small N, Cov is a better predictor for large N. Also visualized are checkpoints selected via the tournament procedure of Eq. (14) (marked Q) and by minimizing KL (marked red), demonstrating that the former selects better models for P.ss@NN
reward. Several prior works have demonstrated that the performance of BoN is strongly indicative of how well the model will perform after post-training via reinforcement learning (Yue et al., 2025; Wu et al.. 2025).
Our starting point is the observation that cross-entropy alone cannot provide meaningful answers the questions above; see Figure/ 1, which illustrates that cross-entropy can be anti-correlated with BoN performance, echoing Chen al. (2025). Instead, we show that the missing link is the coverage profile, a refinement of cross-entropy that explicitly quantifies the model's ability to assign sufficient probability to rare but high-quality responses.
Definition 1.1 (Coverage profile). The coverage profile of a model के for a distribution π
( x))
Cov (T [t) == P Pr-ny~rt (1m)
2 N
(1)
7(y x
where 2 1 is the number of Best-of-N sampling attempts.
Here, is the full response when prompted with represents the pre-training data distribution, which we presuppose covers downstream tasks of interest, and स is the pre-traines model. We prove that good coverage profile is necessary and sufficient for Best-of-N to succeed (see Section 2, as well as Propositions F.6 and F.7). This is highlighted in Figure 1, where we find that the coverage profile is correlated with downstream performance for Best-of-N (which is exactly Pass N), even when cross-entropy is not.' Motivated by this characterization of BoN performance, we ask: When, and through what mechanism, does next-token prediction produce a model ã with good coverage?
1.1 CONTRIBUTIONS
We develop theoretical understanding of the coverage principle, whereby next-token prediction implicitly optimizes toward a model with good coverage, inheriting the training corpus' coverage over tasks of interest.
Cross-entropy: Scaling laws and limitations (Section 3). We begin by deriving provable scaling laws that link cross-ente cross-entropy- controyy-specfficlllll specifically, a certain sequence-level notion to coverage and hence downstream performance, but show that cross-entropy can be sensitive to sequence length and other problem parameters, leading to vacuous predictions; this motivates our main results.
Next-token prediction implicitly optimizes coverage (Section 4). The first of our main theoreti- cal results (Theorem 4. 1) is generalization analysis for next-token prediction (more generally, maximum likelihood) that exploits the unique structure of the logarithmic loss to show that coverage can generalize faster than cross-entropy; we refer to this the coverage principle. Concretely,
Formally, the coverage profile refines cross-entropy/KL. divergence; the former the cumulative distribu- x(vlx) while KI. divergence the mean; Remark F.I. tion function (CDF) of the log density ratio log X[y[a)"
2

---

Published as a conference paper at ICLR 2026
analysis shows that the coverage profile for models learned with next-token prediction (i) avoids spurious dependence on beblle-dependem parameters such as sequence length (in contrast to cross- entropy), and (ii) converges faster still as the tail parameter N is increased. Our seillysiswhichh similar in spirit to Mendelson's 's small ball method (Mendelson, 2014; 2017)-can be viewed as giving new; fine-grained understanding of maximum likelihood.
Stochastic gradient descent through the lens of coverage (Section 5). The preceding results apply to general model classes II, but consider the empirical maximizer of the next-token prediction (maximum likelihood) objective, in the vein of classical techniques in learning theory. For the second of our main results, we focus on a specific model class oeerssrorrrrraeteriieeee autoregressive linear models (3)-but but take a more realistic approach and analyze stochastic gradient descent (SGD) on the next-token prediction objective, in the one-pass ("computc-optimal)) regime. We show that while SGD provably optimizes the coverage profile, it experiences suboptimal dependence the sequence length H. We then show that gradient normalization (which is loosely connected Adam-like updates (Bernstein & Newhouse, 2024)) provably improves coverage, removing dependence on the sequence length.
Interventions for better coverage (Section 6). Finally, we look beyond standard next-token prediction and explore families of new interventions aimed at improving coverage in theory. (i) Test-time (Section 6.1). We show that for standard token-level SGD, a decoding strategy inspired based on test-time training (Krause et al., 2019; Sun et al., 2024; Akyürek et al., 2025) provably improves coverage.
(ii) Model/checkpoint selection (Section 6.2). For selecting the best model (or checkpoint) from small number of candidates, we give lournament procedures that enjoy significantly better coverage profile (particularly with respect to the tail parameter N) than naive validation with cross-entropy.
Additional results (Appendix G). Beyond the results above, we show that: (1) MLE can find models with low coverage even in the presence of severe misspecification; (2) coverage can general- ize better under additional structural properties of the model class such convexity (Appendix G.2).
In summary, we believe that coverage offers a new perspective on the connection between pre- training objectives and downstream post-training success. Our results demonstrate that this per- spective is mathematically rich and fundamental, opening the door to a deeper understanding, cf. Appendix A.
PROBLEM SETUP
We now introduce the formal problem setup for the remainder of the paper.
Next-token prediction and maximum likelihood. We work in the following setting, which sub-
sumes next-token prediction: X' is the prompt space, y is the response space, and To X - A())
is the data distribution. We are given a dataset D = {x",, ,o',,")) H where x'' μ and y nol. x').
We consider the maximum likelihood objective
[n( = Σ log ( xx))
(2)
and refer to मे := arg max (x) as the maximum likelihood estimator for a user-specified model class II. This is a generalization of the next-token prediction, where y = y token sequence and ( | x) = I1 m x, v1 is explicitly autoregressive, so that L,( = El= D log (vi, xi 1)) We specialize next-token prediction at certain points but otherwise focus on the general setting. We make the following realizability assumption throughout.
Assumption 2.1 (Realizability). The data distribution Tu is realizable by some model π E II.
This formulation captures pre-training and SFT, with some caveats; see Appendix A.1.
Post-training and the coverage profile. Given a reward function rT(2, y) (0. representing success at a downstream task T, the goal is to fine-tune - reinforcement learning or test-time scaling-to to obtain near-optimal reward. We show (Propositions F-6 and F.7) that for any task-specific comparator policy कं X + A())), Best-of-N sampling with @(N) samples
3

---

Published as a conference paper at ICLR 2026
satisfies Easy (r, (T())) rr(2, (x))] x COVN R),, good coverage profile for कт sufficient for high reward. Further, while less well understood, some form of coverage is thought to be necessary for the success of post-training methods like GRPO (Yue et al., 2025).
Returning pre-training, it is clear that there is little hope that next-token prediction will produce model R with good coverage with respect to a downstream task unless the data distribution To itself has reasonable coverage with respect this task. We therefore posit that the data distribution covers such a downstream task, in the sense that it includes high-reward responses with some bounded- below probability. Since coverage satisfies a transitivity property, it follows that coverage with re- spect to implies coverage with respect to the optimal policy for the downstream task. For example, To has a 10% chance of generating correct response, and C/VN/IO R) = E, then we get 10s error. Thus, going forward, we focus on understanding when next-token prediction achieves good coverage N f) relative to the data distribution To itself, and avoid concerning ourselves with specific details of the task policy #y or the specific relationship between wy and To-
Autoregressive linear models. We analyze next-token prediction and maximum likelihood for gen- eral model classes II, but our running example throughout the paper will be the class II of autore- gressive linear models, defined by a known feature map á X x V + R For each parameter 0 E θ Idd the model ma (TT) is defined by
To U/ 2, V1:h-1) x exp( (0, Φ(x, Vist)))-
In practice, autoregressive sequence models such as those based on transformers- each token by sampling from softmax distribution whose logits are given by linear combination of learned features (Radford et al., 2019). Eq. (3) simplifies this by freezing the feature map, yet remains expressive enough to model complex non-Markovian dependencies, depending the choice of features.
Assumption 2.2. We assume S {0: :|0||| < 1} is convex, and suph 110(2, yra))ll s B.
3
CROSS-ENTROPE AND COVERAGE: SCALING LAWS AND LIMITATIONS
natural approach to understanding when next-token prediction achieves good coverage is ap- peal to cross- cross-entropy- perhaps first showing that next-token prediction achieves low cross-entropy (which is true asymptotically), and then relating cross-entropy to coverage. In this section we mo- tivate our main results by showing that while this is possible in weak sense, it does not yield predictive guarantees for downstream performance in the finite-sample regime.
-utropy--perhap
Define the sequence-level cross-entropy for के as DCE(TO || () == Em, Σ log F(yn (=
Since E biid mo [T)] = DCE(TO T). one expects that as scale up compute, number of samples n, and model capacity IT, DCE mo || R) - DCE(TD |l mp). or equivalently DKL (TD R) + O.
where (TL (] == E,,, Σi= log
Tolvn, IZ,11A 1) H(ya A1))
is the sequence-level KL divergence.
A simple scaling law for cross-entropy. We show below that if the model म has reasonable KL. divergence the data distribution, the coverage profile can be bounded:
Proposition 3.1 (KL-to-coverage; see Proposition F.1). For all N 2 e, Co ()
Diu (mo f)
[og(N/e)
Combining Proposition 3.1 with Proposition F.6 and our assumption that ao has good coverage with respect to the downstream task yields a simple "scaling law" for test-time compute with BoN:
Consider a task of interest with reward ri(x, y), and suppose the data distribution mp itself has constant probability of success (i.e., sampling y Tp(- a) with rT(x, y) 1). To achieve sub-optimality with Best-of-N, it suffices to choose the compute budget N as
DKL
N RJ exp
(4)
"See Proposition F.5 for formal results.
4

---

Published as a conference paper at ICLR 2026
20
H-8
1.00
H-8
200
H-8
H-16
H-16
H-16
H=24
Covy-16
H=
Ratio Com
150
H-24
KL
R 10
0.50
100
5
0.25
Ok
10k
20k
30k
40k
0.00
0k
10k
20k
30k
40k
0K
10k 20k
30k
40k
Iterations
Iterations
Iterations
The coverage profile avoids spurious dependence on sequence length. We train models in graph reasoning task and record their KI. divergence and coverage profile, measured w.r.t. To as we vary the problem horizon (sequence length); see Appendix D for details. Left: Convergence of KL over training for three horizons H, demonstrating that KL. at convergence scales linearly in the horizon H. Center: Convergence of CovN over training, manifesting no dependence on H at con- vergence. Right: Ratio of KL over CovN, showing that Proposition 3.1 can be overly conservative.
Figure 2:
That is, for fixed model and KL-diivegeece level DKL (TD || .)) S ( || [)), Eq. (4) predicts that test-time compute should increase exponentially with the desired accuracy e.²
Insufficiency of cross-entropy. At first glance, this seems to be in line with empirical test-time scaling laws (OpenAI, 2024), but there is an issue: While token-level cross-entropy has been ob- served to be modest in contemporary language models (Kaplan et al., 2020; Hoffmann et al., 2022; Xia al., 2022), the sequence-leve cross-entropy (and KL-divergence) generally grows with the length H of the sequence, so that Eq. (4) predicts exponential test-time scaling in the sequence length. Moreover, such a law cannot hold if we only assume token-level cross-entropy is bounded; Proposition F.7.
this the end the story? On the one hand, it is simple to show (Proposition F.2) that Proposi- tion 3.1 is tight for worst-case pair of models. Moreover, even for the autoregressive linear model in Eq- (3), sequence-level KL divergence scales linearly with the sequence length H, as shown in the next result.
Proposition 3.2. Fix H E and d = 1. There exists 6 X' V - [-1,1] and induced autoregressive linear class 11 with parameter space θ = 1- 1]. distribution μ over X, such that for any proper estimator मे = #(D) E 11. there exists data distribution To E 11 such that w.p. at least 0.25, DKL(mL ( 2
H
In
This behavior is reflected empirically in Figure for graph reasoning task. Yet, for this task, find (Figure 2) that in spite of large cesss-entopp/KL, next-token prediction learns model f with good coverage profile across a range of sequence lengths and that downstream Best-of-N succeeds. Why is this happening? In light of the discussion above, it must be related to specific inductive bias of the next-token prediction objective itself.
A glimmer of hope: Case study in Bernoulli models. To why large cross-entropy may be a barrier to coverage, consider perhaps the simplest setting, Bernoulli models, where X' = {1}. (0, 1}. 11 (Ber(p))pEtto. 1/2). and To Ber(p*) for some small p* (0, 1/2).
The maximum likelihood model is मे = (fr( where P is the empirical frequency of y = =1 in the dataset. We observe that with positive probability (and constant probability if n < 1/p"), the dataset D will only contain examples where O, so that the maximum likelihood model is मे = Ber(O). This implies that expected KL divergence is infinite: E[DKL (To 7)] +00. However, the cover-
age profile turns out to be well-behaved; direct caleulation shows that CO ft log(s 2) n with probability at least 1 - for all N 2 2: this gives hope that even though cross-entropy itself infinite, maximum likelihood may actually learn model with good coverage in the background. In what follows, we will show that this is not a fluke, but a general phenomenon.
Remark 3.1 (Missing mass). The underlying issue in both of the preceding examples is missing mass: there are responses that even well generalizing learner will fail to cover, and for these
Neither KI. divergence nor the coverage profile are observable quantities (though cross-entropy is es- timable upper bound on KI.). so this theoretical prediction rather than practical one as-is; Remark E.2.
5

---

Published as a conference paper at ICLR 2026
may incur large contribution the KL-divergence. More generally, KL-divergence and cross- entropy are susceptible to contributions of the scale log Wnax where Wnax = max Il (which could be large H, as in Proposition 3.2) when the model does not have enough information to generalizefexrrpollate. This phenomenon is particularly pronounced when the prompt distribution heterogeneous.
4 NEXT-TOKEN PREDICTION IMPLICITLY OPTIMIZES COVERAGE
We now present our main result (Theorem 4.1): due to the unique structure of the logarithmic loss, maximum likelihood can learn models with good coverage profile even when the cross-entropy is vacuously large. Henceforth, we abbreviate CovN(T) == COVN |] We make use of the following covering number.
Definition 4.1. For class II and a 2 O, we let No (II, a) denote the size of the smallest cover IT' such that for all π E II, there exists n' E such that suPrex. seyyloog a x)) log x (y x) S a.
Theorem 4.1 (Fast generalization for coverage). Fix N 2 8 and let c. > 0 be an absolute constant. Suppose Assumption 2. holds. With probubility least 1 ō- the maximum likelihood estimator has
log. No (11,))
Nog el1,clog N) + log(8-1)
CoV (f) x
inf <>0
(5)
log N
n
=: Ceine (II,m)
Ceoarse (TI.N,t)
Eq. (5) has a fine-grained term Crine (II, r)) and coarse-graines term Cenarso (II, N, n);; we interpret each below.
Fine-grained term. Crane (TI, n) evaluates the covering number No (IT, at a small scale (typically poly(1/n)), which matches typical bounds for conditional density estimation (e.g.. Bilodeau et al. (2023)) in KL divergence; however, unlike KL-based bounds this term has no explicit dependence
sequence length H density ratios log W The term is further scaled by log N, which implies that coverage enjoys faster convergence as we move further into the tail by increasing N; this reflects the unique structure of the logarithmic loss, and may be viewed new form of implicit bias.
Summarizing, the fine-grained term in Eq.(5) witnesses the phenomenon term the coverage principle: the coverage profile enjoys faster generalization than cross-entropy: roughly, the rate is what we would expect (via Proposition 3.1) 1) if we could somehow control KL without paying for the sequence length I/ or density ratio log W Waxx- See Appendix C for a detailed comparison to standard (asymptotic and non-asymptotic) generalization bounds for maximum likelihood based on Hellinger distance and KL-divergence.
Coarse-grained term. The coarse-graine term Cenarse (II, N. n) captures the missing mass phe- nomenon exemplified by the Bernoulli example in the prequel. This term is not explicitly normalized by 1/ log N (compared to the fine-grained term), but depends on the covering number Noc (TT, a) only at a very large scale s log N. As such, the dependence on the complexity/richness of 11 in this term vanishes as we increase N.
Overall, while the guarantee in Eq. (5) might look surprising at first glance (particularly the coarse term, as we are not aware of any existing generalization bounds with dependence on covering num- bers at such a large scale), we show in Proposition G.1 (Appendix J) that both terms are tight in general.
Overview of analysis. The proof of Theorem 4.1 is given in Appendix J (with high-level sketch in Appendix J.1). The basic idea to interpret the condition Cov (r) 2 & as a small ball-like anti-comcentration condition the vein of Mendelson (2014; 2017). That is, for models where coverage is large, the condition ( 2 witnesses a one-sided bound which implies that the empirical likelihood of π is too large with high probability, and thus π cannot be maximum -iiimuu--ikelhoood likelihood solution.
The coarse- coarse-grained term Ccoarse (II, N. n) enters because we only need to show that the coverage profile concentrates, not the log loss itself. The fine-graines term Crine (11, n) enters from one-sided concentration of the empirical likelihood, with the 1/ /og N scaling arising from the following form of implicit bias: If an example (x',y") has po(w"l=')/<<(e 2 N, this witnesses a negative contribution of order log N to the difference T.](T)) (To).

---

Published as a conference paper at ICLR 2026
Discussion. We emphasize that while covering numbers are fundamental and widely used noton of capacity in statistical learning and estimation (van de Geer, 2000; Zhang, 2002; Rakhlin & Sridharan, 2012; Bilodeau et al., 2023), they are conservative from a modern generalization perspective. Nonetheless, Theorem 4. shows that they are sufficient to capture rich aspects of generalization for coverage, and we expect that core analysis techniques can be combined with contemporary advances in generalization theory for overparameterized models (Belkin et al., 2019; Bartlett et al., 2020).
4.1
EXAMPLES
To build intuition, we analyze the behavior of Theorem 4.1 under a growth assumption on the covering number, then specialize to autoregressive linear models, showing how they exemplify the coverage principle.
Corollary 4.1. (i) Parametric regime: Suppose that there are parameters d 2 and C 2 2 such that log Noc (II,a) a) S ddlog(C/aa) for a E (0, C/2). Then for any N 2 8, with probability at least lon(Co)
1 6, COVN (7) x d[log(C/ log N)], + log +1og(1/5)
1
(ii) Nonparametric regime: Suppose that there are parameters C 2 2 and p > 0 such that lo S (C/a) for a E (0, C/2). Then for any N and n 2 log N. (C/log NyP, log(1/5)
with probability at least 5, Cov ( S log N
This result shows that for sufficiently rich classes (c.g., when > 0), the fine- fin-grained term dominates the naarse-grained term for sufficently large. On the other hand, for simple classes (e.g., when p 0), the coarse-grained term can dominate the fine-grained term
Autoregressive linear models: Low dimension. We now consider the autoregressive linear model in Eq. (3). When the dimension is small, this class satisfies log Nox (II, (IT,
dl log( (BH/o) (corresponding to the parametric regime in Corollary 4.1), and so, coverage generalizes in a (nearly) horizon-indee iitizonnindeeedeentt fashion, in stark contrast to the cross-entropy lower bound in Proposition 3.2. The only drawback (which is fundamental) is that since the class has low capacity, the coarse-grained term dominates for most parameter regimes, and the improvement as N scales is quite modest.
Autoregressive linear models: High dimension. As more interesting example, we next look at the behavior of next-token prediction for autoregressive linear models in an "overparameterized" regime where the dimension d is arbitrarily large (Zhang. 2002; Neyshabur et al., 2015; Bartlett et al., 2017). Here, control the richness of the class 11 by the norm parameter B. In this regime, it turns out that in the worst-case, the capacity log No (II, ax) scales polynomially in H. To address, this we prove refined version of Theorem 4.1 that adapts to the variance in the data distribution To, avoiding explicit dependence on sequence length.
Define the inherent variance for the data distribution
0 = E,,
1 V1.t.) - -.m (z. Bl:h -1-.-111111
(6)
where T. (x. 11:h-1) := mm..mm0 L) 31.))] | is the average feature vector given the prefix (x, V1:k-1). We can interpret the inherent variance of as a notion of effective sequence length; it captures the number tokens that are "pivotal" in the sense that they have high variation conditioned on the prefix; the name reflects a noted phenomenon in language modeling that most tokens are near- deterministic and easy to predict given their prefix, with only a few having high entropy (Abdin et al., 2024). Thus, while of can be as large as B H in the worst case, we expect to be smaller in general.
Theorem 4.2 (Overparameterized autoregressive linear models). Consider the autoregressive linear model (3), and suppose Assumptions 2.1 and 2.2 hold. For any N 2 2, next-token prediction achieves
E CoCov (7) S V n-log N + B²
Similar to Theorem 4.1, the first term in Eq. (7) be viewed "fine- giaenegained"' and the second term as "coarse- "coarse-grained" grained"; the former is typically larger, but decreases with the tail parameter N, while the latter does not decrease with N but is typically smaller begin with. We prove (details
7

---

Published as a conference paper at ICLR 2026
in Proposition K.1) that this result is tight in the sense that if H, 2 H is indeed necessary to achieve good non-trivial coverage in the overparameterized regime.
We view the introduction of the inherent variance of as instance-dependent notion of complexity for autoregressive models to be a non-trivial conceptual contribution, which may find broader use.
5 STOCHASTIC GRADIENT DESCENT THROUGH THE LENS OF COVERAGE
The coverage-based generalization guarantees for next-token prediction in the prequel apply to general model classes II, but consider the empirical maximizer म arg maxxen I,, (TT) of the next-token prediction (maximum likelihood) objective, in the vein of classical techniques in learning theory. For second set of main results, we focus autoregressive linear models (3) but take a more realistic approach and analyze stochastic gradient descent (SGD) in the single-pass regime. This setup motivated by contemporary "ccomputc-optimaa"" language model training, which typi- cally uses one or fewer passes over the training corpus (Kaplan et al., 2020; Hoffmann et al., 2022).
5.1 STOCHASTIC GRADIENT DESCENT HAS SUROPTIMAL COVERAGE
For the next-token prediction objective, single-pass stochastic gradient descent (SGD) takes the form*
0++ Proje (0' V log To (y" [x')),
(8)
for x' 1 and y' To(-1 x²). where 1j 0 is the learning rate. As the next-token prediction loss L(0)= == Eng - log Te (y x)] is convex under the parameterization (3), we can show that SGD converges To KL divergence. This implies a coverage bound, albeit a suboptimal one.
Proposition 5.1 (SGD for autoregressive linear models). Upper bound: Suppose Assumptions 2.1 and 2.2 hold. As long ll S 2HB8. it holds that EE Choosing 7 minimize this bound gives
[ Σ DKL(TD TT:)] S 7/ + 2m
E [ Et= COVN(TTO)]
S ToR (V #
(9)
Lower bound: Suppose that B c. log (TH). Then there exists an autoregressive linear class 11 such that for any constant step size 1) > 0, there exists an instance To E II with a, I 1 such that with probability at least 0.5, the SGD iterates satisfy CovN(Tg) ll mT+) 2 c m-m Tlog N for any
t [T].
The coverage bound Eq. (9) (which follows by passing from KL to coverage through Proposi- tion 3.1) is similar to Theorem 4.2, except that the second term B²H has an unfortunate dependence on the sequence length H. The lower bound shows that this dependence is tight, and SGD can indeed experience poor coverage. This failure of SGD related to heterogeneity across prompts: there are some prompts for which the effective scale of the gradient in Eq. (8) grows with H. leading to divergence unless we use a small learning rate S BB Yet for other prompts, the effective gradient range is small, leading to slow convergence (on the order of (HH) steps) unless 1) >> HB"
Remark 5.1 (Sequence-level SGD). The update in Eq. (8) can be interpreted as a "sequence- level" form of SGD, since perform single gradient step for each full sequence yf (note that V log To. (y' x²) = V log mos (y) x² vish-1)). We view this model for what is done in practice, whereby one performs SGD on sequences of tokens spamning some fixed context win- dow. While this context window may be shorter than the full training example (e.g., a long article), understanding the implications of a limited context window is beyond the scope of this work.
Ei=1
5.2 GRADIENT NORMALIZATION IMPROVES COVERAGE
To address the suboptimality of SGD, we consider gradient normalization as simple intervention. For mini-batch D [(x',,')] " of K samples from To. define the batch stochastic gradient as ((: D)
D E(EY)D V log To w;). We consider the following normalized SGD update:
0+ + Proje (0' +
3D0 x+
(10)
*Proje()) (-) denotes Fuclidean projection onto O this SGD on the loss I(() == IE[ log To [r)]
8

---

Published as a conference paper at ICLR 2026
here D' is mini-batch with K fresh samples drawn i..d. from To. and À > 0 is regularization parameter for numerical stability. We show that this update achieves horizon- indzeendend coverage bound.
Theorem 5.1. Suppose Assumption 2 and Assumption 2.2 hold. Let T, K 21,N 2 3 be given. For an appropriate choice of 17, X 0, the normalized SGD update (10) achieves the following bound:
E T Σ ( 2
of Tolog N +
B² + T
B K-log N
(11)
To achieve E[[ovN (e6ovN(T]]] (T)] S for larget level > O, it suffices to choose T = 0 log B log N + 1), giving total sample complexity n = TK = 2 log N
B E
K (
o²B
B'+
lloog
Theorem 5.1 shows that gradient normalization achieves horizon- priizon-independent coverage with a qualitatively similar rate to the guarantee for next-token prediction in Theorem 4.2: To achieve
of coverage E, both rates scale as poly ( log N B, though the dependence on E for Theorem 5.1 is worse. We view this as another instance of the coverage principle, as the rate achieved by gradient normalization goes beyond what can be achieved by passing through KI. divergence. We emphasize that minibatching alone is not enough to achieve this result; rather, minibatching is necessary to avoid excessive bias once we introduce gradient normalization.
As remark, the normalized SG update in (10) is closely related to SignSGD (Balles & Hennig, 2018) and Adam (Kingma & Ba, 2015) as shown by Bernstein & Newhouse (2024). We believe that similar coverage guarantees could potentially be shown for these methods using our techniques.
Distillation. As an additional result, we show (Theorem G.2 in Appendix G.4) that for a distillation setting, where лi corresponds to a teacher model and we have access its per-token logits, we can derive an improved gradient normalization scheme that fully closes the gap with Theorem 4.2.
6 INTERVENTIONS FOR BETTER COVERAGE
In this section, we develop new interventions that improve coverage (and downstream performance) beyond the conventional algorithms analyzed in Sections 4 and 5. We view these results as promising proofs of concept for further research into interventions driven by coverage.
6.1 IMPROVING COVERAGE AT TEST TIME
In this section, we show that a modified decoding strategy based on test-time training (or, dynamic evaluation) (Mikolov et al., 2010; Krause et al., 2018; 2019; Sun et al., 2024; Akyürek et al., 2025) leads to improved coverage when combined with token-level SGD.
We focus on autoregressive linear models, but depart from Eq. (8) by learning models with token- level SGD update, defined as
g+++1 Proje (0' + 17 log Tgn. (yh a',, vinn-))), for h 0,... H
(12)
and g++1 g++1.0 = He. for t E [T]. and where (z", (et,vi.H) ~ Tp. We will show that-when combined with test-time training-like update that performs token-level gradient updates during test time-the updates in Eq. (12) can circumvent the H -edependence in the lower bound of Proposi- tion 5.1.
Concretely, we consider a distribution T. : X + A()*)) formally introduced in Appendix K.6, which can be interpreted as an augmented version of the autoregressive linear model Ty that uses test-time training to sample. Given a prompt we first sample y1 x-(- II,, then perform gradient step 0' + Proje (6 + v log No (11 z)) to increase the probability of the token we just sampled. We then sample 1/2 To (- x, V1), update 6" + Proje, + v/V log z (Uz | xx,1)) )), and so on. Once the full sequence y1.H is sampled, we reset back (so that we can process the next test-time example). This bears similarity and specifically coincides with the method used in Krause et al. (2019); Rannen- Triki et al. (2024). We show that when augmented with this test-time sampling scheme, token-level SGD achieves horizon-indege corizon-indeeendden coverage bound that matches and even slightly improves upon the bound for next-token prediction in Theorem 4.2.
many test-time training methods in the literature,

---

Published as a conference paper at ICLR 2026
Theorem 6.1 (Token-level SGD with test-time training). Suppose Assumption 2.1 and As- sumption 2.2 hold. For a suitably chosen parameter > 0, token-level SGD (12) achieves E ll
ITT) T
+ and thus [ E COVN (To=((=(=)) (TT= 5 Tog 1
7 B (
This improves Theorem 4.2 by factor of 1/ylog N on the leading term and factor of 1/log N on the second term. Furthermore, the algorithm bypasses the lower bound on KI. divergence for proper methods in Proposition 3.2, demonstrating a provable benefit being improper.
6.2 SELECTING FOR COVERAGE
We last consider the problem of selecting model (e.g., checkpoint) from a small number of candi- dates to achieve the best coverage. We introduce tournament-lik procedure that improves upon maximum likelihood in that it removes the requirement that mi E II; guaranteed find a model in the class with good coverage if one exists, even if то itself is not in the class. As algorithmic in- tervention, we envision using this procedure to select single training checkpoint or hyperparameter configuration use for RL fine-tuning or test-time scaling. Indeed, demonstrated in Figure 1, us- ing cross-entropy as a selection criterion-ar as is standard-may result in poor coverage, while these procedures select better checkpoints. Our results here concern the general setting in Section and are not restricted to autoregressive linear models.
While their main motivation is model/checkpoint selection with finite class II, both estimators can also be applied to general, infinite classes IT. In this case, they improve upon the coverage achieved by the maximum likelihood estimator in Theorem 4.1, even in the well-specified case where To E I1; informally, the tournament estimators allow us to remove the fine-grained term in Theorem 4.1, leaving only a maase-grrained term.
A simple tournament for maximizing coverage.
Given dataset D ((''''(')))""" ve) define
C (r' [| = fi ='o''a) -67 2 |}|.
(13)
which can be interpreted as an empirical version of the coverage profile C [] in Eq. (1)
when n' Zо (see Lemma J.2). For N 1, we consider the estimator
म = arg min ma C T).
(14)
Informally, this estimator chooses the model π that minimizes the maximum coverage against any other model n' in the class II. When II is small, we can implement this tournament by simply evaluating the empirical coverage in Eq. (13) for each pair. The main guarantee for this estimator is as follows.
Theorem 6.2. Let N 2 1 be given. Then, for any a E [0, 1}, with probability at least 1 - 6, the tournament estimator (14) achieves
Cov (7) S minzen Cov + NI-a + log([11]
(15)
This shows that the tournament achieves a coverage profile nearly as good as the best-in-class, except for small polynomial blow up. in that we bound the coverage at level Ni+e in terms of the coverage for the best-in-class at level No.
Infinite class and improving the tournament. Eq. (14) can also be applied general, infinite classes II. In this case, it turns out that it improves upon the coverage achieved by the maximum likelihood estimator in Theorem 4.1 (see Theorem 6.2'). Furthermore, in Appendix G.5, we de- scribe an improved tournament estimator that is able to remove the 1 1/N term from Theorem 6.2, thereby achieving nontrivial guarantees even when the coverage parameter N is constant.
DISCUSSION AND FUTURE WORK
See Appendix for discussion and open problems, and Appendix for additional results.
ACKNOWLEDGEMENTS
We thank Clayton Sanford, Matus Telgarsky, and Nati Srebro for helpful discussions. FC acknowl- edges support from ARO through award 29-1--2-1- Simons Foundation and the NSF through awards DMS-203188 2031883 and PHY-2019786. and DARPA AIQ award.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We provide full proofs for all theoretical results in the appendix. Appendix D includes extensive experiment setup and implementation details for all empirical results. The source code is included in the supplementary material, along with the plotting scripts and data to reproduce Figure 1 and Figure 2.
REFERENCES
Marah Abdin, Jyoti Aneja, Harkirat Behl, Sébastien Bubeck, Ronen Eldan, Suriya Gunasekar, Michael Harrison, Russell Hewett, Mojan Javaheripi, Piero Kauffmann, et al. Phi-4 technical report. arXiv preprini 2rXiv::242.0086 2024.
Ekin Akyürek, Mehul Damani, Adam Zweiger, Linlu Qiu, Han Guo, Jyothish Pari, Yoon Kim, and Jacob Andreas. The surprising effectiveness of test-time training for few-shot learning. In Forty- second International Conference on Machine Learning, 2025.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 3.3, knowledge capacity scaling laws. In The Thirteenth International Conference on Learning Representations, 2025. URL ptttps://oppereviiee www.areeeeeeeeeeruuuiiiee het/forrumidd=Fxiiilggt
Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. arXiv preprint arXiv:2403.06663. 2024.
Lukas Balles and Philipp Hennig. Dissecting adam: The sign, magnitude and variance of stochastic gradients. In International Conference on Machine Learning, pp. 404-413. PMLR, 2018
Hritik Bansal, Arian Hosseini, Rishabh Agarwal, Vinh Q. Tran, and Mehran Kazemi. Smaller, weaker, yet better: Training LLM reasoners via compute-optimal sampling. In The Thirteenth International Conference on Learning Representations, 2025. URL forum?)d=33yyxFQuu1..
Peter L. Bartlett and Andrea Montanari. Deep learning: A statistical viewpoint. Acta Numerica, 30: 87-201, 2021.
Peter L. Bartlett, Dylan J. Foster, and Matus J. Telgarsky. Speraizaly-normaliieed margin bounds for neural networks. Advances in Neural Information Processing Systems, 2017.
Peter L. Bartlett, Philip M. Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences (PNAS), 117(49):30033-30070 2020.
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine- learning practice and the classical viri-vari trade-off. Proceedings of the National Academy of Sciences (PNAS), 116(32)15889--1588 2019.
Jeremy Bernstein and Laker Newhouse. Old optimizer, new norm: An anthology. In OPT 2024: Optimization for Machine Learning, 2024.
Blair Bilodeau, Dylan J Foster, and Daniel M Roy. Minimax rates for conditional density estimation via empirical entropy. The Annals of Statistics, 2023.
Adam Block and Yury Polyanskiy. The sample complexity of approximate rejection sampling with applications to smoothed online learning. In The Thirty Sixth Annual Conference on Learning Theory, pp. 228-273. PMLR, 2023.
Bradley Brown, Jordan Juravsky, Ryan Saul Ehrlich, Ronald Clark, Quoc V Le, Christopher Re, and Azalia Mirhoseini. Large language monkeys: Scaling inference compute with repeated sampling, 2025. URL https::
Fan Chen and Alexander Rakhlin. Decision making in changing environments: Robustness, query- based learning, and differential privacy. Conference Learning Theory (COLT), 2025.
11

---

Published as a conference paper at ICLR 2026
Fan Chen, Dylan Foster, Yanjun Han, Jian Qian, Alexander Rakhlin, and Yunbei Xu. Assouad, fano, and le cam with interaction: A unifying lower bound framework and characterization for bandit learnability. Advances in Neural Information Processing Systems,
2024a.
Feng Chen, Allan Raventos, Nan Cheng, Surya Ganguli, and Shaul Druckmann. Rethinking fine- tuning when scaling test-time compute: Limiting confidence improves mathematical reasoning. arXiv preprint 2502::2502.071544 2025.
Jinglin Chen and Nan Jiang. Information theorretic considerations in batch reinforcement learning. In International conference machine learning, pp. 1042-1051, PMLR, 2019.
Yangyi Chen, Binxuan Huang, Yifan Gao, Zhengyang Wang, Jingfeng Yang, and Heng Ji. Scal- ing laws for predicting downstream performance in Ilms. Transactions Machine Learning Research, 2024b.
Tianzhe Chu, Yuexiang Zhai, Jihan Yang, Shengbang Tong, Saining Xie, Dale Schuurmans, Quoe Le, Sergey Levine, and Yi Ma. SFT memorizes, RL generalizes: A comparative study of founda- tion model post-training. Forty-second International Conference on Machine Learning, 2025. URL https:// https://openreview wttpeenreviieeeeeooouuummmm
Rick Durrett. Probability: theory and examples, volume 49. Cambridge university press, 2019.
Anir-massoud Farahmand, Csaba Szepesvári, and Rémi Munos. Error propagation for approximate policy and value iteration. Advances in Neural Information Processing Systems, 2010.
Bahare Fatemi, Jonathan Halcrow, and Bryan Perozzi. Talk like a graph: Encoding graphs for large language models. In The Twelfth International Conference Learning Representations, 2024.
Mare Finzi, Sanyam Kapoor, Diego Granziol, Anming Gu, Christopher De Sa, J Zico Kolter, and Andrew Gordon Wilson. Compute-optimal Ilms provably generalize better with scale. arXiv preprint arXiv:2504. 200222544..220.. 15208, 2025.
Dylan Foster and Alexander Rakhlin. Foundations of reinforcement learning and interactive deci- sion making. arXiv:2312.167711 2023.
Dylan J Foster, Sham M Kakade, Jian Qian, and Alexander Rakhlin. The statistical complexity of interactive decision making. arXiv:2112.13487, 2021.
Dylan Foster, Akshay Krishnamurthy, David Simchi-Levi, and Yunzong Xu. Offline reinforcement learning: Fundamental barriers for value function approximation. In Conference on Learning Theory, pp. 3489-3484444 3489. PMLR, 2022.
Dylan J Foster, Adam Block, and Dipendra Misra. Is behavior cloning all you need? understanding horizon imitation learning. arXiv preprint 1xXi::2407.15003 2024.
Dylan Foster, Zakaria Mhammedi, and Dhruv Rohatgi. Is a good foundation necessary for efficient reinforcement learning? the computational role of the base model in exploration. Conference on Learning Theory (COLT), 2025.
Samir Yitzhak Gadre, Georgios Smyrnis, Vaishaal Shankar, Suchin Gururangan, Mitchell Wortsman, Rulin Shao, Jean Mercat, Alex Fang, Jeffrey Li, Sedrick Keh, al. Language models scale reliably with over-training and on downstream tasks. In The Thirteenth International Conference on Learning Representations, 2024.
Samir Yitzhak Gadre, Georgios Smyrnis, Vaishaal Shankar, Suchin Gururangan, Mitchell Wortsman, Rulin Shao, Jean Mercat, Alex Fang. Jeffrey Li, Sedrick Keh, Rui Xin, Marianna Nezhurina, Igor Vasiljevic, Luca Soldaini, Jenia Jitsev, Alex Dimakis, Gabriel Ilharco, Pang Wei Koh, Shuran Song, Thomas Kollar, Yair Carmon, Achal Dave, Reinhard Heckel, Niklas Muennighoff, and Ludwig Schmidt. Language models scale reliably with over-training and downstream tasks. In The Thirteenth International Conference on Learning Representations, 2025. URL https: //openreview.net wwtp:eeeetooooee http://wwww..eeoooooo/o/////////
12

---

Published as a conference paper at ICLR 2026
Kanishk Gandhi, Ayush Chakravarthy, Anikait Singh, Nathan Lile, and Noah D Goodman. Cogni- tive behaviors that enable self-im improving reasoners, or, four habits of highly effective stars. arXiv preprint arXiv:2200.01130 2025.
Zhaolin Gao, Jonathan D Chang, Wenhao Zhan, Owen Oertell, Gokul Swamy, Kianté Brantley, Thorsten Joachims, J Andrew Bagnell, Jason D Lee, and Wen Sun. REBEL: Reinforcement learning via regressing relative rewards. arXiv:240..16767 2024.
Behrooz Ghorbani, Orhan Firat, Markus Freitag, Ankur Bapna, Maxim Krikun, Xavier Garcia, Ciprian Chelba, and Colin Cherry. Scaling laws for neural machine translation. In International Conference on Learning Representations, 2022. URL https://openreviit/ http://pnnreeieeeeeeeo/rurnn/// fet/oormmidd hR_SMu8cxCV.
Elad Hazan. Introduction to online convex optimization. Foundations and Trends@ in Optimization, 2(3-4):157-325, 2016.
Jordan Hoffmann, Sehastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Train- ing compute-optimal large language models. In Proceedings of the 36th International Conference Neural Information Processing Systems, pp. 300000006300 10016-30030, 2022.
Audrey Huang, Adam Block, Dylan Foster, Dhruv Rohatgi, Cyril Zhang, Max Simchowitz, Jor- dan T Ash, and Akshay Krishnamurthy. Self-improvement in language models: The sharpening mechanism. International Conference on Learning Representations (ICLR), 2025a.
Audrey Huang, Adam Block, Qinghua Liu, Nan Jiang, Akshay Krishnamurthy, and Dylan J Foster. Is best-of-n the best of them? coverage, scaling, and optimality in inference-time alignment. International Conference on Machine Learning (ICML), 2025b.
Audrey Huang, Wenhao Zhan, Tengyang Xie, Jason D Lee, Wen Sun, Akshay Krishnamurthy, and Dylan J Foster. Correcting the mythos of kl-regularization: Direct alignment without overopti- mization via chi-squared preference optimization. In The Thirteenth International Conference on Learning Representations, 2025c.
Yuzhen Huang, Jinghan Zhang, Zifei Shan, and Junxian He. Compression represents intelligence linearly. First Conference on Language Modeling, 2024.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gener- alization in neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
Xiang Ji, Sanjeev Kulkarni, Mengdi Wang, and Tengyang Xie. Self-play with adversarial critic: Provable and scalable offline alignment for language models. arXiv:2466.04274 2024.
Nan Jiang and Tengyang Xie. Offline reinforcement learning in large state spaces: Algorithms and guarantees. Statistical Science, 2024.
Hangzhan Jin, Sitao Luan, Sicheng Lyu, Guillaume Rabusseau, Reihaneh Rabbany, Doina Precup, and Mohammad Hamdaqa. RI fine-tuning heals ood forgetting in sft, 2025. URL httpt://////xiv org/abs/2500 12235.
Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline ri? In International conference on machine learning, pp- 5084-544450 4096. PMLR, 2021.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arxiv preprint arXiv:2001.08361 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations, 2015.
Ben Krause, Emmanuel Kahembwe, lain Murray, and Steve Renals. Dynamic evaluation of neural sequence models. In International Conference on Machine Learning, pp. 2766-2775. PMLR, 2018.
13

---

Published as a conference paper at ICLR 2026
Ben Krause, Emmanuel Kahembwe, lain Murray, and Steve Renals. Dynamic evaluation of trans- former language models. arXiv preprint urXiv:1904008377 2019.
Hong Liu, Sang Michael Xie, Zhiyuan Li, and Tengyu Ma. Same Pre-training Loss, Better Down- stream: Implicit Bias Matters for Language Models, 2022.
Zhihan Liu, Miao Lu, Shenao Zhang, Boyi Liu, Hongyi Guo, Yingxiang Yang, Jose Blanchet, and Zhaoran Wang. Provably mitigating overoptimization in RLHF: Your SFT loss is implicitly an adversarial regularizer. arXiv:2405. 16436, 2024.
Sanae Lotfi, Mare Finzi, Yilun Kuang, Tim GJ Rudner, Micah Goldblum, and Andrew Gor- don Wilson. Non-vacuous generalization bounds for large language models. arXiv preprint arXiv:2312.17173, 2023.
Sanae Lotfi, Yilun Kuang, Mare Finzi, Brandon Amos, Micah Goldblum, and Andrew G Wilson. Unlocking tokens as data points for generalization bounds on larger language models. Advances in Neural Information Processing Systems, 37:9229-995 9256, 2024.
Nicholas Lourie, Michael Y Hu, and Kyunghyun Cho. Scaling laws unreliable for downstream tasks: A reality check. arXiv preprint ar arXiv:2507.00885, 2025.
Xingyu Lu, Xiaonan Li, Qinyuan Cheng, Kai Ding, Xuanjing Huang, and Xipeng Qiu. Scaling laws for fact memorization of large language models. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Findings of the Association for Computational Linguistics: EMNLP 2024, pp. 11263-11282, Miami, Florida, USA, November 2024. Association for Computational Lin- guistics. doi: 10. (055:////222.ffdding.ss... findings-emnlp. 658/.
URL https::ttttctseeeee /appan//ccll hrwpllanntoooooor/g///eee 2024.
Shahar Mendelson. Learning without Concentration. In Conference on Learning Theory, 2014.
Shahar Mendelson. Extending the scope of the small-ball method. arXiv preprint arXiv: (et::000990333 2017.
Tomas Mikolov, Martin Karafiát, Lukas Burget, Jan Cernocký, and Sanjeev Khudanpur. Recurrent neural network based language model. In Interspeech, volume 2, pp. 1045-1048. Makuhari, 2010.
Vaishnavh Nagarajan and J. Zico Kolter. Uniform convergence may be unable explain gener- alization in deep learning. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Vaishnavh Nagarajan, Chen Henry Wu, Charles Ding, and Aditi Raghunathan. Roll the dice look before you leap: Going beyond the creative limits of next-token prediction. In Forty-second International Conference on Machine Learning, 2025.
Behnam Neyshabur, Ryota Tomioka, and Nati Srebro. Norm-based capacity control in neural net- works. In Conference on Learning Theory (COLT),
OpenAI. Introducing openai o1. Blog, 2024. URL https::/o/pepaiii..oo.m com/
Yury Polyanskiy. Channel coding: Non-asymptotic fundamentt limits. Princeton University, 2010.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 1(8):9, 2019.
Alexander Rakhlin and Karthik Sridharan. Statistical learning and sequential prediction, 2012. Avail- pdf.
able at http://ww/ mit. http://www.iittddddrraaaia////re http:///wwahiihhoorreeessstt///a
Amal Rannen- Triki, -rinen Jorg Bornschein, Razvan Pascanu, Marcus Hutter, Andras György, Alexandre Galashov, Yee Whye Teh, and Michalis Titsias. Revisiting dynamic evaluation: Online adapta- tion for large language models. arXiv preprint arXiv:2403.015// 2024.
Dhruv Rohatgi. Adam Block, Audrey Huang, Akshay Krishnamurthy, and Dylan J. Foster. Computational-sttinn statistical tradeoffs at the next-token prediction barrier: Autoregressive and im- itation learning under misspecification. arXiv preprint arXiv:2502. 12465, 2025.
14

---

Published as a conference paper at ICLR 2026
Clayton Sanford, Bahare Fatemi, Ethan Hall, Anton Tsitsulin, Mehran Kazemi, Jonathan Halerow, Bryan Perozzi, and Vahab Mirrokni. Understanding transformer reasoning capabilities via graph algorithms. Advances in Neural Information Processing Systems, 37:78320-78370, 2024.
Abulhair Saparov, Srushti Ajay Pawar, Shreyas Pimpalgaonkar, Nitish Joshi, Richard Yuanzhe Pang, Vishakh Padmakumar, Mehran Kazemi, Najoung Kim, and He He. Transformers struggle to learn to search. In The Thirteenth International Conference Learning Representations, ICLR 2025, Singapore, April 24-28, 2025. OpenReview.net, 2025.
Nikhil Sardana, Jacob Portes, Sasha Doubov, and Jonathan Frankle. Beyond chinchilla-optimal: Accounting for inference in language model scaling laws. In Forty-first International Conference on Machine Learning, 2024. URL https: /owpnrvvieweee.ooooeoeeee net htttfoorumidd=0mmX tTDU..
Charlie Victor Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling L.M test-time compute optimally can be more effective than scaling parameters for reasoning. In The Thirteenth Interna- tional Conference on Learning Representations, 2025, URL id=4FWAwZtd2n.
net/forum?
Yuda Song, Gokul Swamy, Aarti Singh, Bagnell, and Wen Sun. The importance of online data: Understanding preference fine-tuning via coverage. Advances in Neural Information Processing Systems,
2024.
Vladimir Spokoiny. Parametric estimation. finite sample theory. The Annals of Statistics, pp. 2877- 2909, 2012.
Jacob Mitchell Springer, Sachin Goyal, Kaiyue Wen, Tanishq Kumar, Xiang Yue, Sadhika Malladi, Graham Neubig, and Aditi Raghunathan. Overtrained language models are harder to fine-tune. In Forty-second International Conference on Machine Learning, 2025. URL https: //ppnn/vviw.......//e net/forum2jdt-WWedduur
Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen, Xiaolong Wang, Sanmi Koyejo, al. Learning to (learn test time): Rnns with expressive hidden states. arXiv preprint arXiv:2407.0466 2024.
Jianheng Tang, Qifan Zhang. Yuhan Li, Nuo Chen, and Jia Li. Grapharena: Evaluating and exploring large language models on graph computation. In The Thirteenth International Conference on Learning Representations, 2025.
Alexander K Taylor, Anthony Cuturrufo, Vishal Yathish, Mingyu Derek Ma, and Wei Wang. Are large-language models graph algorithmic reasoners? arXiv preprint arXiv:2410.22597, 2024.
S. A. van de Geer. Empirical Processes M-Estimation. Cambridge University Press, 2000.
Aad W Van der Vaart. Asymptotic statistics, volume 3. Cambridge University Press, 2000.
Martin J Wainwright. High-dimenssonal statistics: A non-asympiotic viewpoint, volume 48. Cam-
bridge University Press, 2019.
Heng Wang, Shangbin Feng, Tianxing He, Zhaoxuan Tan, Xiaochuang Han, and Yulia Tsvetkov.
Can language models solve graph problems in natural language? Advances in Neural Information Processing Systems, 0.:30840-306611 2023.
Xinyi Wang, Shawn Tan, Mingyu Jin, William Yang Wang, Rameswar Panda, and Yikang Shen. Do larger language models imply better generalization? pretraining scaling law for implicit reasoning, 2025. URL http::////axxx..o/aabbbbso/
03635.
Wing Hung Wong and Xiaotong Shen. Probability inequalities for likelihood ratios and convergence rates of sieve mles. The Annals of Statistics, 1995.
Fang WU, Weihao Xuan, Ximing Lu, Zaid Harchaoui, and Yejin Choi. The Invisible Leash: Why RLVR May Not Escape Its Origin, 2025.
Mengzhou Xia, Mikel Artetxe, Chunting Zhou, Xi Victoria Lin, Ramakanth Pasunuru, Danqi Chen, Luke Zettlemoyer, and Ves Stoyanov. Training trajectories of language models across scales. arXiv preprint 22.:::2222.098000 2022.
15

---

Published as a conference paper at ICLR 2026
Tengyang Xie and Nan Jiang. Q* approximation schemes for batch reinforcement learning: A theoretical comparison. In Conference on Uncertainty Artificial Intelligence, 2020.
Yuhong Yang and Andrew R Barron. An asymptotic property of model selection criteria. lEEE
Transactions on Information Theory, 44(1):95-116, 1998.
Gilad Yehudai, Noah Amsel, and Joan Bruna. Compositional reasoning with transformers, rnns, and chain of thought, arXiv preprint arXiv:2503.0144, 2025.
Yang Yue, Zhiqi Chen, Rui Lu, Andrew Zhao, Zhaokai Wang, Shiji Song, and Gao Huang. Does re- inforcement learning really incentivize reasoning capacity in Ilms beyond the base model? arXiv preprint XKV::2504.1337, 2025.
Hansi Zeng, Kai Hui, Honglei Zhuang, Zhen Qin, Zhenrui Yue, Hamed Zamani, and Dana Alon. Can Pre-training Indicators Reliably Predict Fine-tuning Outcomes of LLMs?, 2025.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. International Conference on Learning Repre- sentations (ICLR), 2017.
Tong Zhang. Covering number bounds of certain regularized linear function classes. Journal of Machine Learning Research, 2(Mar):527-550, 2002.
Tong Zhang. From e-entropy to KL-entropy: Analysis of minimum information complexity density estimation. The Annals of Statistics, 2006.
16

---

Published as a conference paper at ICLR 2026
CONTENTS OF APPENDIX
I Additional Discussion and Results
19
Discussion and Future Work
19
A.I Simplifications in the Problem Formulation
19
A.2
Future Work
19
Related Work
19
Comparison Classical Generalization Bounds for MLE
21
D Experiments
22
D.I Graph Reasoning Task
22
D.2 Experiment Details for Figure 1
................................
23
D.3 Experiment Details for Figure 2
25
Properties of the Coverage Profile
27
Supporting Results
28
F.1
Properties of the Coverage Profile
28
F22
Analysis of Best-of-N Sampling under a Good Coverage Profile
30
F3
Properties of Maximum Likelihood
30
F.4
Autoregressive Models:
Coverage and Stopped KL--Dvergence
Additional Results
33
G.I
Tightness of Theorem 4.1
33
G.2
Maximum Likelihood: Better Coverage for Convex Classes
G.3 Lower Bound for Maximum Likelihood under Misspecification
34
G.4 Stochastic Gradient Descent: Improved Gradient Normalization for Distillation
35
G.5 An improved tournament via on-policy generation
35
II Proofs
37
Technical Tools
37
H.1
Concentration Inequalities
37
H.2
Information-Theoln Inequalities
38
I Proofs from Section 3
39
Proofs from Section 4
39
J.1
Proof Sketch for Theorem 4.1
39
J.2
Proof of Theorem 4. (Coverage for MLE).
40
J.3 Proof of Theorem G.1 (Coverage for MLE with Convex Classes)
43
J.4
Proofs for Supporting Results..
44
Proofs for Autoregressive Linear Models
45
K.1
Organization
45
K.2
Proof of Theorem 4.2 (Coverage for MLE for Autoregressive Linear Models)
46
K.3
Proof of Proposition 5.1 (Vanilla SGD: Coverage Upper Bound)
53
K.4
Proof of Proposition 5.1 (Vanilla SGD: Coverage Lower Bound)
54
K.5
Proof of Theorem 5.1 (Coverage for Normalized SGD)
57
K.6
Proof of Theorem 6.1 (Test-Time Training)
62
K.7
Proof of Theorem G.2 (Gradient Normalization for Distillation)
63
K.8
Necessity of Variance Dependence in High Dimension
65
Proofs from Section 6
66
L.1 Proof of Theorem 6.2 (Simple Tournament)
66
17

---

Published as a conference paper at ICLR 2026
L.2 Proof of Theorem G.3 (Offset Toumnament
67
18

---

Published as a conference paper at ICLR 2026
Part I
Additional Discussion and Results
A DISCUSSION AND FUTURE WORK
Our work, through the lens of coverage, takes first step toward clarifying the mechanisms through which pre-training with next-token prediction leads to models for which post-training is effective.
A.1 SIMPLIHICATINS IN THE PROBLEM FORMULATION
In the course of the paper we have made various simplifying assumptions. Some of these can be relaxed in a straightforward fashion, while others more fundamental.
In language model pre-training, the pre-training corpus consists of sequences with varying lengths H. and does not typically split examples into prompts and responses. Our formulation in Section 2 is simplification (one that is closer in spirit to supervised fine-tuning), but we expect that the insights derived here can extend to the general setting.
Much of our analysis focuses the realizable/wellssecifee setting where TD II. We give evidence in Appendix G that the coverage profile more tolerant to misspecification than KL- divergence, but leave deeper investigation for future work.
Our treatment assumes the distribution over prompts μ is the same for pre-training and post- training. This straightforward to relax the cost of introducing an additional coverage or distribution shift coefficient to handle the mismatch between the two distributions.
We show that good coverage profile is necessary for BoN to succeed on downstream tasks. While there is ample evidence current RL techniques can fail in the absence of coverage (Yue et al., 2025; Gandhi et al., 2025; Wu et al., 2025), is not clear what the minimal conditions required for RL are.
Our results focus on coverage at the sequence level. For reasoning tasks, is natural to explicitly factorize the response y (Veot: Vans) into chain-of-thought (reasoning trajectory) component Ycor and an answer component Yans- For this setting, a weaker notion coverage is the following
answer-level coverage profile: Cov)*s(TT || 7)) = Pmo (o(Mana) N The answer-level cov-
erage profile is sufficient for downstream BoN success for tasks where it is only important to
produce the right answer, not a correct reasoning trace. We have Cov (TT, f) s Cov f)). but the former be strictly smaller in general.
A.2 FUTURE WORK
Our results open several new directions for future research.
Interventions for coverage. There is much to be done in understanding and improving existing algorithms such as optimizers through the lens of coverage. Our results Section show initial promise for using coverage to guide design of optimizers and model selection schemes, but the algorithm design space remains opaque, and there may be significant room for futher improvement. More ambitiously, one could imagine re-structuring the entire language modeling pipeline itself around coverage.
Semantic coverage. The notion coverage we focus on, the coverage profile, is mathematically convenient but may be conservative in regard to downstream performance, since it only depends on the model through its predicted probabilities. An important direction for future work to understand pre-training and post-training through fine-grained "semantic" notions of coverage that more explicitly account for the representations learned by next-token prediction.
B RELATED WORK
Related empirical observations. On the empirical side, results are connected line of work that studies scaling laws for zero-shot downstream performance based on pre-training metrics such cross-entropy (Gadre et al., 2024; Huang et al., 2024; Chen al., 2024b; Sardana al.,
19

---

Published as a conference paper at ICLR 2026
2024). Several empirical works have also investigated how specific capabilities scale with addi- tional pre-trate training. including machine translation (Ghorbani et al., 2022), knowledge capacity and memorization (Allen-Zhu & Li. 2025; Lu et al., 2024), and multi-hop reasoning (Wang et al., 2025). Our findings are consistent with Liu et al. (2022); Zeng et al. (2025); Lourie et al. (2025); Springer al. (2025), who observe that cross-entropy is not always sufficient for predicting downstream performance, and in some cases can be anti-correlated.
Perhaps most closely related, Chen et al. (2025) show empirically that decreasing cross-entropy in pre-training does not necessarily lead to better Pass@N performance, and that @ can even degrade as pre- training proceeds-a a finding similar to Figure 1.² Our results can be viewed as placing their findings stronger theoretical footing; conversely, their empirical results provide strong motivation for our theorctical treatment. Chen et al. (2025) also study a modification to the maximum likelihood objective aimed improving coverage (in the spirit of Section 6); their approach targets the structure of outcome- based reward, whereas our notion of coverage profile and results are agnostic to the downstream task/reward structure.
We mention in passing some additional works. Chu et al. (2025) explored the different (syner- gistic) roles that supervised fine-tuning (SFT) and RI. play in language model development, and subsequent work observed that the best checkpoint to start RI. from can sometimes be in the middle of SFT training (Jin et al., 2025). Bansal et al. (2025) empirically identified the coverage of teacher-generated synthetic data reasoning tasks. Several papers have also investigated empirical tradeoffs between model size and reasoning performance under best-of-N sampling (Snell et al., 2025; Brown et al., 2025).
important indicator for how effective distillation can be for
Coverage in post-training. Coverage metrics similar to coverage profile play central role in theoretical literature on post-training and test-time algorithms (Huang et al., 2025a;b;c; Foster et al., 2025; Liu et al., 2024; Song al., 2024; Gao al., 2024; Liu et al., 2024; Ji et al., 2024), which analyze algorithms under the assumption that the base model has good coverage; our work can be viewed as providing theoretical motivation for this assumption. Formally, one can use Markov's in- equality to bound the coverage profile by the -p-like coverage quantities considered in these works.
Various notions of coverage similar to coverage profile have also appeared in the more classical literature on offfine reinforcement learning (Farahmand et al., 2010; Chen & Jiang, 2019; Xie & Jiang, 2020; Jin et al., 2021; Foster et al., 2022; Jiang & Xie, 2024); here coverage is typically used to quantify the quality of an offline dataset rather than model/policy itself.
Generalization in deep learning. Understanding the generalization behavior of deep learning models has been a central focus of the theory community for the last decade (Ncyshabur et al., 2015; Zhang et al., 2017: Bartlett et al., 2017: Jacot al., 2018; Belkin et al., 2019; Nagarajan & Kolter, 2019; Bartlett et al., 2020; Bartlett & Montanari, 2021). Our approach is somewhat complementary, in the sense that it focuses on the specific objective of next-token prediction with the logarithmic loss, and aims to understand when minimizing this loss leads to generalization for an alternative objective, coverage profile. We expect that our techniques can be combined with these contemporary generalization results to provide a more refined understanding of generalization for the coverage profile with deep models.
From this line of work, perhaps most closely related are Lotfi et al. (2023; 2024); Finzi et al. (2025), which aim to provide non-vacuous generalization bounds for the cross-entropy loss itself for autore- gressive models.
Analysis of maximum likelihood. Our theoretical results are closely related to classical line of work in statistics ( (Wong& & Shen, 1995; van de Geer, 2000; Zhang, 2006), which shows that maxi- mum likelihood can converge to the true model in Hellinger distance (or other Renyi divergences) under minimal assumptions, even when KL divergence is poorly behaved (large or infinite); see Ap- pendix C below for a detailed comparison. Our results in Section 4 are similar in spirit, but provide a more fine-grae ineegrained perspective. showing that the coverage profile can converge even faster than these results might suggest, particularly as one ventures further into the tail. Our analysis has some conceptual similarity to the small ball method of Mendelson (2014; 2017), which we elaborate on in Appendix J.1.
$ Note that Chen et al. (2025) also uses the term "coverage", but as a synonym for Pass@N; this is not specifically related to the notion of the coverage profile we consider here.
20

---

Published as a conference paper at ICLR 2026
Our techniques are also related to recent work of Foster et al. (2024); Rohatgi et al. (2025), which specializes the general techniques above to autoregressiiv models (c.g., under Hellinger distance).
C COMPARISON TO CLASSICAL GENERALIZATION BOUNDS FOR MLE
In this section we briefly compare our main coverage-based generalization bound for maximum likelihood to classical generalization bounds for maximum likelihood based on Hellinger distance and KL-divergence.
Comparison to KL concentration. For general model classes 1I, the best non-asymptotic KL- based generalization bound we aware of Proposition F.9 (Appendix F), which scales roughly
DKL(TD R) S log Wnax Crine m)
under the assumption that all Il obey sequence-level density ratio bound I W Wnax- Note that autoregressive linear class, we have log Waax BH, matching Proposition 3.2. Combining such guarantee with Proposition 3.1 gives coverage bound of roughly
log Wmax log N
COVN S
this is rather uninteresting since Cov N (T) O for N 2 Wnaxi in other words, we do not get a meaningful improvement as scale N.
The Hellinger distance is a standard metric of distribu- (IP, Q) = f(VP - /Q) The guarantees of maximum likelihood estimation (Wong & Shen, 1995; Van der Vaart, 2000; Zhang, 2006) also imply convergence in Hellinger distance. For general model classes IT, the best non-asymptotic Hellinger-basee general-
Comparison to Hellinger concentration.
tion estimation, defined via (IP,
ization bound we are aware of is Proposition F.8 (Appendix F), which scales as roughly
D2 (TO, R) X Crine (II,
Combining such guarantee with Proposition 3.1 gives coverage bound of
CovN (f) K Crine (II, n)
for all N 2 Compare to the KL-based result above, this result gives a non-trivial bound on coverage when N is constant (comparable Theorem 4.1), but the issue is that it gives no further improvement as we scale N.
We also note that the classical theory of maximum likelihood (e.g., Van der Vaart (2000)) provides asymptotic convergence rates for dddimensional parametric classes IT which have the following form:
Asymptotic bounds for maximum likelihood.
d Crina (II, n),
D (To (T)
S ni
+00.
While this upper bound does not scale with log W sufficiently large burn-in cost no, which itself will typically scale with log Wa, or similar problem- dependent parameters; see, e.g., Spokoiny (2012) for non-asymptotic bounds of this type. Our lower bounds (e.g., Proposition 3.2) imply that there is no hope of removing such a burn-in cost in general.
only be attained with > no for
21

---

Published as a conference paper at ICLR 2026
D EXPERIMENTS
This section presents details for the experiments in Figure and Figure 2. We describe the general graph search task used throughout experiments in Appendix D.I, then detail the specific setups used for Figure 1 in Appendix D.2, and for Figure in Appendix D.3.
D.1 GRAPH REASONING TASK
We evaluate our theoretical predictions using experiments in graph reasoning tasks, in which transformer models are trained to find paths between source and target nodes in graphs. Both graph reasoning benchmarks and synthetic datasets have seen increasing use as abstractions for reasoning problems and for probing language modeling phenomena (Sanford et al., 2024; Nagarajan al., 2025; Saparov et al., 2025; Bachmann & Nagarajan, 2024; Yehudai et al., 2025; Taylor et al., 2024; Wang et al., 2023; Fatemi al., 2024; Tang et al., 2025). These tasks provide minimal abstractions of core reasoning problems, yet are expressive enough to capture pre-training and fine-tuning phenomena. They also offer flexibility in problem structure and difficulty: by specifying different graph topologies and path depths, we can modulate difficulty and expose sources of hardness.
D.1.1 GRAPH SEARCH TASK DESCRIPTION
The graph search tasks for all of our experiments in Appendix D.2 and Appendix D.3 share the same high-level components, and are comprised of
Problem instances. A set of graph search problems G that map bijectively to a set of prompts X.
Data distribution. A distribution over the prompts μ E A(X). and data collection policy To X - A())
Dataset. The training dataset {(x,,)) y)} is comprised of prompts æ ~ y and y To (z).
Next, we describe the general details of the graph search task common to all experiments, as well as how the graph search task is converted to a sequence modeling problem for language models.
Graph problem instances. Each graph search problem in G G is specified by tuple G st). Here, G = V, E) a graph structure with nodes (or vertices) and edges E {(u, u. E v s E V is the source node, and / is the target node. The nodes V are represented as integers, so that V C [m]] for some fixed m E
For all experiments, we utilize a layered directed acyclic graph (layered DAG) for each graph struc- ture (G....) E G. in which nodes are organized into sequential layers with edges flowing only from one layer to the next. The graph G = (V, E¹) has L + 2 layers with disjoint sets of nodes, so that V = (1...(..... V where V denotes the set of nodes in layer i. The first and last layers contain only the source and target nodes, respectively, so that V = {s} and 2 = {t}.
The edge structure E connects only subset of nodes in each layer to the next. We refer this subset in each layer i E {I,...,Le + 2} as its passable nodes V S V or the set of nodes with non-zero out-degree,
V = (o E V* deg (v) > 0}.
The passable nodes in layer are fully connected all nodes in the next layer, that is,
E' = {u,, v) u Vi E Vitl, , {l,..., L + 1}}.
The remaining nodes in V'y \ V have no outgoing edges, and are thus nodes the model must learn avoid in order to output valid paths.
Data distribution. The model's 's task is to imitate the data collection policy which samples only a subset of the (potentially many) valid paths from source to target based on global features of the graph. A valid path from to L is a list of nodes of the form (x, u2, u1.11, where E V for each v (2, ; L + 1): that is, the path must start with the source node & and end with the target node and each intermediate node in the path must be passable node from its respective layer. A graph may have many valid paths, specifically, I |V many. In order for a model to learn valid paths, learning a simple local rule suffices: can output any node in the next layer with > 0 out-degree, which is representable by a fairly shallow transformer.
However, imitating Te is a much harder problem. The data collection policy To samples a subset of these valid paths determined via global rules, complex functions computed over features of
22

---

Published as a conference paper at ICLR 2026
the entire graph that go beyond those required for path validity alone. By varying the complexity of these rules, we can modulate both the difficulty and the nature of the learning problem. This structure naturally maps onto reasoning tasks: following passable nodes corresponds to taking "reasoning steps" that make progress towards the solution, while selecting non-passable nodes corresponds to reasoning errors that lead to invalid solutions. Moreover, when иp selects among valid paths via such global rules, this corresponds to learning high-quality solutions that accurately reflect desired properties for the problem.
Dataset. Recall that the model learns to imitate TTp from dataset D ((r, where each prompt corresponds to graph search problem G = (G, s) E G, and each response y ~ To( is an expert response. formatted as follows.
We convert given graph search problem G = (G, 8, t) E g with graph structure G = (V, to prompt by concatenating the edge list E, the source node s, and the target node 4, formatted as
I u v_1 2_2 v_2
u_k v_k stt
where (u,, 0;) E [m]] are the vertices of the i-th edge in the edge set E. For formatting, the special character separates two edges, the character / separates the adjacency list from the source and target nodes, while the character marks the end of the prompt.
As an example, for edge set E = {(10, 23), (86, 47), (45, 32)}. the prompt is
x: 10 23 | 86 47 1
45 32 / 10 45
Next, each response encodes the path from the source to the target node in G as a sequence of nodes. That is, the response takes the form of string
V_1 v_2 v_3
v_H-1 V_H
where E [m] is the i'th nodes in the path for each i [M]. and 01 while ви f. Here, the horizon H corresponds to the path length in G, and in the layered DAG we have H = =I + 2.
Summary: Graph search to sequence modeling problem. In summary, a graph search task with
set of problem instances & induces autoregressive sequence modeling problem with vocabulary space V = [m] U {1./. =) prompts X' C V corresponding to search problems in a layered DAG graph structure with L layers, and responses y y corresponding to paths with length H = + + 2. In addition, the task is equipped with μ E ((X) and To X' A()) that is used to collect the training dataset D = {xx,)) }) where x ~ p and y ~ mp()).
D.1.2 MODEL DETAILS
Next, we describe the common implementation details for the models we train to solve the graph search task.
Tokenizer. We use numeral tokenizer, which is standard for graph reasoning tasks (Sanford et al., 2024; Bachmann & Nagarajan, 2024). Each node e E [m]] tokenized as its integer node value, and the special characters I, /, and are tokenized m 1, m + 2, + 3, respectively.
Transformer model. We train causally GPT2-like transformer models to minimize the cross-entropy loss using the Adam optimizer with fixed learning rate, and perform a grid search over the parameters displayed in Table 1. Parameters with fixed values were chosen based related papers such as Bachmann & Nagarajan (2024). In both experiments, the model architecture with 4 heads, hidden layers, and 384 hidden dimensions worked best. We use absolute positional encodings. Training iterations and grid search values for the learning rate are different for each experiment, and discussed further below.
D.2 EXPERIMENT DETAILS FOR FIGURE
The graph search task for Figure exposes natural properties pre-training data under which cross- entropy reduction comes at the cost of worse coverage profile. The key idea is that because the pre- trainining data is diverse (with multiple distinct modes or graph classes), the model is unable to perfectly fit the distribution. As result, when one mode of behavior is better-represented than another, cross-entropy minimization, which is an average-case distribution distribution-matching matching metric, can sacrifice coverage across the different modes in order to increase performance single mode.
23

---

Published as a conference paper at ICLR 2026
Hy perparaaeeee
Values
Number of heads
[4, 6, 8) (3, 4, 8]
Number of layers
Hidden dimensions
384
Activation function
GeLU
Batch size
128
Weight decay
0.01
Table 1: Hyperparameter grid search values for transformer models in graph search.
Concretely, the graph search task for Figure 1 is a mixture of two classes of graph structures. Due to representational and finite-sample constraints, the model is unable to fit both perfectly during training, and, in particular, fitting one class well (in the sense of cross-entropy loss) comes at the cost of worse performance on the other. The checkpoint with the best coverage arises at some middle point in training when the model learns both classes of graphs equally well, and has good coverage over both classes (the dip Cov, the leftmost subplot of Figure 1). Further reduction of cross-entropy loss over the latter half of training requires the model to lose coverage over Te in the less-represented graph class (observed as the increase in Cov the latter half of training iterations).
Even though the task cannot be learned perfectly from the supervised learning feedback, the model can still learn policy that always samples correct path matching To 's with N O Best-of-N sampling attempts, which means that it leads to efficient downstream post-training (e.g., on one of the modes or with reward-based feedback), and also achieves optimal performance with test-time scaling methods.
For the experiments in Figure 1, we first pre-train a model on a larger set of graph structure classes so that it learns diverse set of behaviors, then finetune its behavior two. The performance on the fine-tuning task displayed in Figure 1, and we first describe the fine-tuning dataset, followed by the pre-training dataset.
D.2.1 TASK DESCRIPTION
All graphs in follow the layered DAG structure described in Appendix D.1 with I. 8 interme- diate layers that each have 4 nodes, i.e., |V*| = 4 for layers E (2....,9) (recall the first and last layers contain only and respectively).
> 0} denotes the set of passable nodes. For each graph problem G = (G, s) G with graph structure G (V, E), subset of the layers indexed by Is C (2....,9) with [g]]
Recall that in layer i, V = {u V*
randomly selected. Then, the edges are defined so that the layers in 12 have two passable nodes each (i.e., [V 2 for E T22, while the remaining layers have only one passable node each (i.e., |V.
for i E (2.....9) .9 T2). The passable nodes in each layer are chosen at random, but for the layers in 12 are guaranteed to have one even and one odd node. For each graph in G, there are 2 = total valid paths since |/2| = 2 layers have two passable nodes each while the other layers have one.
Data distribution. The set of problem instances G -i of problems, 9i and G2. The prompt distribution in the fine-tuning task is a skewed mixture over the two classes with ğü E ({11, denoting the probability of each class in the data; within each class, the graphs are drawn uniformly at random (described at the end of this section). Although there are 4 valid paths from souree to target, in each class Gn or G2 the policy To chooses one path based on different global rule, described below.
Gi L G2 is comprised of two disjoint classes
Class Gi (probability ((1) 0.9). For integer j E Z, let the function p(j) = (j mod 2) denote its parity. For layers with |Vil = 1, To deterministically selects the unique passable node. For layers i E In (where [V]| 2), the set V contains one and one odd node, and कe deterministically chooses the node u E V such that p(u) = p(t); that is, the node whose parity matches the parity of the layer index.
24

---

Published as a conference paper at ICLR 2026
Class G2 (probability (2(2) 0.1). For layers i with [V = 1, To deterministically selects the unique passable node. For layers i E Iz (where |V'| = 2), To chooses the node v E V. such that p(u) 1 + p(t); that is, the node whose parity is opposite the parity of the layer index.
The class of a graph is technically identifiable from the prompt by computing a parity-based feature over randomly selected subset of the nodes, but this problem is too difficult for the model to learn in the fine-tuning stage. Let V C V be fixed subset of nodes whose cardinality is half the total number of nodes the graph (i.e., [V*| = |V1/2). Then all graphs Gi satisfy 1 = Duev p(u), while all graphs in G2 satisfy 0- Ower VEv p(u). However, determining which nodes belong V requires complex reasoning over the graph structure.
Dataset. Each sample in the dataset D = {xx y)} is then generated via the following procedure.
1. First sample an index ~ Jt.
2. Sample G E G; by randomly drawing V C [m]] without replacement, and instantiate the edges according to the description for each class above.
Format the prompt æ per Appendix D.I.
4. Draw y To(: ær) according to description for each class above.
D.2.2 PRE-TRAINING DESCRIPTION
The graph problem instances the pre-training task, Gpre are a superset of the graphs in the fine- tuning task, that is. Use 94 = Gpre with K = 3, and 91 and G2 defined as in the previous section for the finetuning dataset. The data distribution is uniform mixture of these classes, ji(it) 1 for each E [K]. and the third class G3 shares the same layered DAG structure as Gi and G2 (with L= 8 intermediate layers, where two layers are randomly chosen have multiple passable nodes). However, in Gs. To is stochastic policy and samples one of the 2 = 4 valid paths at random. The dataset then drawn using the same data generation procedure described for the fine-tuning task above.
D.2.3 TASK- IMPLEMENTATION DETAILS
The transformer model first pre-trained on fixed dataset drawn from the pre-training distribution, with 8 x 64, 000 prompts in total, using a learning rate of 1e-4 4 for 200k iterations, which was chosen based on a grid search over learning rates (5e-5, 1e-4, 5e-4).
The final checkpoint is then finetuned for 50k iterations in an online fashion, where fresh samples are drawn for each batch (this is equivalent to offline training with a dataset that has an equivalent number of samples). The learning rate is 5e-6, which was chosen based grid search over learning rates (5e- 6, 1e-5).
D.3 EXPERIMENT DETAILS FOR FIGURE 2
For Figure 2, we consider family of tasks that is parameterized by the horizon H, in order to expose the fact that cross-entropy is sensitive horizon, but the coverage profile is not. This construction leverages the intuition from Remark 3.1. The training data is heterogeneous, with fraction con- sisting of difficult graph problems that the model cannot learn to cover with the given number of training samples. This un-learnable subset of the data contributes to the large KL-divergence, but does not affect the coverage profile.
D.3.1 TASK DESCRIPTION
For Figure 2, we devise family of tasks parameterized by the number of intermediate layers H E (8, 16, 24). For a fixed H. each task GH utilizes the layered DAG graph structure described in Appendix D.I with L H intermediate layers, each containing 4 nodes, so that each graph has H total layers (including source and target). The response space y >HH+ corresponding to paths of length H (including the source and target nodes).
Data distribution. The task is heterogeneous mixture over 3 classes of graphs described below that we refer to as GH.1 1 U GH. U GH. = -H- The classes GH.2 and GH.3 are significantly harder to learn and the model will fail to do with the given number of training samples, even though GH. is learned quickly (and also provides useful features for learning the other two tasks). The distribution over these 3 classes is fixed for all H and specified by j ((1,2,33)
25

---

Published as a conference paper at ICLR 2026
Class GH. (probability j=(1) = 0.94). All H intermediate layers have only 1 passable node each (i.e., VV = 1 for all i E (2...., H + 1 }), so each G IIr, has only one valid path from source to target. For prompts corresponding to graphs in this class, Tp deterministically selects the unique valid path.
Class GH1 (probability (2(2) - 0.05). For each graph, half of the intermediate layers (or H/2) are randomly selected to have two passable nodes, while the rest have one. More formally, subset I1/2 C {2..... H + 1] with //// = H/2 is randomly selected, such that [V!| for E In/2 and |V for i E ( 1] TH/2-
There are 2 valid paths from source target, and иp deterministically selects one of them. For layers with |V'| 1, selects the unique passable node. For layers E 41/2 (where |V = 2). To selects the node E V by following difficult, deterministic rule. This rule requires To to select the node e whose parity matches the parity the layer index, XOR'ed ed with the parity each passable node in the entire graph. More specifically, recall that p(j) denotes the parity of an integer #f+1 = V denote the of all passable nodes across all intermediate layers (including those with just one passable node). Then in layer E /H/2, mn selects the node V such that p(v) = p(i) @ uEEv p(u)).
j E [m]] and let V. :-
Class 9H.3 (probability x(3)) = 0.01). Regardless of H, for each graph a subset 1T C { ... H 1} with |4| 4 is randomly selected, such that [V_| for i 14 and |V-|| = for i (2...., H/ + 1} \ 4. There are 2 = 16 valid paths from source target. The policy samples uniformly at random from these valid paths.
Note that prompts/graphs from each class are distinguishable from each other (or, identifiable) based on prompt features alone, so a powerful-enough model can achieve perfect performance across all of them simultaneously. GH,2, for example, has more edges and thus longer prompt than GH.1; similar statements apply to 9H.3. Dataset generation occurs in the same manner as described in Appendix D.2.
D.3.2 TASK-SPECIFIC IMPLEMENTATION DETAILS
Lastly, describe experiment-specific implementation details top of those previously described in Appendix D.I, which are common to all experiments. In addition to grid search over the parameters in Table 1, we perform search over learning rates (5e-5, le le-4, 5e-4], for which the learning rate of 1e-4 exhibited the best validation performance. The model is trained for 40k iterations over fixed dataset of 8 x 64, 000 samples.
The results in Figure 2 are computed from evaluations of training checkpoints on per-class validation datasets of 1024 prompts from each GH,; for E [3]; these metries are then averaged according to the probabilities in to obtain the final result. In total we ran 16 seeds, and plot their median. The I6 quantile.
shaded region Figure displays the region between the T6 quantile and 15
26

---

Published as a conference paper at ICLR 2026
E PROPERTIES OF THE COVERAGE PROFILE
Before proceeding, we briefly discuss some properties of the coverage profile that will be helpful to keep mind.
Remark E.1 (Coverage profile as a refinement of cross-entropy). The coverage profile can be viewed as a fine-grained, inference budget-sensitive refinement of cross-entropy. Concretely, if we write
To x)) x)
CO || ②) = Pm log
log N
(16)
#(y
it becomes clear that the coverage profile is simply the camulative distribution function (CDF) of To(u)z)
the log density ratio X == log #(ylz)' while KL-diiergence corresponds to the mean: Em, (X). It is well known that the CDF of a random variable is more informative statistic than mean (Durrett, 2019); the former can be much more sensitive to the model's 's behavior at the tail than the latter. Indeed, the coverage profile can behave very differently across scales, as shown by Figure 1.6
Remark E.2 (KL divergence and coverage profile are not estimable). We emphasize that KL. divergence and the coverage profile are not estimable quantities in general, due to the fact both depend on the unknown density To(y 1)) for the data distribution. This motivates the use of cross- entropy in practice, as the former is an estimable upper bound on DKL(TT) || R). Analogously, we show in Section 6.2 that various estimable proxies for the coverage profile can be used to select mod- els with good coverage. One exception the expert distillation setting (see Appendix G.4), where To teacher network for which the pog-probabilittes log to x) are available.
'Interestingly, show (Proposition F.1) that if the coverage profile satisfics certain growth condition uniformly for all scales M, then it implies bound Kl--divergence-a weak converse Proposition 3.1.
27

---

Published as a conference paper at ICLR 2026
F SUPPORTING RESULTS
This section presents technical results used throughout the paper. Appendix F.1 presents basic properties of the coverage profile. Appendix F.2 analyzes the performance of the Best-of-N algo- rithm under coverage. Appendix F.3 presents properties of the maximum likelihood estimator, and Appendix F.4 presents structural results relating the coverage profile to "stopped" KL-divergence, which are useful for analyzing autoregressive models.
F.1 PROPERTIES OF THE COVERAGE PROFILE
This section presents elementary properties of the coverage profile.
Proposition F.1 (KI-tt--ovverage conversion). For all models To and л and M 2 2, we have
COVN(T)) S DKL x)
log N 1+ I
Proof of Proposition F.1. Lemma 27 of Block & Polyanskiy (2023) states that for any N >> 1 and
any convex f [0, x]] + [0, oo] with f(1) = (1(1) = O.
(y x) T
(T 1l
COVN (T) mg
N
(17)
x))
f(N)
where (T || m) Ex [f dm ]]
Applying this with KL-divergence, which corresponds to = log æ, we have that
dr
fæ) = r log x with f (z)
N J(N)
1 log N 1 1
(18)
which gives result.
Proposition F.2 (Tightness of KL-to-coverage conversion). For any N 2 2, there exist models To and To such that
DKL (TD ()
CovN log N
2N
Proof of Proposition F2. Consider Ber(p) and 7 = Ber(p/N) with S Then COVN p and
1-p ( 1 -1)
- p
< plog N + +(1 (p)
DKL || () = ploog N + +(1 - p) log
p
1-
N
1
(11- = p log N -(1 - p)
Kv
1
Sp. log N
+
2
2N
This is the desired result.
Proposition F.3 (Uniform coverage decay implies bounded KL). Given n, mp - A()), define
Wnex == sup x(sla)
So(g)) and
C: sup (C) N),,
N21
where we note that C' S log Wmn.- It holds that
Dia (m ())) C. (1 log(log( (CC).
(19)
28

---

Published as a conference paper at ICLR 2026
Proof of Proposition F.3. Let 0 fixed parameter, and define X To/ Then we have
DKL (To || x) E [log(X)] s (r,]log(())((lggg > 83] +
(20)
Since X W almost surely, write
Bm log( (T)[(lo(XX)) > 8)) =
P [lo((() X) > t/dt
(21)
Wog
IP mm
X e e'Ict
(22)
Wan)
I C
dt
(23)
5
log W nax)
= =Clog
(24)
The result now follows by setting 8 = C.
Proposition F.4 (Hellinger-to-coovraae conversion). For all models and and N > 1, we have
2N
COVN(TT ll (x)
(VN 1)2
Proof of Proposition F.4. Without loss of generality, we assume y is discrete the following
proof. By definition,
Vzly x)) 2
DH x) Exonto
Σ
To x²)
y1
2
Enraso
Σ
Toly xt)
(1- 2
VN )²
i if x) S T m)
2
P..
To
x) > N
x))
2
where the inequality follows from the fact that v T x) implied by x (y x) < 3/4 To x). Re-organizing completes the proof.
V
(
(1 (1-vx)
v/x) V x)
Proposition F.5 (Chain rule for coverage profile). For any models np. कi, and म., and any M1, M2 > we have
Cov (TT ()) S M2 COVM, (TD #) + COV (TI To).
(25)
Proof of Proposition E.5. We can write
TT w) F(y r)
Covar Mr (TT)) || 7) = P...
Mi
TT x))
TT (x)
TT (x)
कT x)
Pm
M1.
$ M2
+ Pr.
> Mr
> M2
7(y x)
(r)
S x))
Tp x)
To xx)
क (y
x))
< MzF n
M1 M2
+Pm
M2
F(y x)
T
x)
M2CoV M,, (mD
())
Cov (7T
mT).
29

---

Published as a conference paper at ICLR 2026
F.2 ANALYSIS OF BEST-OF-N SAMPLING UNDER A GOOD COVERAGE PROFILE
In this section we analyze the performance of the Best-of-N algorithm under a good coverage profile. Let a base model î be given, and let a reward function rr y) E [0, 1] be given. Let क. X + A denote an arbitrary task- susk-specific comparator policy.
We let N (z) denote the distribution of the Best-of-N algorithm with parameter N, which draws N responses y² ;-... y ii.d #(- (r) and returns y = arg max, v. r (T(2,,e).
Proposition F.6 (Coverage implies success for BoN), Let M 2 be given. For any e > 0, if N 2M log(e-) 1 and ( f) S 1.2 then we are guaranteed that
B क (1)) rr "N (z)) s ( +
(26)
Proof of Proposition F.6. This is an immediate consequence of Lemma F.1 in Huang et al. (2025b),
noting that we can bound EM #) < Cov (x ?).
Proposition F.7 (Coverage is necessary for BoN), For any model म and reference mi, and for any N 2, there exisly a reward function rr (2, v) E {0, 1} such that
Be [M((,, (x()) - rr BON (x)) 2 2
CoV2N (7)
(27)
Proof of Proposition E.7. For any X, we define Se == {y yy sigle) rr( y) = Lyy S.}.
2 2N} and let a(paa)
By definition, for any fixed a E X, it holds that r2(2, (x)
(x)) P. y Sa) yl (... (3i [N]. V E Sx)
- P (0 E Sa)) & N. ...P
Sz))
= N Σ 7(v x) N. Σ
2N T x)
1 -
PP you (- (Sz),
2
yES,
BES,
where we use the fact that 7ly | x) S 2N (T | x) for any y E Sz. P (1)p E Sx) = CoV2N ( T). Therefore,
We also note that
Exsp (x(((,,T(x))) ra( #BON (x))) 2 2
COV2N ( (7).
F.3 PROPERTIES OF MAXIMUM LIKELIHOOD
In this section, we specialize standard guarantees for maximum likelihood (Wong & Shen, 1995; van de Geer, 2000; Zhang, 2006) to derive bounds on the coverage profile; as discussed in Ap- pendix C, these results not tight compared Theorem 4.1.
Proposition F.8 (Convergence of maximum likelihood in Hellinger distance). Assume that To E II. With probability at least 1 - 6, the maximum likelihood estimator म == arg max Ln( satisfies,
log Nog (II. +
Di 7) S inf
(28)
E20
and consequently
Noo eII,e)
COVM S inf
+
(29)
2>0
n
for all M 2 2.
Proof of Proposition F.8. The first bound follows from Proposition B.2 of Foster et al. (2024). The second bound follows from applying Proposition E.4.
30

---

Published as a conference paper at ICLR 2026
Proposition E9 (Convergence of maximum likelihood in KL). Assume that II, and that all и E II satisfy m
S Wnax. With probability at least 6, the maximum likelihood estimator
= arg max Ln (T) satisfies,
log No (TT,
(TD ||
S log W nax
inf e20
tEg.
(30)
and consequently
log W nax
log No + ejj
COVM
2
inf c20
(31)
log M
n
for all M 2 2.
We remark that the log( considered in Proposition 3.2, it holds that log. Noo (II, E) S log(1/e) V and In S 2H
Ong(Wnns W na factor in Eq. (30) can be tight in general. For example, for the class
Proof of Proposition F.9. By Lemma 4 of Yang Barron (1998), holds that (KL( (T) (2 log W (Wnax)) D ( (To. R).
Therefore, the first bound then follows from Eq. (28). The second bound follows from applying Proposition F.I.
F,4 AUTOREGRESSIVE MODELS: COVERAGE AND STOPPED KL-DIVERGENCE
This section shows that we can relate the coverage profile to "stopped" KL-divergence defindd Eq. (32). This useful result in the context of autoregressive models because the stopped KL-divergence is always bounded, even when KL-divergence itself may not be.
Proposition F.10. Define the stopped Ki-divereence for parameter N as
H
}
Dseq.N [T) = E (a Jaath
min
log N. Σ DKL(TO(- x, (1..ii-1) 1) [[ x. 11-h- 1))
h=1
(32)
Then as long as N > e, holds that
2
Co ()))
D Deeq, NoT IT).
(33)
log N
Proof of Proposition F.10. Consider the stopping time
7: := min
h: ha = H or DKL(TTO(UHH+ = (1.)) (1 = (1.))))) > log N
jSh
Then, for the process Y = (x, M1r)), we have the chain rule:
DKL - =)
Em
DKL (To(Un = 1 10) T(h = -I (1.8):n2))
"
< Exg min log N,
Σ h=1
DiL (mo(un = 11 T(m 12, 11n
where the inequality uses Ejer DKL (mo(5t+ which follows from the definition of r. Therefore, by Proposition F.1, have
x. )).) ( (
x,,))))) 5
log N,
(YT) x(Y*)
DKL DRL (Y = (Y =))
TP no
> log N)
log N 1/N
31

---

Published as a conference paper at ICLR 2026
Finally, we bound
TD
(U1.
TD() +(+))
TD
2 N
S Pm < H1) + Pmo
log N
(V1:H
By Markov's inequality.
H1 DKL (Ta(- 2t (1.h || (T I (1--1)) > log NV)
(T < H1 S Pmo
M==1
Bn min log N, Σ (To(- a M1ste-1 (+ 2 31-h-1)))
log
Combining the inequalities above completes the proof.
The following result sort of partial converse Proposition F.10, showing that the coverage profile can be lower bounded in terms of the tail behavior for sum of step-wise Hellinger distances.
Proposition F.11. For any N 2 1 and 6 (0, 1), it holds that
H
Cov (To || T) 2 P.o.
Di æ, 11---11), (T( x. v1:h-11) 2 log(AV/6)
Proof of Proposition F.11. By definition.
1 (ym x. 11:h-11)
Byn exp - log
π ( r, 11=h-1)
Σ
(D(h x1,11:-11 ( I. V1:h-1)
V4 EVEY
= 1 - (= 1T, ([- BLA- I exp((DR(To(- Ix, 11,1)A-1), ((- 11)
Therefore, holds that
Di(TOO x. 1,Nin-1), a( (1:-11) log ( 2. 11:h--1
Emg exp
I 1.
2 T x, B1:A-1)
By Markov incquality, this implics
H
( x) log
P mo
DĀ (mo( x. 10.8h.--): T(. (1.h log(1/6) $
x)
h=1
To conclude, we note that
H Da(To(: 11, -1:h 2 log(N/6)
H
1
To( x) T x))
< IP mb
(= x, (1: w1:h 1))h
2 log
+ log((1 16)
2
h=
TV x) + log(1/0) 2 log(N//))
+ P mu
log
( x)
< + Cov (Tp T))
Re-organizing gives the desired result.
32

---

Published as a conference paper at ICLR 2026
G ADDITIONAL RESULTS
G.1 TIGHTNESS OF THEOREM 4.1
To conclude, we show that the coarse and fine- ane-grained terms in Theorem 4.1 are both tight in general.
Proposition G.1. The following lower bounds on coverage hold for the maximum likelihood estima-
tor.
(a) Coarse rate: For any n, d 2 and B Z log(5n), there exists a class with log.Noo (II. a) 2 d dlog( B/a) V and To E II such that with probability at least 0.5, it holds that for any N I e
d
COVN 2c.
(b) Fine rate: For any 2 d N,N 1, there exists class II and no E II such that [TT] = 2d +
and Noc (II, a) & 2 for any 2 V and with probability least 0.5, it holds that
d
COV (R) 2
n log
Informally, case (a) shows that for the class II under consideration, the coverage does not decrease with log N until N is trivially large such that log No: (II, log N) = 0: this is precisely the behavior of the coarse term in Theorem 4, so this implies there is no hope of removing this term. Meanwhile, case (b) can be interpreted as showing that there is no hope of replacing the high-precision covering number found in the fine-gru tin-igrined term in Theorem 4.1 with coarser notion (e.g, at the scale in the coarse-grasd -ggainndd term), since the rate grows with = log III even though log (II, a:) is constant for
We note that Proposition G.1 is an algorithm-specific lower bound, not an information- theoretic lower bound; show in Section 6.2 is that it is possible to improve over Theorem 4.1 with algorithms explicitly designed to optimize for coverage.
G.2 MAXIMUM LIKELIHOOD: BETTER COVERAGE AGE FOR CONVEX CLASSES
In this section, we give an extension Theorem 4.1 which shows that maximum likelihood can achieve a faster convergence rate for coverage as well as strong tolerance to misspecification when the model class is convex.
Assumption G.1 (Convex model class). The class II satisfies II = {mo E Q) for convex, compact parameter space O, and the mapping 0 + To(y x)) concave for all X, y E.
Theorem G.1 (Fast convergence of coverage for convex classes). Let 2 0. N 2 1, N 2 2c be given, and suppose that Assumption G. holds. Lei
8" arg min DKL mo). BEB
With probability at least 1 8, the maximum likelihood estimator मे arg max Ln( satisfies
log Noo (II, - log(8-1)
Ne
Nog (I1, (II,s) 8)
N S Cov (TO-)
inf
+
n
N
6>0
n
(34)
where is an absolute constant.
Note that we allow for misspecification here, as Eq. (34) shows that the coverage of र can be upper bounded by the coverage Tæ* the best-in-class approximator ua with respect to KL-divergence. In the well-specified case where mo E II, the bound simplifies to
1
log. Nog (II,
log (I,,clog N) + log(6-1)
Cov.N (ft <
N1-2c inf
n
Criee (II, n)
Croarse (11, N,
N1-2c
which improves upon the rate Cov (7) of Theorem G.1 is presented in Appendix J.3.
(lm()))))) log
Ccoarse (II, (II,N, in Theorem 4.1. The proof
J.3.
33

---

Published as a conference paper at ICLR 2026
G.3 LOWER BOUND FOR MAXIMUM LIKELIHOOD UNDER MISSPECIFICATION
In the following proposition, we show that without well-specified model class (Assumption 2.1), maximum likelihood may have coverage profile scaling with log M milkee DRL (To || () (cf. Propo-
sition F.3), even when there exists π E 1I such that CovN(T) = O.
Proposition G.2 (MLE under misspecification). For any a E [0, 1]. M e there exists problem instance ny and class II {T1, m2 such that
ca2
supllog To x) - log mı (y 13)5s Sa, a.M/
( 2
log M
and for any 1, it holds that with probability at least 1. the MLE क i.e. COVN(
a²
log M
Proof of Proposition G.2. Let p P. p)) = p, and To is given by
32 M. Consider X (+, {0,13, p(-)
1 (/)
Tn(- +) To(- ++) Ber
We construct the class II = (mi, T2 as
1
का (+I+)= Ber
n2(-1-) = Ber
2ea
2
1
++(++) Ber
(+(-))
Ber
2MM
Given the dataset D = {(xxxx v²) heini sampled from Np. we define N(x,)) = # [n] (x',0))= v)) and N(x))= (xx, 0) Nar, 1).
Then
T. (T22 - (mI) = N(+,1) 1) a N(+,0) log
2ea -1 - -)-,1)) log M + N(-,0) log 2
aa
M
By symmetric, it holds that (+, N)+,0)) 2 1 Further, by Markov's inequality, holds that (P(N(-) 2 4np) S +. Therefore, for the event E (N((, 1) 2 N(+, (,,,,,,N)()))) I 4np), we have PEE)= 2 +. In the following, show that L, (m2) Z.. (TI) 0 under E.
We condition on E. We first note that under this event, we have N(+,1) 2 ((++)) S (+).. Hence,
-1] ) NN() -log M
IT( - (m 2 N(+)
a
log
2
2
2ea -
= N(+) DKL
BBrr
| |Ber
-N()) M
2ea
2 N(+) (1 - ea-2 -N(-): -log M.
Finally, using the fact that a > a and N(+) 2 (1 -4p)n 2 f under F,, we have
im(TT) [m (a²
02 - 4plog M n O.
Hence, under the event E, we have ãे n2. However, it is clear that
CoVN(T2) = p. (O =
This completes the proof.
34

---

Published as a conference paper at ICLR 2026
G.4 STOCHASTIC GRADIENT DESCENT: IMPROVED GRADIENT NORMALIZATION FOR DISTILLATION
In this section, we focus on autoregressive linear models (3), and consider variant of our setting inspired by distillation We assume that for each example (x", Vi.H), for each h 1...., H, we have access to the true next-token probabilities (t( x², Vi-h-1) for all y/ V. This an unrealistic assumption for general pre-training, but it is natural for distillation, where corresponds to a teacher model (in particular, the next-token probabilities are already computed as part of a standard forward pass through the teacher model).
For the distillation setting, we give an improved gradient normalization scheme that improves upon the rate achieved by Theorem 5.1, closing the gap between SGD and maximum likelihood by match- ing the guarantee for Theorem 4.2.
Define (I 31:A-1) == DKL (T x. 91:h-1) || TT(- x, 11:h-11)): note that for the distillation set- ting, we can compute this quantity in closed form for any prefix z, V1-A-1 in the training corpus. We consider the following sample) iruncated/normalizedd stochastic gradient estimator:
H
ge
x)
ag(2,y1n-1) log (e( x. 91:h-1),
(35)
where A = log N, and where
Lo S A,
ao(x, 91:h-1)
Ejen-1colz, B1.)) A,
(36)
A- E E sola.inis) co(z.nna1)
otherwise.
With this definition, we define the following normalized SGD update:
++ = Proje (0 + n (y² ==)))
(37)
Intuitively, the idea behind the update in Eq. (35) is to truncate the gradient at the point where the KL divergence between the teacher and student model is too large, and then normalize the gradient by the KI. divergence; this is inspired by the structural result Proposition F.10 in Appendix F.4, where we show a close connection between the coverage profile and a certain "stopped" variant of KL divergence.
Theorem G.2. T,N be given. With a suitably chosen stepsize > 0, the normalized SGD update (10) achieves the following coverage bound:
T
B²
CoVN (ToC)
S
+
(38)
T
T log N
T
This guarantee matches the rate of Theorem 4.2 for the maximum likelihood estimator. The proof is presented in Appendix K.7.
G.5 AN IMPROVED TOURNAMENT VIA ON-POLICY GENERATION
We describe improved tournament estimator that is able to remove that 1/N1-a term from The- orem 6.2, meaning it achieves nontrivial guarantees even when the coverage parameter is a con- stant.
mtvla) Note that the term 1/N1-a of Eq. (15) comes from the fact that Fng [)]
2 N) can be as large
1/N in the worst case, implying that the â produced by Eq. (14) may at best achieve coverage of /N. To overcome this, we introduce an offset term:
== arg min max,'enl [Cov (T' (T) - 2N Cov
(39)
2 N for models M', म. This estimator augments the simple tournament in Eq. (14) with an "offset" term that accounts for the fact that some of the models might be quite far from mp. The main guarantee is as follows.
where define Cov ( (| 1 E Py-
35

---

Published as a conference paper at ICLR 2026
Theorem G.3. Fix N 2 1, a such that 11 2 4. Suppose that there exists ग. II such that
[log (V z) - log क ( x)] I a log for any x E X. y. Then with probability - 8, the
tournament estimator (39) achieves Co (i) X [og(
Compared to Theorem 6.2, this tournament eliminates the additive 1/N1-e term. It does, however, require stronger condition on the best-in-class model म that |log To x) log ( / x) - log N,, which implies in particular that Cov (T) O.
Infinite classes: Beating maximum likelihood. While we motivated the tournament estimators through model/checkpoint selection with a finite class II, both estimators can also be applied to general, infinite classes II. In this case, it turns out that they both improve upon the coverage achieved by the maximum likelihood estimator in Theorem 4.1, even in the well-specified case where To E ПΤ; informally, the tournament estimators allow us to remove the fine-grained term in Theorem 4.1, leaving only coarse-grained term. See Theorem 6.2' and Theorem G.3' for the formal statements.
36

---

Published as a conference paper at ICLR 2026
Part II
Proofs
H TECHNICAL TOOLS
Notation. We denote by (R) == {o e R |||||| - R] the d-dimensional Euclidean ball of radius R. We drop the superseript when the dimension is clear from context.
H.1 CONCENTRATION INEQUALITIES
Lemma H.1 (Freedman's 's inequality). Let (Z')ign be real-valued martingale difference sequence adapted to a filtration (Fi)snn 1 almost surely, then for any rj E (0, 1/R,, with proba- bility at least 1 - 8, for all n. S n,
n' Bi (
log(d-1)
Z 1
"
The next result standard consequence of Lemma H.1 (e.g.. Foster al. (2021)).
Lemma H.2. Let Zz))sn be a sequence of random variables adapted to filtration (Fi)snn If Z < R almost surely, then with probability at least for all n' n.
3 12'
Σ
Z
E, [Z]] 4Rlog(200
(40)
and
n'
E [24] 2 Z + 88
(41)
The following lemma is uniform version of, e.g., Lemma 23 in Foster & Rakhlin (2023).
Lemma H.3. Suppose that is a distribution over Z, and let F C (Z + R) be a function class. We
let N(F, || be the c-covering number of F under the norm p(f. == f
Let D [Z' ..... 7 be drawn i..dd from Then the following holds with probability at least 1-8: 8:
1( < n log F (exp( + (18) inf (log N(F + 2nc], V.EF.
e20
Proof of Lemma H.3. Fix e 2 0 attaining the minimum of log NFF, [|-1160) + 2ne, and let fi, ff be e-covering of J of size J NNFNN,F66 ll- For each [J], define gj (z) fz(2) log E, (exp(( Then, it is clear that E,, (2)
and hence
exp [
Vj e [J].
By Markov's inequality and the union bound, holds that with probability at least - 5,
Σ (z") < log(J/8), Vje [J].
(42)
Note that for any f E F, there exists j E [J] such that p(f.f) S c, and in particular
f(2') log (x)e < 2. + f₃(Z") - log (x,ee = 2e + (;(2), Vi [n]]
and hence Eq. (42) implies that Σ f(Z)) Eexpp log(J/6) + 2nc. By the arbitrariness f. the proof is hence completed.
37

---

Published as a conference paper at ICLR 2026
H.2 INFORMATION THERRETII INEQUALITIES
Lemma H.4. For distribution P, Q E A(x), function f: : - (-B, it holds that
Erlsi-Eqr 1 A Varqlf) D
More generally, for any g:t= + B2 B), it holds that
|p - Eglell - 4v Balg Del
(43)
and
- - Ep[g]| <3 B DE
(44)
Proof of Lemma H.4. We denote P(x) (resp. Q(x)) to be the density function of (resp. Then for any function f X + R.
Q).
(/,(())) (f(x)
Er[s] Eoll]
- ((TT)(P()--((11)) $ / (f(x)- Eq(f1) P(c)) + ((x)) dr
(v P(x) ( ((()))
$ $DD Q)- (Varolf) Ep(s Bo(f)
In particular, when h: [0, M). the inequality above implies that
[p[h] [q/hi| s 2D Q) [Mp + [g[hh) _ (Ep(A)) + [Qlh)) + +2M Q).
and hence it holds that [M[ _ 3 E [4] D3 Q).
Now, suppose that f : X + [ B, B]. Applying the above inequality to h(x) = (f - Eo(f))² e
[0, 4B2 gives
Ep(ʃ - | < 3B B2(l)) + 16B (PP,
(45)
Combining the above inequalities implies that
[E][f]-Eo] [olf] s 4/Varolf) (ED)).))
To prove the upper bound for vector-valued function g X (B))), we can apply the above inequality with (z) == (v. (zx)) and take the maximum over u E B2)) The second upper bound follows similarly by applying Eq. (45).
Lemma H.5. Suppose that & : y (B( with B > and for any 0 E B2(1), mo A(Y) is defined as TolY) xx exp((o(y) 0)). Then for any 0000 0 E B2(1), holds that
By (Ø(g) Exg, [Ø], 0 8- < 15BDKL (TT 11 To)-
Proof of Lemma H.5. Denote ō(y) Φ4v) -E,,, (o). By definition.
exp B (00)). 0 04) )
DKL(TO: || (T)) = log Ey [exp((x ((ō(y), - 87)) 2 B log E yyrr
Note that for 2 -4, we have exh > 1 + 3 Therefore, we have
E Ey (o(y), 0722) 2 15B2 E Bynxg (ō()). gn DKL(TTT TO) 2 log 1 10122
where we use log(1 + x) 2 fr for all x [0, $]
38

---

Published as a conference paper at ICLR 2026
1 PROOFS FROM SECTION 3
Proof of Proposition 3.2. Consider the setting where d = 1, = (0, V = [-1,1). the distribution f is given by ((1) = 1a(0) 2n and the feature map Φ : X x V [[1,,] given by ((0,-) 0, and ((1, V1.h) yy-
In the following, we fix any algorithm Alg (2' x 2) + AIT). Let Pre be the probability distribution of (D {(.' )E) #) where xt n.y mal. x²) are sampled i.i.d. and
Alg(D).
Note that under this construction, IP (pro. (1 =0 11 [T]) > 1 (1( 2- Consider the event
F = xx' = 0 AA E (T]} Then, for any E [[-1, 1]. event A, holds that
P Mle (A E) = E Aie (A E),
because for any 0 E O the distribution To distributions and does not depend on O. Furthermore, for any 0 E --,11], 1].
0) Ber } 68
is product of H Bernoulli
|| (T) (TL
= ((1) DKL (To- (V1.H = x0= 10
To (V = 1 x===1 1))
e
c
= =Hp(1) DKL
Ber
cl -e-
Ber
+ +
and hence 0 0+ (ML( m0) + ( (T) is minimized at A = 0, i.e, for any Ω E II,
H
(Ber
H
DKL( ( (1 + DKL ( 2
2DkL
Ber
e C eec-T)
Ber
2
2n
2n
Therefore, consider the event Ag == {DKL(TO) ( ()) 2 in ). and we have shown that Af C A-1- Hence, we can lower bound
(An))))(A)) PT-1- (A-1) PTIMIe (T) PT) E) + P P
(F) (
E)
E]
+
Ale A ET
2
This gives maxx-E(-1,1) Pag. Alc (DKL(TTO ately.
H
2 }. and the desired result follows immedi-
4n
As remark, we note that the construction above can be modified so that the variance of (defined in Section 4.1) can be bounded as o2 S He-2 In particular, as long as B 2 log H, it holds that
2 0. I 1, implying that KL can converge slowly even when the "inherent variance" a, small.
J PROOFS FROM SECTION 4
J.1 PROOF SKETCH FOR THEOREM 4.1
The basic idea behind the proof of Theorem 4.1 is to interpret the condition C 2 E as an small-ball like anti-concentration condition in the vein of Mendelson (2014; 2017). That is, for models π E II where coverage is large, the condition Cov x (T) 2 witnesses a one-sided tail bound which implies that the empirical likelihood of is not too large with high probability, and hence cannot be
solution.
Let c E (0, 1/2) be the absolute constant in Theorem 4.1, and let C 2 log 4 be another absolute constant. Fix N such that log N 2 4C. For each model æ E II, let Sw(m) == Alfi [M]
Tx(="(>") x(r12*) > N1-2=}| denote the empirical probability that म fails to cover Tp. Our first step is to show via covering and concentration that with high-probability, all क E II satisfy
SN(T)= > 7 covn (T) - Ccoarse (II, N,n).
(46)
That is, a large coverage profile implics that the number of points in the data where क fails to cover To is large. This argument only depends on the covering number coarse log N seale-llading to
39

---

Published as a conference paper at ICLR 2026
the coarse-graine -ranseegr term in Theorem 4.1 because only need show that coverage concentrates, not the log-loss itself.?
We now argue that models with large coverage profile must have low log-likelihood compared to xp. In particular, using Eq. (46), we have
(o(y x²) C
"
aty x')
(T) (TO)
Σ
log
log
V V(-C)
Tly
x²)
TT('' xx))
log Tly xx)
[Sx(T)]((1 2c) log N - C) +
V(-C)
(Ty' x²)
* x²) To xx)
S 4
log N COVN (T) + Ccoarse (II,Nn) On N) + log
V (-C),
(47)
as long as c S 1/8 and log N 2 4C. We view step (*) as using a form of implicit bias in the logarithmic loss: If an example (x',y") has moly pol"l=")/ (((("=) 2 N (i.e., fails to cover To on this example), this witnesses negative contribution of order log N the difference Lu (x) Ln (T)).
Next, using a variation of standard one-sided tail bound for the logarithmic loss (van de Geer, 2000: Zhang, 2006),* 8 we show that with high probability. all π E II satisfy
Tt x²) TD(y" x²))
Elog
V (-C) S Crime nI)
(48)
==1
as long as C 2 log 4. Combining Eq. (47) and Eq. (48), we conclude that all π E II have
CO (To) (T) + Crine n), nn + Ccoarse N,,, n log.N N
(49)
Since the maximum likelihood estimator से has (TD) - (+) < O, the result follows.
To summarize the key ideas as they relate the final guarantee in Theorem 4.1: The coarse-grained term Ccoarse (11, N, n)) enters because we only need to show that the coverage profile concentrates, not the log loss itself. The fine-grane grainnr term Crine n) enters concentration of the empirical likelihood, with the 1/log N scaling arising from implicit bias. The reason this argument avoids dependence on the sequence length H or other spurious parameters that would otherwise affect cross-entropy is that the argument is fundamentally one-sided: the conclusion Eq, (49) only shows that models with large coverage profile have low liglliieldood compared xpp
J.2 PROOF OF THEOREM 4. (COVERAGE FOR MLE)
Theorem 4.1' (General version of Theorem 4.1). Let N 8 be given. With probability at least
1 8, any approximate maximum likelihood estimator with ( 2 max (.) ne sox satisfies
log (II, clog N) + log(5-1)
log Nog + +E + Eapx
CO S
+
inf <>0
log N
(50)
where e is an absolute constant.
In the following, for a fixed threshold C' 2 log 4, we define the clipped log loss
I(((T) = Σ max log ( xx)
n
To [x)) -C].
(51)
1=1
LO(T) Σ max 0, log (By x2)) c..
To x
(52)
"hhe (x) is defined with the threshold as N1-2c rather than N to account for approximation errors incurred by covering. 9That the bound is one-sided is critical, as this allows us to avoid paying for the range of the density ratios under consideration. For details, Proposition J.1.
40

---

Published as a conference paper at ICLR 2026
Note that (T) (mo) (x) LT( Furthermore, since To II, the approximate maximum likelihood estimator satisfies I., (7) Ln (mD) apap and hence
L (T) (+) ne
In the following, we show that (T) can be bounded by a one-sided uniform convergence argument, and show that (T) upper bounds the coverage profile (o) for any II and log N > C.
Proposition J.1. Suppose that C 2 log Then, with probability at least 6,6, 6, it holds that for any E IT,
(+) < log(1/6) + 2 inf (log Nox (log.Ns (II, (II,)) ne).
Proposition J.2. Fix any a E (0, log N-) Then, with probability least 6, holds that
C -2 Lo(x)) + 16 1o (II, (T,
S log (T)
The proof of Theorem 4.1 and Theorem 4.1' is completed by combining the propositions above and setting a 4 log N. what follows, we prove the propositions.
Proof of Proposition J.1. This direct corollary of Lemma H.3. For each π E II, we let (z) = _ max log s(vlz) - and consider the function class F {fx E I1I. Then, ;oty|aj.
N(F, e; [1:11=)) < Noe (II, 2c) for any 2 0. Applying Lemma H.3 with Lemma J.1 (stated and proved below) gives the desired upper bound.
Lemma J.1. As long as C 2 log 4, it holds that
( 1x) Ty xx))
-c}) 1.
E (a exp max log
(53)
Proof of Lemma J.1. We denote - and E = (x. y)
7(glZ) mo(u
2 u
Then it holds that
F x))
c))
B(z,y)/m exp max log xx x
((y x) V
B(z,y)~<m
If(x v) E E}+ VuI((E,,)) & E}
TO x
EnzTo
Σ
VI(y ( [ + V Poo (E)).
(((N)EE
For X, denote F. == {y (2. E E}. By the Cauchy-Schward inequality, we have
Σ M 1y x)
Σ ④ ())
Σ
moly x)) S (Pnnattos)FEa))
{((ry)E
VEE.
DEE.
Therefore, as long as u < (or equivalently, C log 4), holds that
T r) To xz)
-c -c}) S V P (E) + Pm (E9) (E9) 1,
E ((r.y))~m exp max log
where we use 1-p= (p) vp) < 2(1 vp) for any E [0.1].
41

---

Published as a conference paper at ICLR 2026
Proof of Proposition J.2. Fix any N 2 1, 0. By definition, for any II,
n (T((T) max 0,
Toly' xi) TTy x²)
2 (log N- -c) c7ff1 E [n]
To xx)) क (3 xx)
log
2 log N
nllog N C) -Coo (To m),
where we recall that (see Eq. (13))
Cov [T) = 1 f
Toly' xx Tly x²)
[m]
Then, by Lemma J.2 (stated and proved below), it holds that with probability at least - 6, for any E II,
COVN (To m)
1
Cov 1| []
8log(2N6 (11, a)/8)
Rescaling +- e -2a N and reorganizing completes the proof.
Lemma J.2. For any model n, n', we consider the quantities
1 [ t [n]
<'y² (y
N}}.
T'Vy xx)
CovN (|
COV ( [| Pr..
M)
T④y x))
Fix a 2 0 and model म. With probability at least 1 8, for any E 11. holds that
8 120 (II, a)/6)
Co [T)
Cov c
(क [T)
Similarly, with probability at least 1 for any I1, it holds that
(T F) 8log(2N (II, a)/
COV (T [|) < 2 Cov To
a0 N
Proof of Lemma J.2. We only prove the first inequality. Let IT' C II be a-covering of II with [I]'] = (II, a). Then, by Freedman's inequality (Lemma H.2) and union bound, it holds that with probability at least - 8, for any x E II'
Cov N ( || 2 Cov To
ca
( ||)
Estats
where we denote Estat (log(2(1T'/88) Then, note that for any ı II, there exists n' E TT' such that log x)) - log n' (y x)) for Vzz V. know
T(y² xxt)
*ly" (x) T(y² xx)
t E [m]
2 c
2 N
('y zt)
and hence Cov ( m') < COV ( || x). Similarly, Cov th under the above event, it holds that
en N ( || x') 2 Cov N ( T). Hence,
COVN ()
Cov N [|
Cov ( 1²) Sstat
Cov To cs
n (T x) Ettat
2
Since π E arbitrary, the proof hence completed.
42

---

Published as a conference paper at ICLR 2026
J.3 PROOF OF THEOREM G.1 (COVERAGE FOR MLE WITH CONVEX CLASSES)
Let a 2 0, N' 2 I. N 2 2e N' be fixed. By definition and concavity of 0 + To(y x), we know O" is an optimal solution of the following concave problem
E arg max (z.g)=_, log (e(y 2>]]
Ag0
Hence, the optimality of Hr implies that
( E V
log To. (y x>1> 2
vB E O.
Consider the function F(O) Ema
xuly m1 mex (v]2
1, which is also concave by Assumption G.I. For
any E O, (0 - 8" VF((F))
8, -,
E,,,
Va (U (1 |x)
(0 00-0 0 E,, -.-mm, V log To. (y x>)) 2 0.
TT. (y x))
Therefore, F attains its maximum over at 0", i.e., F(0)= S F F(0*) = 0 for any e.
Similarly, it is also clear that 0 + Ei= log To (v x²) is concave, and hence म = Tg. where E θ satisfies
000 7
v log TolY' |xa)
A0 E O. 2 0,
In particular, consider the function
11
no (yy x²)
F(0)=
1
Tzly y² x²)
Under Assumption G.I, F concave, and for any O,
( 0 - Vag(y x²)
0 8. Σ
V log Toly" |x²) 2 0.
Taly'
Therefore, F attains its maximum over at 0, and in particular, F(O*)= < F()) 0. This implies
Toa (y'
(+)) Taa π
(y' xxf)
n
log
log #ly' |x Σ
log no. (y' x²).
(54)
x²)
73y
x+)
In the following, we use that N > 2. Note that . log x increasing for x 2 1. Therefore, Eq. (54) implies that
2 0 for any > 0, and x + x log
(N log - 1) n CovN (Tg- [] S (m) L (T..).
(55)
Then, by Lemma J.2, we have with probability at least 1 - 8, for all π E II,
To- (y xx) 7 x))
log (TT, (TT,
N () IP...
22 N
VT II.
Further, by Lemma H.3, the following holds with probability at least 5: For any 0 E O,
(V x²) (TT) - In (TT log
Tea (yr² xii
To (x)
$ n log Em
+ inf (log(Nee (II, c)/8) + 2nc}
Tg- (y x)]
.20
< inl inf (log ( e) + 2ne): <20
where we use E,,, with probability at least
molyyan
F(0) + +1 1 for any 0 O. By union bound, we have shown that
Te- xx)
20.
IP
Tg (y)
x) > 2a N <
xx)
log(Nso (I,,
inf
Nog (TI,e)
+ te].
f(y
N 200
n
43

---

Published as a conference paper at ICLR 2026
Note that
(y x
NN (T)= = -Pa
2 20 N
iT ax)
<
Paa
TT- (y x))
2 e 2
+ 1r
To x Tga (y x
N²
fly xx)
Therefore, the proof is completed by rescaling N + Ne 2n 2N/N,, , + o and combining the inequal- ities above.
J.4 PROOFS FOR SUPPORTING RESULTS
Proof of Proposition G.1 (a). Assume that B 2 log(5n) and n 2 d 2 2. Consider =-1, = [d]
and let the feature map be given by o(y) = Be, for E V, where (eiy...,ea is the coordinate
basis of Rdd We consider + {0 E Rdd |10||x S ]]. and we set
0+ log( (n)
Σ ej
e1
2B
1=2
Then holds that
4n
(1) d - 1 + 4n
To(
Vy > 1.
d 1 4m
Given the dataset D = {y' ......... .. we consider the random variables ny = li [n] y = yH-
Note that under D To. it holds that
nld d --1
E
Σ n
IF
Σ Th
+1}
d d 4n
particular, with probability at least 0.5, it holds that Σy>1 my {y [d] : ny 0} has cardinality least d-1
S 2
d=1
i.e. the set 3/
In the following, we condition on this event analyze the MLE By the definition of MLE,
E arg max -n log
Σ
B8,
B
my
yEld
VE|d]
e nay
We denote Py == TT(1)
Then, the KKT conditions imply that for each E [d]. either
Eiera nn,
Py ny Og -1 and Py 2 Ty. or Oy 1 and Py 2z In particular, for any Vo.
Py > 0
n. and hence it must hold that Oy = -1 Then, because Eye(d Py = = Σ ΣyE(a π
ny
there must exist j E [d] such that Pj
< By and by the KKT condition we have 0, = 1. Therefore, 2B
for any E Vo. it holds that Py, S - S _ and in particular implies that
s>(u) mg(y)
Antd-1 2 B This
(o(y) Tg(r))
d --1
d -
Cov. (TT)
F To
" Pro (26) 2
2(d + 4n)
10n
This is the desired lower bound.
Proof of Proposition G.1 (b). Let e = coy 4 and p = cuz²
1o for a sufficiently small absolute
constant co 0, X'
(0, 1,... d), 10, and the distribution μa be given by m0()) p.
((1) = 1ي
Let the data distribution To be Tp( i) = Ber(1/2) for E [d] and To 0) = 1. For any
0 E θ = (+1,-11 -1dd, define Ta as
( () = Ber 1
Te( |t) Ber 1+ ct. [
Vi E [d].
N
2
44

---

Published as a conference paper at ICLR 2026
Consider the model class II {mо} U {Io E O). Note that for any 0 E O COVN(TTo)) (0))= = p.
Then, calculate
[T( (TD) C(O, log N
Σ Iii log( + clc) C(i, log( log(1 - <0:i];
iE(d
where we denote C(x,y) y) 1it E [n] (x', v²) (r. We further write C(x) = C(2, 0)
Car, 1). Taking maximum over 0 E (-1, 1] gives
max Ln (TO) (Tp) IE0
l+ O]] T-
l(g log N Σ
C(i, -(i,))
C(i) x²)
iE(d)
C(0) log N nE² +
E
Oiii,)) C(i,
2
EE|d
In the following, we denote A, = C(i, 1) - C(i, 0) and A == Eiejd A. Note that for any i [d]. condition on C(i), A is sum of ((i) i..dd random variables drawn from Unif({ (n,(((-1,1)) and hence
C(i)
B[( C(t) C(t),
E|A,1 C()) V
where we apply Khintchine's s inequality. addition, we note that C(i) B(n, q) is a binomial random variable, where q = " Hence, E[C(i)] = ng. and to lower bound CC),. we invoke
Lemma J.3 (stated and proven in the sequel) to show that E VC(I) 2 Vng 1 (because 2d and hence ng 2 1). Therefore,
1-0
Pg
2ng
d n4
E|A|
Σ E
N
Σ
El
C(1)
V2
2
sE|d
EE(d]
and we can also bound E( < a ie(d by Paley-Zygmund inequality, it holds that
I(A,) d d
Lield] B[C(i)) dn(1 - -p) = dnng. Then,
(1 61-
P(A > (E(A)) 2 (1 b)²
Ab E [0.
T A
We choose b = v 0.88 to be a numeric constant so that PPA > EE(A) 2 0.11. By Markov's inequality, it also holds that (0) > 100np) < 0.01. In the following, we condition the event E = A 6 [EA n {C(0)= < 100np} (note that P(E) 2 0.1). Then, we have
beynd 100mp log N ne² 2 0,
max Ln (Te) Em((T) 2 C(() log N ne² c - BE
8
as long as cp < 10-4 This implies that there exists Cov ( 2 p. This is the desired lower bound.
such that = To. and hence
Vari2 Lemma J.3. For non-negative random variable Z, it holds that BlVZ| 2 E
2[E
Proof of Lemma J.3. Note that the inequality Vu > 3u-u² holds for 0. Setting u E and taking expectation completes the proof.
K PROOFS FOR AUTOREGRESSIVE LINEAR MODELS K.1
ORGANIZATION
This section contains proofs for all of the results in Sections to 6 concerning autoregressive lin- models (3). We begin with the proof of Theorem 4.2 (MLE for autoregressive linear models).
45

---

Published as a conference paper at ICLR 2026
We then present the proofs for various SGD methods, starting with vanilla SGD (Proposition 5.1; upper and lower bounds). followed by normalized SGD (Theorem 5.1), test-time training (Theo- rem 6.1), and expert-guided gradient normalization (Theorem G.2). The final subsection provides an additional lower bound, showing that the dependence on the parameter of necessary high dimension.
Throughout this section, all upper bounds are derived under Assumptions 2.1 and 2.2, i.e., we assume that + C B2(1), Φ X V* (2)) and mp we- is realized by some parameter 0* E e.
Notation and preliminaries. For any S X x V IR and dataset D (( vi.)))) we write
ED[s] ==
For notational simplicity, we denote
Te (z, 31-A-1) = 18 (oa v1.h))
and
o" (1r, 6(x, v1sh) /e. (x., y/:h 11:
Vara, (x, (1:h-1)== -1 (x, 31?
Then, by definition,
V log To (y1. (x)
Σ (0(x, v).h)) (0, v1.h-1))
H
Φ (a,, V1:h)) + Σ
(oo (x,,
δo(s, (1:h-1)),
(56)
h=1
and it holds that o3 = ENT E/m == Var (2, In addition, we write
cØ (a, V1:h-1) (To( x, (1:h-11) mo(. (1:n-1))
(57)
For any 0 E O, the key quantity of interest is Dseq, (TD To), defined via
Dseq, (TD me) Er, min log N,
Σ
DRL (To(: | x.311n11 TT(- x, (1-k-1))
h=1 H
Eo (x, BLA Σ
Egg min
log N,
h=1
By Proposition F.10, it holds that Nov
2 S log N
D Dseq (To mp).
Further, by concavity, we have Eo (z, 11:h-1) S (1, 5o- (2, (1).v 0*).
(58)
By Lemma H.4, it holds that
11 (I,, (1 50(2 Vi:h x)ll 4 Var (x, 11:h-1) c0( 8Beg (1:hh))
(59)
K.2 PROOF OF THROREM 4.2 (COVERAGE FOR MIL FOR AUTOREGRESSIVE. LINEAR MODEL.)
We prove the following slightly stronger result. Theorem 4.2 follows immediately by combining Theorem K.1 and Proposition F.10.
46

---

Published as a conference paper at ICLR 2026
Theorem K.1. Suppose that Assumption 2.2 holds. Then the MLE ā achieves
03 log N
Ba logg N
Ep Ep(Dseq. [1 S
+
n
for any parameter N 2 2, where the divergence D=eq.N(-1) -I) is defined in Proposition F.10.
We begin with central technical lemmas, which are proven in the sequel. The first lemma is consequence of the fact that the MLE मे "g maximizes the empirical likelihood, i.e.
arg max Epllog Te (y x)), aEe
(60)
where we recall that for any dataset D {(x', Vi.H (..:)/esi Vest write ED[/] = # Ei= M(x vi.H) for any f X' V* - Lemma K.1 shows that in expectation, sum of per-step conditional KL divergences between To and से is bounded (this does not imply bound on sequence-level KL. divergence, since i is dependent the data D).
Lemma K.1. Recall that we denote cœ(x, 11:h-1) = DKL(TO(- x, (1.h-11 1) mo(. 11sh-1)). Fur- ther, define
E1 = ED
Σ E v1
(61)
Then holds that F
20 VT
Define A == log N. The next lemma is a uniform convergence-like argument which shows that the quantity Ej above-when truncated at certain level A-conccetratess around its expectation up a multiplicative factor. This argument is inspired by the fractional covering method introduced in Chen al. (2024a); Chen & Rakhlin (2025).
Lemma K.2. Fix any E (0, 200B 8 E 1). and let J = =exp( ( + log(1/8). Let O' == (fI.... O5}, where 0,,... O2 N(0,, A²1) are sampled i...d. Then the following holds with probability least 8 over the randomness of e' and D:
(1) For any j E [J], it holds that
H1
$
(z, BLA-1) 8A lo (47/6)
Bm, min A, Σ
co, (x, BLA-1))
< 2ED min A, Σ hml
Am1
(2) There exists E [J] such that
}
H
Ex, min
A,
Σ
ca(2 v1-h11)
< 2En, min A, Σ
co
(2, (2,31:h-1)
+ CA
(62)
hml
hml
and
H (π, 31 vn A,
ED min A, Σ
cg, (0.
2 min
h=1
h=1
(63)
CA
2E
Var no (2., 11:h-11
where C' 1000 is numeric constant.
Above, the distribution of т. under ~ (0 A²1) can be viewed as a fractional cover for II in the sense of Chen et al. (2024a). In particular, working with the fractional cover offers the following technical advantages:
The fractional cover (0, A²1) incurs error o (see Lemma K.3) that depends only the variance at the ground- truth parameter 6*, This contrasts with classical coverings, which enforce uniform bound for all 0 e.
47

---

Published as a conference paper at ICLR 2026
For θ = (1)),, the Loe covering number of II (ef. Definition 4.1) scales with the dimension A standard approach to deriving dimension dimension-indeeendde independent bounds is to apply symmetrization techniques and use data-dependent L2 covering show uniform convergence. In contrast, fractional-covering approach avoids the (technically subtle) symmetrization step because the cover (01....,0J) 0,) ~ N A2T) is drawn independently of the dataset D.
Completing the proof. Equipped with the lemmas above, complete the proof as follows. First, we condition on the success event of Lemma K.2, and let j E [J] be an index such that (62) and (63) hold. Then,
upper bound (recall that A log N and N || .) is defined in
Proposition F.10)
N1
Dseq, N (To (ă) Emg min
A,
h=1
Ex min
A,
Σ
ED, (Is, 11:h-11 1)
+ CAPO3
H
16A (GAlog(4J/0)
S 4 min
A,
Σ
<o2 (I,, 92
+
+ C2
n
h=1
< EE min
A,
Σ
cg \1:b-1))
h=1
16A logg///))/0) CA
4CAAEE
Σ
Var (2t, 11-h-1)
n
Li=1
where the first inequality uses (62), the second incquality uses Lemma K.2 (1), and the third incqual- ity (63). Therefore, denote d²( ED Σ Varm, (I, v1-h-) and we have shown that for any 8 (0. 1), any A (0, 2008l holds that
D Deeg.N TD TT) 2 Ci E1 () A( 1 A
+ +log(1/0))) <
where C1 > 0 is an absolute constant.
Since 5 E (0,1) is arbitrary, integrating the tail inequality above yields the following bound on the expected value:
+ A 1
A (111)
E E[Deqq (1 [T)] < C1
E
A A3E(o2(D)
07 +
A
<
2
vA < A <
n
200B
1/4 oen
Choosing A min 2
completes the proof. The coverage upper bound follows
immediately from Proposition F.10.
K.2.1 PROOFS FOR SUPPORTING LEMMAS
Proof of Lemma K.1. Recall that म Tg. where Ø = arg maxage ED (log ve( x)). Then by concavity of the likeliid,,,,d, we have that
E V log (yne 1)) 50,
1e e.
Using the expression (56) and 0* E O, we know
Σ (d(2, vn.h 2 1p 1)
Ep
0.
48

---

Published as a conference paper at ICLR 2026
Therefore, combining the inequality above with Eq. (58), have
H
11
Σ
c
= ED
DKL (TT( 1 11- 20 11:h-11))
h=11
E
Φe- (ar, (i:h111h)) 0g(2, 11;h-1); -0)
S ED
Σ
co- (Ir, 11:h-11) - ((0,, y1.h): 0 03
H
- 2
Ep
Φ (z,, v1sh))
=: 1i
where we recall that x (x, V1-h) = Φ(x, (1:h) - ã (x,, V1:1-1-)).. By definition, it holds that E,, (r, y1a)) x, 11,hh 1] = 0, and hence
E(EE) E
E
Φ" (a,, 11:h))
A=
H
H
o?
E mp
Φ (ar, v1.h)
Em
Σ
o (a, )1
This gives the desired upper bound.
Proof of Lemma K.2. By Freedman's inequality (Lemma H.2) and the union bound, follows that (I) holds with probability at least 1 2. In the remainder of the proof, we prove (2).
Define the following weight function a = og X' x V + 1].
Ejsh-1 E,(2, Myy 5 A, E Ejen M1.)) 2 A,
ag(I,
A 1 sg(m.)ns))
otherwise.
c2(2,01a11)
We also define F(a,b) Ja a. The properties of F(s.)) and the weight function summarized in Lemma K.4 (stated and proven in the sequel).
Then, by Lemma K.4, holds that for any O,
H
E,, min A,
Σ co(x, w1m-1)) < 2 Em, min
A,
co ((..Viih))) m1)
Aml
h=1
E
Σ
a1 Ul F(egl, vo(o.
-=11
and
H A,
H
Ep min
Σ c0(a, (Ln-1) Eoola h=1
< 2Ep min
A,
Σ
c (1(=,8an11)
y
h=1
a(x, 1 (i (lee). F(da 00
+ ED
h=1
Therefore, it remains to control the error ΣH_ az. 11 V1sh-1), co (z V1:h-1)) under both E, [ and Ep|.]. We next state the following lemma (proven in the sequel), which leverages
IInspredd by the analysis here, we also adopt this weight function in the SGD update (35) with the truncated stochastic gradient estimator.
49

---

Published as a conference paper at ICLR 2026
the structure of Gaussian distribution. This result can be viewed as fractional covering number bound (Chen et al., 2024a) and hence generalizes the argument of Chen & Rakhlin (2025, Proposi- tion C.4).
Lemma K.3. For any K 2 1, E (0, m e B2(,), distributions PL PK over Z ==
X' x Vr, and weight function 2 + [0. 11. it holds that
1
llg Po Per (Vi E [K], 2x (<=)F((e(=))))))))) Eo (z)) S 70 Varmo (z)) S
+ +2,
2
where we recall that F(a. b) a bl 91-33 a.
In the following, we apply Lemma K.3 with K = 2, parameter 0 = 0, weight function and the distributions pi, P2 defined as follows:
Let p1 be the distribution of x' = (2, U1:h 1) under x ~ 14, V1.H ~ To (- æ) and h ~ Unif
Let P2 be the distribution of x' (x', vinh-1) under Unif and h Unif((H).
By definition, it holds that
H cx
1
EzppI (o))
E
H
Banpi Var
(z)
1
Em
H
Varm, (2r, 11:h1
o² H
H
1
Ezopp (o))
ED
H
H
E Ezmpa Var (z) = H
Ep
Var (I. (c..---1)
Lh=1
Now, consider the following set for any 0 E O: + = {Vi E (1,2), Ez-p, (x)))((
egr S 300 3 (a)
By Lemma K.3, it holds that
(0)) Pe E0T) 2 exp (- 1 N2 22).
1
40 VOE E (0,
200B
Therefore, we have
P(vjEMM.00 $ $ 10) = Po.... 10..0,0~~0000111
Vij E [J],0, $ 03
< (1 - (0)) S exp -Jq(()) S 5
and hence IP(3)
E [ E ) 2 1 2. The proof Lemma K.2 (2) thus completed, as Eq. (62) and Eq. (63) hold for any j E [J] such that f₂ E O
Proof of Lemma K.3. We first fix any h E [H] and = = (x, (1sh-1) E X' x ph-1 and analyze the behavior of log To (yk 2) under 8¹ N000 T)...
By definition, we have To (Bn 2) &yn ( 2) exp((0' - 0, Ø(2 (2,2())))),)) iee, log nor (yn =) - log Tolyn | 2) (0' 0, (2. log (yn exp((0' 0, (5(2,,n)))) Therefore,
(0(2) eær (2) = (xo(n = 2) || TOe Te(yn = =1 DKL (To = 12) || T (yh = 2)) (Exo(-1e) (0' - 0, 0(2 yn)) - log (- exp( ((0((i 0 612 vn)))
(0' - 0, Φa. (2) Ōo(2)) log Byn exp (00(((00 0, 0(2, vn) oo(2))),
50

---

Published as a conference paper at ICLR 2026
where we recall that (0(2)) Byn (y) yn)]. In the following, we denote Φ0(2, yh) ((2, vn) (2) and
Egi (z) == log En (-(()))) exp( 0, @0(2,8n)),, Eg (2) = (0' 0, do (2) Φu(z))
We first bound (z). By definition, we have E (z) DKL(TO(: 1 (x) Tar (- 2 0. Further, using Jensen's inequality, for any 2 Z, we have
Ber-N(0,221) & log ~(0.33)) B Pynnmoo(ls)loppl 0, un)))
log E yn x exp
E By (1) (z vn))
where the last inequality follows from et 11 for E [0, 1]. Further, using Lemma H.4, we have
E ( vn)ll
mn) 6b())
< 3 33 Wh) - ou (=)11 + 16B (6B2DKL(0> (To(- 2) me(- |))
33arr (z) + 16B2
Next, we bound E (z)]. Under 0' ~ N(0,A21). it is clear that (0' - 0 (2) - -0(2)))
(0 A3 (z) (0(2)]1 for any fixed z. Therefore, it holds that
2
Eo-N(0,A21) E3 (x)
11 (2) - x0(2))
π
A
(4v Var (2) c0(2)) + 8BEe(2))
+
SBA
co(2) + 32K A Var (2),
8K
where the second line uses Eq. (59).
Combining the inequalities above and taking expectation of 2 ~ Ps. we know that for i E [K], holds that
Ee-~((0 ( BeNpi [a(z) < A Eampi (Var (z) + 16B² (z))o(())
EV Benpa [a(2) Ex (z) < Ez-p.
32KA (2)
1
SBA
a(2)co(1)
8K
and hence by Markov's 's inequality and < where we denote θ = 0, and
B holds that = Fo For (0' #0T)
o7 = =
( Rd Egup, ((=))e((=)) - ear (x)l 2 Esup, (6K
( 66K++64 (z) 1 (
Note that DKL (N A21) (1, A21)
|1112
< ZA Hence, by data-proggssing incquality, we
2A
can bound == Po-N (f' ( as
2 DKL (NL(WN(BB A2T) [[, A²π) 2 (Ber(p) || Ber(g))
2
P + (1 - p) log 1- P
= pooog
> log(1 (1/) log
1-4
This implies that - -log q 5 3z + 2, giving the desired result.
51

---

Published as a conference paper at ICLR 2026
Lemma K.4. Suppose that a1,""
by 2 0, 0. Define F(a, b) la bb -2a Let
1,
Eish 0j Ejer 03> A,
ah
0,
A-Egeh
otherwise.
Then clearly on E [0, 1] wh [H], and it holds that E Ondh min A, E an and
H
H
H
min
A,
ah
2 min
A,
bh
+
2
on Fan, bn):
1=1
h=1
h=1
and
H O F(an. bn)-
min
A,
bn
2
min
A,
ah
h=1
Proof of Lemma K.4. Fix the sequence "
We first prove that
H
a4oA min
A,
(64)
h=1
To do so, we consider two cases.
Case 1: E an < A. In this case, aj = 1Ah E [H], and the equation holds trivially.
Case 2: Σ ag A. In this case, we let l E [H] be the maximal index such that aę 0. Then,
by definition, Ejet aj S and Ejsr aj A. and at A-E aj Hence.
at
Σ agah
ahap
Σa
aj arar A.
jel
We also note that from the proof above, we also know that for any sequence (c1,"" ,CH) such that ch 2 an for [H], we have
H
min
A,
Ch
Cnchh
(65)
1=1
1=1
Equipped with these results, prove the inequalities in the lemma statement. We note that
H
On F(an, bn)
Σ
on an
bnl
Og
2
or equivalently,
H
anlar - bbl = Σ ag F(a,, bn) + min A, Σ ap
h=
Therefore,
H
H
H1
H
min
A,
Σ
aj
ana S min
A.
by
an lan bul
hm1
H1
Ff
H
1 Or F(an, bn) + min
min
A,
bn
A. h=1
h=1
h=1
52

---

Published as a conference paper at ICLR 2026
Re-organizing yields the first inequality. Similarly, have
min
A,
bh
min
A,
Cah + an bnl)
an (an + lan 6%.1)
h=1
h=1
min
A, A=1
ah
On F(an, bn)-
h=1
The proof is hence completed.
K.3 PROOF OF PROPOSITION 5.1 (VANILLA SGD: COVERAGE UPPER BOUND)
We first invoke the following standard lemma.
Lemma K.5. Suppose that the sequence (0', (1)) satisfies 0
(roje ngg') for t 2 1.
Then it holds that for any A E O T 221,
T Σ
(-g'8 8 - 0r) 5 10 - 1
1
Σ
T
(66)
2n
2 1=1
Specializing Lemma K.5 the SGD update (8) and taking expectation, we have
E
V log no (y' 1x), 0 8*)
S
2
1
V log To (V' x')l?
(67)
2
Note that (x', 3²) 0r To. and hence
E V log naz (y² x²)) 0*) = (
V log net (y |x:]] = V D molle=or.
Further, by convexity, it holds that for any 0 E O
G(0)= E., K|
log (e x). 0 0* = (V (T)) 0 00) 2 DKL (TD TT).
Therefore, have
T
T
E
DKL (TD TT()
<
E
Σ
G(
2
+
E
E (2.o))ao Vog To (y x)]
17
On the other hand, using the fact that log To |x) is concave and (HB3))smooth (i.e., - HB2T 1
2 log (t xx) = 0),
|| log T x) v log no- (y/ x)] < HB² (0 8*, V log To. (y x)) - V log To(y) x))
Taking expectation of (xr, v) To and using the fact that E,,, log no (y) x) 0, we have
Em log ( x) Vlog To- (y x)]1 < HB² G(Ø),
AO O.
Further, note that E.g V log T (y x) of. it holds that
Exp V log To (y =)|| < 2 + 2HB² G(O),
(68)
Combining the inequalities above, we can conclude that
T
T 2 + 1HB E
Σ G(O))
< "
Σ G(O))
+ tIToO?
We conclude that as long 17 S 2HE4 it holds
T
4 + 2T >
Σ G(O')
2 EE
Σ
DKL (mo TTc)
"
This is the desired upper bound.
53

---

Published as a conference paper at ICLR 2026
Proof of Lemma K.S. A standard result (e.g., Hazan (2016)) is that because the projection operator Proje is an contraction, we have that for all E [T]. the update satisfies
| - 0 2 ||0 - 0-1
1 0*l 2 110 + "19" 0-
(69)
2ny( g 0' 0*) - 1
Summing this inequality across steps 1, T. telescoping, and taking expectation, have
E - 8*) 0 0- 100 0 1|2
7
1 II?
(70)
2m
2
This gives the desired upper bound.
K.4 PROOF OF PROPOSITION 5.1 (VANILLA SGD: COVERAGE LOWER BOUND)
In the following, construct X [ +00) L (-, V {-1,0,1} and e B2(1) with
We fix parameters 2 B 1.
Construction of d.
We first construct map X x V + R² follows. For any HB
8
R and let define On
2nHBB11
wn 0) [1:0],
u(n, 1) [ani / a v(n, 1) [a,; / a?,
We further define
1
v(+ a) B]; v(-,0) = T (0; Ba]
Va - -
For & E X, 3ish E Vh, we define @(r, v1sh) = Bu(x, 1/h).
Under this construction of d then prove the lower bound by considering cases based on the value of 17.
Lemma K.6. Suppose that v HB. log N S HB and B en log(TH) for a large constant
8
en > 1. Then, with the distribution being supported on x and 0* [1;0]. the following holds.
(1) The variance of such an instance is bounded: 1.
(2) There exists 0 ② such that with probability least 0.5, the SGD sequence (0') satisfies ( 2 1- 2F for all E [T].
& Lemma K.7. Suppose that 1) S HB. log N < HB
5 8. and B B 2 cp log(TT)) for large constant
CB > Then, there exists distribution μ and E such that the following holds.
(1) The variance of such an instance is bounded: 1.
(2) There exists 1 E O such that with probability at least 0.5, the SGD sequence (0') satisfies
HB
(TTr 2 emin
Vt [T].
T- B2
The proof of Proposition 5.1 (lower bound) is then completed by combining Lemma K.6 and Lemma K.7.
Proof of Lemma K.6. Fix the parameter
We denote n = n HB and a = an = HB
8
Denote
Up = [1:0].
2i = [a: / 11-a7 a?],
1-1 = [a; / a?].
In other words, for any 0 E O Vi.H na(: x) sampled i.i.d. with Po x), where Pe is
defined as Po(a x)) Ea'ey (op
exp (up(B(((2.))00)
54

---

Published as a conference paper at ICLR 2026
Under construction, we have
exp( ((, vyy.))
Ty
17 92:h-1) = EeEy Cxp(B(O, va))
=: Po(yn).
We study the SGD update starting from 0 = 01 By definition, 6(1,, 31sh1) = Bo(n 3hd, and hence
H
H
V log To (11:1 (n)
Σ
Bu(n,
E [Bu(n. a~Po
B
Σ h=1
Uyn
E
[va]]
a Pao
In the following, we denote
F(H.n)
1 H
Σ Uga
F(O) =- E [Va]
Eeev a eeppe va))
am28
Enev exp ((0,va))
h=1
Then, the SGD update can be written as
u' = 0' + n F(ei.))=((01)
0 = (roje(
We make the following claims.
1 To. Claim 1. For a E {-1,0,1) 13 and 110-u,) vall S [[F(0) vall 2e1.
it holds that Pa(a) S 2e- R/4 =: 6 and hence
Claim 2. Suppose that S min THH 5 }. Then it holds that o, S 1. Further, with probability least 0.5, it holds that F(vi. H) eo for all e [T]. In the following, we condition on this event.
Claim 3. By definition, for a E {-1, 1). we have I| + 7/000000000 - va)ll 1-- and 8- + 01
8-1 00-
Claim 4. Let 16 Suppose that to- Then for (-1, if |00 v. 6. then it holds that ||001 - -v-all S
Claim 5. Suppose that 61 S 2T and log N HB Then COVN mo) 2 1 2r for E such 8
that min m - 11, |10 1-1ll < 16.
Combining the above claims, we know that there a constant C such that as long as сB (og(T'H) it holds that a, < 1. Further, under the success event of claim 2, it holds that for
[-1,1), 110-~va]] S i for all f [2]] such that 2 t a. Therefore, by Claim 5, this gives T6
HB Cov (πØ²) 2 1 as long as log N < 8
Proof for Claims 1-5. To prove Claim 1, we note that (0,va)
va) 2 1 - ||0 - -vall 2 15 and for o < 11 Therefore,
16
(0,v,) v)) < (Va, v2) |10 va
I < +
16
Eijo B(0,.s)) BB(,vo))
2
1 1 Po(a)
BB/1
C1.
This completes the proof of Claim 1.
Next, we prove Claim 2. Recall that 0* = [1: = Up- By Claim 1, we know 1 Po- (0) S 61, and
hence Varawrg- [va] < 56 This implies o² HB [va] 5HB2, < 1.
We also know Pm, (yh 0 wh [H) (0) (0) > (1 c1 )
Hę1- Therefore, taking the
union bound, we know Ply = vh E [H]. t E [T]) 2 THe 2 1 This completes the proof of Claim 2.
Furthermore, for any 0 such that min v2 100 vー1 ||} < +/- as long as log N < (log(1
G) log(G)),, we have
Cov. (To mo) 2 (1 1 2n
1
[(H log (0) H log (0 2 log N}
2n
55

---

Published as a conference paper at ICLR 2026
In partieular, this is ensured when log N HB This completes the proof of Claim 5.
Claim 3 follows immediately from the definition of a, vo, and E-1-
Finally, we prove Claim 4. Recall that u' = 0' + π F ((iin)) - -(())) Then it holds that
||u' - (1) - 10--ll ||us 1700 + (17 1jval $ 110 Uaa + 7/ (()...)) coll + =l)F(()) Va
< 2 =:
In particular, it holds that |||2*|| - (7--1) 1)1 < é' and hence || 2 n -
e' -11-26117-112
1. Therefore, 0+ Proje(2') T and can bound
110 v-al
ux² (77 1)w-a
+ v-a
1
(1)
D)u-al
In7
1
||4|||
11'
2
4e'
4
8E1 Se.
1
f
1j
Proof of Lemma K.7. We again denote π = HBn S 8.
We choose 8* = [32 }1, and let the
distribution μ be supported on { {-+
BH
m(+) = 1 p1-) min
512enB log N
Note that for Z {-,+), To(1 2)) c-B/2414ch77 5/2 and hence - To( = x)) s 2e B3/2 Therefore, similar Case 1, we have the following claims.
B/2
Claim 1. Suppose that B 2 cg lo for large constant CB > 0. Then it holds that S 1, and with probability least 0.5, it holds that E I(x² +] S 4Tp(1), and yi for all h E [H]+ f E (T].
In the following, we condition on this event. We choose 00 = fr I #.
< z such that e Flo N and we let
Claim 2. For any 0 E + C IR2 it holds that - Po | +) S
2 (where w|l] denotes the
2
first coordinate of vector w R²2) Hence, when x" +. using yi 1. we have log ग T x²) [2) = 0 and
4HB e
< V log ( [1) HB(1 E & 2HB(1 P&(1 ++) &
Similarly, when x² we have
4H
v ( xx)[1] = 0,
0 S V log (o |x <
Then, combining the inequalities above with Claim 1, we can inductively show that for any / E [T),
AnHB BuB
16 BerB s(())
512eBn log N
0'[1 [[[]-00 [ < 1x +}
< <T ((+)
BH
B
Therefore, we have [1] < r for < T. It remains to prove the following claim.
Claim 3. Suppose that aate 81]BB
H Tlog N.
Then it holds that Co m(-)
56

---

Published as a conference paper at ICLR 2026
To prove Claim 3, note that similar to Claim 2, (yh wh [H] I +) 2 Further,
log (b(1 = 1 11+) +) 2 log(1 2e- be2 2 -3e -3e- and log गo = 1 +) S 3
B/2
Hence,
for y E yH being yi 1 for h E [41, it holds that
log (b +) log To 1+) 2 H
3 3c-B/2
2 log N.
The immediately yields
COVN(TT8) 2 m(+) Pr y² x== +) 1u++)
2
K.S
PROOF OF THEOREM 5.1 (COVERAGE FOR NORMALIZED SGD)
We denote M := log N. We analyze the normalized SGD iterates assuming 8BM and An
M v
1
16
Denote
(): = x+ (l]00:))
Then the normalized SGD update can be rewritten as 0++1 = Proje (0 + n))(" Specializing Lemma K.S to the normalized SGD update and using θ B2(1) yields
T
2 + [0'00; (- - 0-)
Σ
1
t=1
Taking an expectation on both sides and noting that D ~ Tp is generated independently, we have
2 S + n
E
E E (- 0*)
ED []1
(71)
In what follows, prove number of upper and lower bounds for the expressions involving ğ(0; D) above, then combine them with Eq. (71) to complete the proof.
Intermediate bounds. Recall
that
write c B1:h-1)
DKL (-o( st, 11) || (- x, Vish -1)) Also recall that we adopt the notation that for any function f and dataset D, we write [D[s] == কे (mmmaae f(z w1.H).
Denote (recall that Deeq. (- is defined Proposition F.10)
"
c0 (D) = Ep
Σ
cz 31th-11)
Ag E,,, min(M, co (D)).
Using the key structural result in Proposition F.10 (recall M log N), can bound the coverage in terms of the expected sum of stopped KL divergences as follows:
2
CoVN S
M Deeq N (TD To)
En min M. Σ co(a, (a(2,V1nh11)
M
(72)
HH
2
<
M - Exan min M, Ep
Σ
BLM-1)
M Ag.
h=1
Therefore, remains to derive upper bounds on Ag for A E {0), ... 0r).
Lemma K.8. Suppose that 8BM. It holds that for any 0 E E,
E [00: <
2A
4Mo²
6.
M
X2
AV
57

---

Published as a conference paper at ICLR 2026
Lemma K.9. Suppose that À 2 8BM. Denote Ag (-g(6; D), 8*). Then it holds that for any θ,
240B
Mox 2
Ag - AAAg +
+
K
Putting everything together: rewritten as
Under the notation of Lemma K.8 and Lemma K.9, Eq. (71) can be
T
2
Ag
+ 2
E
ED
10(0)
(73)
1
Applying Lemma K.8 and Lemma K.9, we have
1
240B
Mo.
2
84
E
A
61
E
Age
T
K
T
16A
4n1i
+
ED
((00D))1'
Tn
T
16A
811A
16
4nj.
E
A
0r
+
Tn
MT
VK
where the first inequality uses Lemma K.9, the second inequality follows from Eq. (73), and the last inequality uses Lemma K.8. Therefore, long as v7À
< M
it holds that
16"
B K
Mo.
A
M
E
Ag
+
T
T
X
VK
Simplifying the upper bound. In the following, we require n) S it holds that
1288 and choose A =
M Then,
16y
A g
B K + (110.)
M
no.
B K
+ (m1a))
+
M Tn?
T
Tnp
VK
1/4
where we use AM-GM inequality and B 2 1. Finally, we may choose min
M
128B
T.T
Recall that M = log N, and hence our choice of gives
T
o log N
B log
N
B
E T
D Diaq. N (TD (&x)
5
T
E
gt
S
T
+
T
+
K
which implies (by Eq. (72))
T log N
o²
B
T E
CoVN (
D Deeg. (TD ((:)
Tot))
+
T
Toog N
T
K log
This is the desired upper bound.
Remark K.1 (Comparison to standard convex optimization analyses). On technical level, we find the proof of Theorem 5.1 to be interesting because it does not pass through KL divergence as an intermediate quantity. More broadly, we do not know how to derive the result as an application of standard analysis techniques in optimization (e.g. via a gradient dominance or PL-type condition), but would be interesting to see if there is a connection." It
We further note that the inherent variance of corresponds to the gradient variance at the true parameter 0*, and hence is tighter than typical analyses that depend global notions of variance.
58

---

Published as a conference paper at ICLR 2026
Proof of Lemma K.8. Note that |]((; D)Il
- min { 1
Recall that
@(0; D) = Ep V log (o 2)
V log (o x)) =
(0(x, vish) - 50(2, (1:h1)):
with the notation introduced at the beginning of Appendix K. We decompose ((;; D) by introducing
B1th-1) (20(2,11111)) Φ0- (z, (c..na--))
To (D() = DD
Σ
(do
(74)
and
=(D) ED
Σ Φ² (a, (1:h1)))
=ED (o(x, y1:h)) 70- (ar, (1sh-1))
(75)
Then, by definition, (0; go (D) =D)) In the following, we first analyze l19o (D)l and
||x[D]] separately under D = (i.") [e(K] ~ and summarize the corresponding upper
bounds on in Lemma K.10 (stated and proven in the sequel).
Now, using ([(0:D)I & min 1, [000
we know
10(0
|T(0 < 1lE (D) > 3 + (D) S M}
< {D) > M + (D) S M (4Vo2(D) co (D) + 8Be (D)) 1 (D))I
M min cu (D) + v Vo2(D) min M, c0 (D)) + I
where the second inequality uses @ | - 11 (D)Il + [x( and Lemma K.10 (2), and the last inequality uses À 2 8BM and M min M. (D)) = when ca (D) M. Taking expectation of 'D ~ TD. we have
Em 2)12
< Em, min min,M,cs (o + }
4
Em Vo²(D) co + d.
M
AvK
S
Em, min (M, to (D)) +
maintMee,
40,
v E min,
a. XVK
M
A
40.
/Do +
M
AVK
where the second inequality follows from Cauchy-Schwae inequality, Lemma K. (3) and the fact
that (D)] = o.. By AM-GM inequality, it holds that X VDo S + 4Maz and the desired upper bound follows immediately.
Lemma K.10. For any 0 E θ the following holds: (1) It holds that
(go (D),0 0r) 2 ED
Σ -=n
ca co (al,,una-)) (D)
(2) Denote o²(D) == EP Ele Varn, (I,, 1/:1111 : Then
10 (D)Il S 4402(D) co (D) 8Bea
(3) It holds that Ep [|()1 and
1
30B
ED (z(D), 8*)
55 (D)
<
a.
2
K
59

---

Published as a conference paper at ICLR 2026
Proof of Lemma K.10. Lemma K.10 (1) follows immediately from Eq. (58):
no (D), 0*) = En
(oh V1:h--1 - 20 (2, M1:h 11:h--1. 0*)
EpD
11thh (o1
co
=: (D)
Lemma K.10 (2) follows immediately from Eq. (59):
1 (D)I < ED
(a,, (1, Φ0- (a,, vl.h v1 )
x==1
< ED
4V (ar v1)b-1) + &Be (2, V1:h-1)
$ 44 co (C)
Ey Ei 0² (x, It remains prove Lemma K.10 (3). Note that K 2(D) K K.Ep
Σ Em ('',,, vink) is sum of the martingale difference sequence (st (x', Fin Di Therefore, we can calculate
1 2 2 = K
Em 10 (x. v1:h) 1|
Σ
H
of K
Furthermore, by Freedman's inequality (Lemma H.1), for any fixed vector parameter E (0, B and (0, 1), it holds that
K H
Σ
ΣΣCH
(60 (x', vind,, v) - (E (ol, o) Ix',,,,, vi.--i=)) 2 x- log(1/6) S
Note that for 0r, by Lemma H.5, we have
(eb (x', vinn)s w)² x', winh-a] ynnt wi, (d(x' Vi:--1.9h) yn) -o- (x" winn-11 3h), 0 0n2 < DL (To(- xf vi- T- x vi-nz)) 15Bec (x', vi-h-1)
Therefore, setting 7 30/ we have shown that for any 8 E (0, 1). holds that
30B log(1/0) co(a, v1.h +
P& (z(D). - 0-) 2 2
En
S 6.
K
Σ co(a,, 91 Then, for the random variable V == 30B (=)) 0-) f (D)) the above inequality implies that for any u > 0, IP(V > u) Se-u and hence IP(()+ 2 u) S Therefore, integrating out the above incquality gives E((V) S
Recall that denote (D) : Ep
K
or cquivalently,
Em, (x(D), 8*) 20(() I =: a
30B K
60

---

Published as a conference paper at ICLR 2026
Proof of Lemma K.9. Recall that decompose -0(0; Go (D() x(D), where go (D)
and x(D are defined in Eq. (74) and Eq. (75), respectively. Then, we know
Ag ET
[0; 0*)
(0o 8*) (z(D), 0 8*)
E
X 1000 (D)
co (D) (2(D), 8 8*)
+ | (D)ll
Eg (D)
1 (x(D), 04) -2 2 (D)
IV
Eno
2
+ 2 (D)l
130 (D)lL
Eo (D)
IV
Ems
A (D)l
1 (D)I
X
where the first inequality uses Lemma K.10 (1) and the last inequality uses Lemma K.10 (3) and we recall that a
308
Note that by Lemma K.10 (2),
K
A + (D)]I + (D)
A (D)] + 02 co (D) 8Bea (D)
eg Eg (D)}
(D)
2A + ||=(D)I 4M
V min M, eo (D))
M
where we use min min(M,a) (M,x) we have
max M, x²
Mr, and X M 8BM. Combining these two inequalities,
20
<0 (D)
2A
2 E
A [D(
1 (D)
min M, (D))
E m
2A +
M,
(E,g min (g
Em,
+ 4M min (M, (I())
(D) Min
EE min ci,
21 E min (M, e M / + 4M Vo min(M, (D))
A
2A4, + Mo, +/2 +4VA,
where the two inequalities follow from Cauchy-Sch -Schwarz inequality. Therefore, there are two cases: (a) Ao I k. and the desired upper bound is trivially true. (b) As 2 K. and then it holds that
s 2XA + 555. S 3AA + 8(Ma.)²
2AA0 Mo.
VK K + +4V40
where AM-GM inequality. Hence, it holds that
Mo. 2
A < 8(AAo + a) max Do,8
A
and reorganizing yields
Mo,
2
As S 8 max (AAo +a), / (AAo + a)
< 8(AAg + aa) +8 Mo,
This is the desired result.
61

---

Published as a conference paper at ICLR 2026
K.6 PROOF OF THEOREM 6.1 (TEST-TIME TRAINING)
Recall (from Eq. (12)) that we consider the token-level SGD iterates defined as logTTaa. Tgn. (Vi. x", via-))), for 0, 0.... 1, and = 0 := 6 for t E [T], where (a', vi.:) TT.
gr.++ Proje (0th + log
(76)
To define the guarantee on 0' which we are able derive, we next define the following test-time parameter update 8 (z, V1:1: 0), for a parameter 0 and prompt It is defined recursively for = 1=0,1, 1....... H.
8 (z. (1.: 0) Proje (0 (a, 91:h-1:0) 1/V log (r (yh æ. v1:h-
(77)
We then define a distribution
To X A() as
(- 2r, V1-h 1) TH (1.) ( x, M1:h-1).
(78)
The distribution TIT can be interpreted as an augmented version of the autoregressive linear model Ty that performs test-time training during sampling.
Proof. While the algorithm in Theorem 6.1 might seem somewhat complicated and mysterious, the proof is actually based on fairly simple online-to-batel conversion argument. We use number of basic inequalities already found in the proof of Proposition 5.1 (cf. Appendix K.3),
We first note that expectation gives
specialize Lemma K.5 to the token-level SGD update (12), and taking
H
H
2 s E
E
V log Tgr. (vh x', Vis 1),
vi,h-1), 8thh 100h-00
v log Tg (vh x vi
17
(79)
In the following, we denote
ceA : E[(-VVll log Tan (vh x², Vi 1i, - -8-]
By triangle inequality,
| log (V) x. 2. )]
S V log T (yh,
x. V1:h 11ll + 2 log (8 xx, 11:h-1) V log Th (3h a. 11
Jr, 11:4-1) is concave and B3-smooth, it holds that for any Using the fact that 0 - log No (yh
V log (yh x. (1=A-)
log no- (yh, x, 31=1-1)1l
< B (0 0", V log Tg- (3/1 x, 11 v log गe (yh, x, 11)
Combining the two inequalities above gives that for all E [T), [A].
V log Tge. (yh xx 1)
< 2 V log TT- (v), x² 1)
2B² (V log To- (v/, x² vish-1) log aan. (3h x', bisn- 1), 0in -00)
Note that the conditional distribution of yf. (a*, Vinn OCn)) is given by sh. ~ Tol. xt, vi Hence, taking the expectation over the entire learning process, we have
E V log Tar. (vh xx 11 < Em. V log To- (vr. x. 11:h-1111
2B E
Vog Fgr. (wh x' win-1), 0+h 0*)
E.g Var (I, 2BB 2B2
Plugging the above inequality to Eq. (79) yields
T
H
2
17 E
Σ
T
H
V log Tga. (Vh xt, Vi vi-h-1111 1)1?
CLPb <
+
2
T H
2 +
r/ Em.
Var (x, v1nb-1)
+ nB²
Er,h-
1
h=
62

---

Published as a conference paper at ICLR 2026
Therefore, as long as 77 S 282 it holds that
T H
M Varm (x, Vich 1) = + 2
CLu
S + 2nTEn
By Eq. (58), it also holds that
Er,b = =E( log Tar. (vh x', -i 0rh -00 2 DeL (TO(- | x² Vin.iinn Tai vinn)in-)))
Combining the inequalities above, as long as 1) S 202 we have that
H
T
DKL (To(- x', vin Tgena (: x', vish
S + 2niTo (80)
EtDb
(=1 h=1
Finally, we note that
0thh =-~ -T (x', 1 1
and that for all and h, x², vh-1 H To. Therefore, have the following key identity:
E [KL (To(- x", vininnn) || Tonin xf vinnnin)) 01]
((x,))~m DKL (To (x, 31:h-1) (o. 100 x. Mish v1sh-1))]
B(a,y)= DKL (To(: x 31:.-11) || T I v1x-
Combined with Eq. (80), this implies that
4 To
DKL (To(: ax vist-1) Tana ( x', wisb--l))
E
n
H
DKL (To(: xx vinnnin mara (: x², vish-1)) 0-
H
E
Bm
DKL (To(- x2, B1:h-1) ll TTt (-1 B1Ah
E
DKL (TTD || m.!!
where the last equality uses the chain rule for KL divergence. In particular, we may choose = min
1/2
to derive + E Σi= DKL (TD ll (
28
o.T
02 V
K.7 PROOF OF THEOREM G.2 (GRADIENT NORMALIZATION FOR DISTILLATION)
Specializing Lemma K.5 to the update (35) and taking expectation gives
E
E (a.y))= To (y x)), 0- - 8-)
2 + 77
E
E ( 13 x))1?
(81)
In the following, we analyze (- ((r.)) 10 x)), 0 - gr) and E(2 To 2)1 for any 0 E θ. following the proof of Proposition 5.1 (cf. Appendix K.3).
Relating the gradient to stopped KL divergence. Recall that the estimator ğ is defined in Eq. (35):
H 2oly x) Σ
ao(r, V log To(yn x, 91:h-1),
h=l
63

---

Published as a conference paper at ICLR 2026
and the weight function ag is defined in Eq. (36).
We first recall an elementary property of the quantity ag. By Lemma K.4, we have
H
Ea (z, (ann)-nnn)))
ag Vin- = min
A,
(82)
h=1
and hence
H
H v15-11)
E(
Σ
an (2, y1sh-1)coo v1--1)
= (x, min A,
Σ h=1
cola,
(83)
= Deeq.N (To To),
where we recall that Deeq, (To mp) is defined in Proposition F.10 and we denote Hence,
log N.
( (lx,y)
[y x)], 0*)
a0(a, (000 v1:h-1) 60 (x,, (1,, 0*)
(84)
Dseq. (TD To);
where the inequality uses Eq. (58).
In addition, the following lemma shows that (B.,.))) -na (v |=)1² is well-controlled. Lemma K.11 (Gradient error bound). For any θ it holds that
E(z.))= ( x) < (64A + 2) + 256AB 256AB"Dseq.vv seqq io (To (T)).
Putting everything together.
Finally, combining the inequalities above, know that
E
Deeq (To )
E
(() ~m 3or (y x)). 0: 0+)
" E
E Ea 10a (y x))
+
2
T
< + 1/T (32A + 1)0² + 128AB E
Σ Deeq N (To mor)
where the first inequality uses Eq. (84), the second inequality follows from Eq. (81), and the third inequality uses Lemma K.11. Therefore, as long as S 20324-178 it holds that
T D Deeq. (TD Tot)
E
S + Ao?
1/2
In particular, we may choose min
(o4lor N+2)
To
1
and derive
o log N
B log N T
E
N(o
S
T
By Proposition F.10, this implies
o²
B² T
E
Cov (Toi)
$
+
T
T loog
64

---

Published as a conference paper at ICLR 2026
Proof of Lemma K.11. Fix any 0 E O. By triangle incquality, holds that
[ x) Go- (y x)
as(2 yi.h - 28 (z, V1.)) Φ0 y1a-1))
(2 V1:h 11 (4 (r, v1-h co( 01:h-1) (1 + ( (1.
1=1
H
1/2
4
ag (2, 11:h Var m1sh-1))
Σ ao( (00 ():
+ 8B
Org (I,
h=1
A Var (z, 11: 8B min A, Σ co(a, v2n-1) h-1
where the second inequality follows from Eq. (59), the third inequality follows from Cauchy- Schwarz inequality, and the final lines follow from the property (82) of the weight function ag E [0, 1]. Hence, using (a + 6)2 < 2d² + 262, we have
10o(y x) - -o- (v x)
H
s 32.4
128AB2 min A, Var (z, (4 Σ
E h=1
co(ar,
Therefore, taking expectation of (ar, v) ~ and using E 1100r (y x)) oē and Eq. (83), it holds that
(x.))~ (0 m| (64A + 2)o + 256AB2 Dseq, N (TD TT).
This is the desired upper bound.
L
K.8 NECESSITY OF VARIANCE DEPENDENCE IN HIGH DIMENSION
We generalize Proposition 3.2 to show that in the worst case (where o² HB²), the scaling COv 2 1o v) can be unavoidable for autoregressive linear model. This implies that the dependence on of is generally necessary to achieve upper bounds that do not explicitly scale with H.
Proposition K.1. Let H, N,n, 2 1, and assume log N < emin B2 for a sufficiently small constant e 0. There exists an instance of the autoregressive linear model class 11 with d H. Φ: 6:X x V + B2B)), and θ = B2(1), such that for any proper algorithm Alg with output रे = Tī for 0 θ there exists up E II, such that under Tg, holds that
H
To m "CCoo (TD [)] 2 c min
n log N S
(0,1), 1). and the distribution μ be given by (((+) = pa()) ( where E [0, 1] is a parameter to be chosen later, Let the feature map ó be given by Φ(-, M1:h) = 0, o(+.
Proof of Proposition K.1. We consider X' (+:-).
Bуje en. where (e,, ,eH) fixed orthonormal basis of H. Note that with this construction, we have To(2) = V1-h Ber(1 (1/2), and
BBn
Tyh +. 11) Ber
B0,
-T..b-
65

---

Published as a conference paper at ICLR 2026
Note that for any h E [H] can bound
Co 88 $ DH(TTOn Tor C, - -8/1
as long as 8 El
B Bl.
1
We fix
E 10,
max (V. to be determined later, and for any E {-1,1) we let 0, =
H k=
U/e and
θ = (0 E (-1,1) c B2(1),
Ho = (mo 0 Oa)
Then direct argument (see e.g., (Wainwright, 2019, Section 15.3)) shows that when pr S B for a sufficiently small constant cp, there exists E Oo such that under To ng. holds that
H Σ
P
0hl 20 2 eH.
Therefore, with probability at least 5, it holds that Σ 1610, -0%1 2 c 2 2 implies
eH
and this in turn
Σ
D no Tamn) C
Then, by Proposition F.11, we know that under the above event, long as log N < c.H we
have COVN 2 S. Choosing e bound.
V
4log N and p min 1, - gives the desired lower ea HB
L PROOFS FROM SECTION 6
L.1 PROOF THEOREM 6.2 (SIMPLE TOURNAMENT)
Below state and prove generalization of Theorem 6.2 which holds when the data distribution To is not necessarily in the model class II.
Theorem 6.2' (General version of Theorem 6.2). Fix N 2 L, and consider the estimator म from Eq. (14);
म arg min max Cov (n' []
(85)
sell
x'Ell
For any 8 (0, 1), parameter a, c 2 0, with probability at least 1 - 6, it holds that
2ot (7) S min Cov (T) + 1
log Noc (II; clog N) + log 8-1
(86)
#ell
N1-a-26
Proof of Theorem 6.2'. Fix N. N' 1,a > 0, and let म E arg min.en Cov (TD (). We study the estimator
arg min max CovN ( m) HEll
(87)
x'Ell
Recall that we denote Cov (α' x) P.o
(x(l=)
=()|²) 2 N) (cf. Lemma J.2). By Lemma J.2, with
probability at least 1 it holds that
6 2
Cov ( ④)
Cov
( π)
Estati
VT E 1I,
(88)
where Estat
log (8)))/8)
Next, again by Lemma J.2, with probability at least
2: it holds
that
Cov (T F) < 2Cov To 20 n (TT F) + Estat:
VI II.
(89)
66

---

Published as a conference paper at ICLR 2026
In the following, condition on the success event of Eq. (88) and Eq. (89). Then, can bound
Cov 2 20 ( () Ettt I COVN []
I max Cov (I (1
x' ''Ell
min max Cov N (TT' ||)) I max COVN aell e'ell
a n'ell
2 max Cov Th r'ell
6-20 x ( π) + Estat
Reorganizing yields
Cov Th 2 ( TT) S 4 max Cov Th 2 || T)) + 4Estat-
(90)
tell
Note that for any N" and models nr, n.', x",
Cov N' (T (|)) S Cov (1'l11" Cov xx).
(91)
Hence, for any model π E II,
Cov a N N (TD T) Cov Ny (To π + Cov zan n i TT,
(92)
Cov N (Tr Cov N (π m
Cov N/ (To nt.
(93)
Therefore, combining the inequalities above,
that
Cov.20 N (R)
Cov
22 N' N (TD || .)
Cov C.VD (TD π Cov N []
Cov (TD (T) + max Cov 2- N (TT (7) + 46stat #Ell
50ov%; ( (To
4 max Cov To Hell
e 20 N/N To) 4estat
2a N'
5Cov Ta (To T)
4E
N
where the first inequality uses Eq. (92), the second inequality uses Eq. (90), the third inequality uses Eq. (93), and the last inequality follows from the fact that Cov (m || To) Pmo
<(ylx)
A <
mlly
The claimed bound (86) follows by setting a clog N, and N' = No
L.2 PROOF THEOREM G.3 (OFFSET TOURNAMENT)
Divergence. For distributions P,Q E A()), we define the following divergence for N 2 1:12
(x( = max Eynp dQ
dP
NN)
E
E [0,
dP
dQ
Then, for models m' X A()), we further define
EN (TT n²)= Enpo (n((T(: x)) [| xx)).
Under this divergence, it holds that for any event E.
P Pu. (E) < N. IP Pa (E) Ev (T n'),
(94)
P P (E) S N. P (E) + EN (T || x'),
(95)
where P Pu, is the probability under x ~ and y ~ M(. It.. Furthermore, we can bound
T x)
COVaN (TT) P Ho
2N < ENN, (Tp IT)
(96)
イ④ x))
17This divergence is inspired by Huang al. (2025b), but our definition differs slightly from the standard Exy-divergence (Polyanskiy, 2010; Black & Polyanskiy, 2023).
67

---

Published as a conference paper at ICLR 2026
Theorem G.3' (General version of Theorem G.3). Fix N, 2 such that N 2 8 Consider the
estimator
arg min max {Cov ( [] 2y Cov (T' [] nell x'ell
(97)
Then with probability holds that
CoveN- S min (TD [|>]||]>>>>>))) /
=ell
n
Note that &, (TT || T) when [ x) Ty 2) I logy for any x E N, E V.
Therefore, Theorem G.3 is an immediate corollary by setting No.
Proof of Theorem G.3', For n' II, we define the set
Ty x) T'y x)
(x( = (1)
2 N
Suppose an i.i.d. dataset D [(x*, (o),,')) )iE To is drawn. We write P. = E 6 o w/) and Nn Σi= 6z. denote the empirical measures (i.e., P, is the uniform distribution over D), and let PRn.= be the probability under the distribution x ~ pn. y ~ x(- x). Under this notation, we have COvx (T' x) (C (IT', T)) and we also recall that
( xx) T x²) 2 NN = Pm (C''
Cov (T' (T) = (((+))))))))))
Thus, the tournament estimator Eq. (97) can expressed as
मे arg min max [(T,, n'), mell
(98)
x'ell
where
C(a,T) = (C((((7', - 2n P (C'('' (T'(())))))))))))))) ,))).
(99)
As an immediate consequence Lemma H.2 and the union bound, we have the following lemma.
Lemma L.1. Fix E (0, 1), and define Estat (66((T)))) bounds hold simultaneously:
With probability 6, the following
(1) For all n, T' E 11, it holds that
mm (Cn. ( T)) Estat 2 PA(CN(TT', R)) 2 P n,TO (CN(TT' (,''TTT)))))
T))
Estats
(100)
2P po (Cx (x', m)) + Estat 2 (C T)) 2 P (Cx (Cx( aann,'n
xx)
Estat
(101)
(2) For any II, it holds that Engpa (To π) s 2 ,a (To πτ) + Estat-
In the following, we fix 8 E (0,1) and condition on the success event of Lemma L.1. Let e arg min (T || x). We denote Eapx (TD ()) and E'px Eygu, (To (). Note that by Lemma L. we have e'spx < 2aap + Estat-
Then, for any ne ET,
L(T,T) < 2F pm. (CN ((TTT))))))))))))) - 27P un. (CN(TT' #)) + Estat
2P
$ 2E, (To T)) + Estat = 2spx + Estat-
where the first inequality uses Eq. (101), and the second inequality uses Eq. (94). Therefore, we have
max x'ell
L(R, n') = min max L(m, x') ) S max L(T, x²) S Estat + dapx
wEll r'ell
w'ell
68

---

Published as a conference paper at ICLR 2026
In partieular, know L(f, R) Estat + d'axx. Then, we can bound
Pn(CN (7, R)) L( π)
2y|p nn. (f)
27
$
PP
Pan. (CN (T,, R))
N
27
[yI pu (CN (T, R)) + d'apel
2y
2
(C T)) + Estat + <'sp
N
where the first inequality follows from the fact that ( the second inequality uses Eq. (95): IP ++,.=(E) yIp po
x Nf x)) for (I, E (T,
E
< En (TD πt) Sap for any event
E. and the third inequality uses Eq. (101). Therefore, using 2 8y2, we know ( S 55stat 2px Then, using Eq. (100), have
(T, Cov [T Fu (CN (T,, (T) S (C, R)) + 2t I 12 45
By Eq. (91), it holds that
CoV2N (7) = Cov To (TD | ()) S Cov (To || T) + Cov / t f),
and we also have Cov22 (To π) Cov2, (T) S En (To ( Eapx by Eq. (96). Combining the inequalities above, we can conclude that
CoveN-( S Cov ( (t) + Eapx I 12 + 45 Eapx
Finally, using Lemma H.2, we have Eppx S 2p + Estat- This is the desired upper bound.
69


Reviews:
[{"review_id": "VvwWoyor9Q", "reviewer": "Reviewer_Au4Y", "summary": "This paper proposes the coverage principle: next-token pre-training implicitly drives models toward good coverage over high-quality responses, which better predicts post-training success than cross-entropy alone. The authors provide theoretical justification using autoregressive linear models and some assumptions such as assumption on convexity.\nBased on the findings, the paper further presents a test-time training decoding scheme that updates logits while sampling; this improves coverage and can outperform pure MLE-based guarantees.In addition, it also provides a method to select checkpoints with better pass@k than cross-entropy-based selection.\nThis paper poses new insights on understanding the relationship between pretraining and downstream performance and provide effective ways to utilize these insights in practice.", "questions": "1. Would it be possible to demonstrate the checkpoint selection or decoding methods using an LLM and complicated tasks, such as HLE, AIME, and factual QAs like SimpleQA, etc.?\n2. Your results often rely on the data distribution \\pi_D covering the downstream task. Can you formalize how much misspecification in \\pi_D is tolerable while keeping useful coverage guarantees?\n3. For the gradient normalization trick, what’s a plug-and-play recipe for transformers and how to implement it?\n4. How robust is the proposed TTT sampler to distribution shift, long contexts, and temperature samplig?\n5. Results focus on graph reasoning with GPT-2–style models. Can you provide additional evidence on natural-language tasks (e.g., code-gen Pass@N or reasoning benchmarks) to test the coverage–performance link?\n6. Can you provide an intuitive example where cross-entropy improves while coverage worsens, and a rule-of-thumb for detecting this in training curves?\n7. Any other suggestions for LLM pretraining/mid-training before postraining?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper provides a solid theoretical justification on the proposed coverage hypothesis.\n2. The insights can be translated to useful applications, including better test-time sampling methods and better checkpoint selection methods.\n3. This paper can inspire future works to develop algorithms that can be more effectively used for improving coverage using methods in addition to traditional next-token-prediction loss.", "weaknesses": "1. As a general limitation of the theoretical justification of deep learning, the proofs are shown with linear autoregressive models and assumptions on convexity, therefore the takeaways may have a risk on directly being translated to complex deep neural networks.\n2. The hypothesis does not directly indicate a better way than next-token-prediction for pretraining to improve the coverage during training. The proposed distillation-only training is unrealistic in pretriaining.\n3. The experiments are conducted on a relatively toy cases, e.g., a graph reasoning task, while real-world tasks are more complicated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sadh2SEKXX", "reviewer": "Reviewer_RfJt", "summary": "Motivated by the observation that downstream results are not always correlated with cross entropy results, this paper introduces coverage that is better related to downstream performance. The authors show through theoretical results that next-token prediction leads to models with good coverage. From a practitioner's perspective, the authors then demonstrate a decoding intervention to improve coverage, and a way to use coverage for selecting models.", "questions": "Coverage seems to focus on diversity more than quality - could over indexing on this during pre-training make post-training more difficult?\n\nWould you recommend using this instead of perplexity during pre-training? If so, how would you expect it to behave in a scaling law study?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- well motivated and tied to prior work, figure 1 makes this clear\n- theory is sound and directly relates to BoN and post-trained models\n- also contains practical recommendations (section 6)", "weaknesses": "- lack of experimental evidence. would be nice to expand on the downstream tasks and post-training methods.\n- coverage assumes a target distribution, which may limit its applicability for pre-training\n- coverage relies on being able to verify samples, which may be difficult for certain (more open-ended) tasks", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "w7kbEJYjJw", "reviewer": "Reviewer_w8sn", "summary": "In this paper authors introduce a new quantity called coverage profile which is better aligned with a model's performance on downstream tasks that require Best-of-N sampling and/or RL. They show that autoregressive next-token prediction implicitly optimizes coverage. Coverage has better generalization and convergence than CE. While CE can continuously improve during pre-training, coverage can drop at some point which may reflect the decrease in the BoN downstream performance.\nThe authors perform experiments with graph reasoning tasks to support their theoretical results. They present a way to estimate coverage profile in practice by defining a Tournament procedure to select checkpoints with the best in-class coverage profile. Moreover, they show that test-time training (TT token-level SGD) improves coverage.", "questions": "- Can the scaling law (eq. 3 and therefore Proposition3.1) be validated empirically (e.g. using the tournament procedure)? \n-  Figure 1, why coverage has non-integer index (e.g. $Cov_{5.7}$), if $N$ is the number of sampling attempts (eq. 1) or does $N$ mean something else here?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The theoretical justification of the coverage principle that can explain insufficiency of cross-entropy for accurately predicting downstream BoN performance is an important contribution.\n- The paper is well-written and solid. Theoretical results are sound.\n- Algorithmic improvements like gradient normalization in SGD can lead to better coverage and therefore better BoN performance.\n- The work opens a new interesting direction in connection of the base model and post-training performance by utilizing the notion of coverage profile.", "weaknesses": "- The only concern that I have is that the authors validated their theory on only one task - graph reasoning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)