Coverage Generalizes Faster Than Cross-Entropy, Unlocking Best-of-N Post-Training

The Coverage Principle: How Pre-Training Enables Post-Training

Fan Chen, Audrey Huang, Noah Golowich, Sadhika Malladi, Adam Block, Jordan T. Ash, Akshay Krishnamurthy, Dylan J. Foster

stat.ML, cs.AI, cs.CL, cs.LG, math.ST

2025-10-17

Next-token prediction implicitly optimizes coverage of high-quality answers. Coverage generalizes faster than CE, without a sequence-length tax, and decides Best-of-N success.

What problem this solves

Driving pre-training cross-entropy down does not guarantee that post-training will work. Several papers (Liu, Zeng, Chen, Lourie, Springer) have already shown that a better next-token predictor need not be a better starting point for downstream tasks. Best-of-N and RL with verifiable rewards are the two main post-training routes today. Both stall on the same requirement: the pretrained model must put enough probability mass on high-quality responses.

The coverage profile measures that requirement. For a data distribution πD and a model π̂, CovN is the probability that the model underweights a data response by a factor of at least N. It is the CDF of the log density ratio; KL is the mean of the same random variable. A clean mean can hide a bad tail. In the graph-reasoning experiment of Figure 1, KL falls monotonically through training while coverage can get worse. For large N, coverage predicts Pass@N better than KL.

Coverage is also transitive. If πD produces a correct response with 10% probability and Cov{N/10}(πD ∥ π̂)=ε, downstream error is about 10ε. The analysis therefore studies coverage against the pre-training corpus itself.

Method

The coverage principle says that next-token prediction (more generally, maximum likelihood) implicitly pushes the model toward good coverage, and that coverage generalizes faster than cross-entropy, without paying for problem-dependent parameters such as sequence length H.

A weak scaling law is available: CovN ≤ KL / log(N/e). That bound says test-time compute should grow exponentially in the target accuracy. Sequence-level KL, however, grows linearly in H. For autoregressive linear models, any proper estimator (one forced to output a member of the given class) satisfies DKL ≥ H/(4n) with probability at least 0.25. Figure 2, with H in {8, 16, 24}, shows KL at convergence tracking H linearly.

A Bernoulli toy case shows why a huge KL need not destroy coverage. Data is Ber(p) with small p. With constant probability the training set is all zeros, MLE returns Ber(0), and expected KL is infinite. Coverage still satisfies CovN ≲ log(1/δ)/n. The missing mass blows up KL; coverage only writes that mass off once.

MLE on a general class avoids the length tax. Theorem 4.1 splits coverage into two terms:

The proof is close to Mendelson's small-ball method: a model with bad coverage contributes a log N likelihood penalty on many training points, so it cannot be the MLE. For overparameterized autoregressive linear models the bound refines to E[CovN] ≲ √(σ★²/(n log N)) + B²/n. σ★² is inherent variance, an effective sequence length that counts only pivotal tokens whose features still move a lot given the prefix. Most tokens are nearly deterministic, so σ★² is typically much smaller than B²H.

One-pass SGD is less lucky. The coverage upper bound carries a √(B²H/T)/log N term, and a lower bound shows CovN can be as large as min(H/(T log N), 1). The failure is prompt heterogeneity: some sequences have gradients that grow with H, forcing a learning rate of order 1/(HB), while others then learn too slowly. Gradient normalization (rescaling by the batch gradient norm, close in form to SignSGD/Adam) removes the H dependence, leaving √(σ★²/(T log N)) + B²/T + B/(K log N). In a distillation setting, truncating token-level gradients by running token-KL recovers the MLE rate.

A test-time-training decoder also works: after sampling each token, take one gradient step on it, then reset at the end of the sequence. Combined with token-level SGD, the coverage bound matches Theorem 4.2, with an extra 1/√log N on the leading term. Because the decoder can leave the model class, it also bypasses the KL lower bound for proper methods. For checkpoint selection, two tournament procedures pick the candidate that minimizes empirical coverage against the hardest rival in the class, even when πD is outside the class. In Figure 1, the diamond (tournament) checkpoint beats the red (min-KL) checkpoint on Pass@N.

The experiments are synthetic graph reasoning, not real LLMs. A GPT-2-style causal Transformer (4 heads, 6 layers, 384 dim) is trained with Adam. Pre-training uses 8×64,000 prompts for 200k steps; fine-tuning then runs 50k steps on a skewed mixture (90%/10% of two graph classes). Figure 2 reports the median of 16 seeds.

Results

SettingMetricResult
Any modelCovN vs KLCovN ≤ KL / log(N/e)
Autoregressive linear, proper estimatorSequence-level KL lower bound≥ H/(4n) with probability ≥ 0.25
Maximum likelihoodCovNFine term / log N + coarse term / n, no explicit H
Overparameterized linearE[CovN]≲ √(σ★²/(n log N)) + B²/n
One-pass SGDCovN lower bound≳ min(H/(T log N), 1)
Gradient-normalized SGDAverage CovN≲ √(σ★²/(T log N)) + B²/T + B/(K log N)
Graph-reasoning experimentKL vs coverage vs Pass@NKL falls, coverage can degrade; coverage predicts large-N Pass@N better; KL scales with H, coverage does not

Why it matters

The paper gives an operational answer to what pre-training actually banks for post-training: not how low cross-entropy goes, but whether rare correct responses are covered. Best-of-N regret is the same order as CovN, so BoN succeeds exactly when coverage is good. RL methods such as GRPO are widely believed to need some form of coverage as well; otherwise there is no trajectory to reinforce.

Three things practitioners can try: do not pick checkpoints by validation cross-entropy alone (the tournament selected better Pass@N in the synthetic experiment); gradient normalization, which Adam already roughly does, has a coverage benefit in theory; at decode time, a single test-time-training step per generated token is provably useful. All three are validated only on a graph-reasoning toy. Treat them as incremental interventions, not a production objective swap.

Limitations

The setup is closer to SFT than to pre-training: data is split into prompt and response, and the prompt distribution is the same at both stages. Real pre-training is not cut that way. The main analysis assumes realizability (πD sits in the class); misspecification is handled better only for convex classes. Coverage is necessary and sufficient for BoN; for RL the paper cites empirical evidence from Yue, Gandhi, and Wu, without a minimal condition.

Every experiment is layered-DAG path finding with a 6-layer, 384-dim model. There is no real language model and no real RL post-training comparison. The link from Adam to coverage is stated as speculative. Coverage looks only at predicted probabilities, not at semantic coverage in the representation, which the authors note may be conservative. The tournament evaluates pairwise empirical coverage, which gets expensive as the candidate set grows.

Terms

Source

What people are saying

Related papers

All paper explainers