cs.AIOct 8, 2026
SaveEstimating great expectations under autoregressive language models with potentials
Organizations: ETH Zürich · Columbia University
Abstract
Many applications of language models hinge not on individual samples but on the expectation of a test functional under the model. Estimating such expectations reliably can be computationally expensive. In this paper, we show how to make estimation more efficient by exploiting the next-token conditional probabilities which are available as a by-product of sampling. We do so through potentials: real-valued functions on prefixes that decompose the test functional additively. We construct an estimator whose variance depends on the chosen potential, and derive conditions under which a potential reduces this variance. We then develop practical potentials for several estimands and applications, and demonstrate substantial variance reductions across several estimands at comparable computational cost.
Figures & tables
Figure 1: Variance reduction in the proper-loss experiment (lower is better). (a) Variance ratio of the potential estimator to Monte Carlo across power scores ( gives ); entropy and Brier are shown on the right. Each series is one policy–reference pair (see 10 ): solid lines are the four main policy–reference pairs and dashed lines are consecutive Gemma-2 SPPO pairs. (b) For , variance ratio against the estimated , with one colored marker per pair and grey markers for individual prompts from all pairs. The shaded regions indicate prompts below and above the Monte Carlo variance. (c) Variance ratio of the cross-fitted control-variate estimator to the potential estimator; hollow markers use the full-sample oracle coefficient.
Figure 2: Cost-adjusted variance ratio for -gram frequency estimation. (a) Empirical cumulative distribution over targets with bootstrap RSE at most , with approximate log-ratio confidence bounds. (b) Variance reduction versus estimated frequency for the targets with at least expected occurrences in the independent frequency bank. (c) Fraction of targets whose frequency estimate reaches each RSE threshold, grouped by token length, at the same sample size for both estimators. Dashed parity lines in (a–b) mark equal variance at equal cost under the proxy; in (c), line styles distinguish RSE thresholds.
| Estimator | Var. ratio vs. MBAR [95% CI] |
|---|---|
| Fitted selection | 0.88 [0.86, 0.91] |
| Terminal | 0.92 [0.90, 0.94] |
| Terminal, | 0.92 [0.90, 0.95] |
| Logistic | 0.90 [0.86, 0.91] |
| Mixture model | 0.87 [0.84, 0.90] |
Table 3: variance relative to MBAR, median over the populated ARI bins; intervals from paired bootstrap replicas. Full version in 19 .
Appendix figures & tables20 assets
Supplementary material from the paper’s appendix.
Appendix
| median | worst prompt | |||
|---|---|---|---|---|
Table 6: Per-prompt variance of against Monte Carlo, prompts.
Figure 3: Panel (a) shows how sampled trajectories update the green-token count over 64 generation steps. Red trajectories are those whose final green-token count exceeds the detection threshold, while black trajectories remain below it. The background encodes the potential’s assigned probability of eventual detection at each step and green-count state, with white indicating probability zero. Red trajectories tend to pass through higher-probability regions, whereas black trajectories concentrate in low- or zero-probability regions. Panel (b) shows the variance reduction at different detection probabilities
| bank | prompts | sampling policy | reference | horizon | |
|---|---|---|---|---|---|
| IMDb-256 | IMDb prefixes | gpt2-imdb + DPO | gpt2-imdb | 256 | 4000 |
| IMDb-50 | IMDb prefixes | same policy | gpt2-imdb | 50 | 4000 |
| Tülu-3 | AlpacaEval 2 | Tulu-3-8B-DPO | Tulu-3-8B-SFT | 256 | 100 |
| Gemma It3 vs It0 | AlpacaEval 2 | SPPO-Iter3 | gemma-2-9b-it | 256 | 100 |
| Gemma It1 vs It0 | AlpacaEval 2 | SPPO-Iter1 | gemma-2-9b-it | 256 | 100 |
| Gemma It2 vs It0 | AlpacaEval 2 | SPPO-Iter2 | gemma-2-9b-it | 256 | 100 |
Table 10: Policy–reference banks of the proper-loss experiment. Labels are those used in the figures; is the sampling policy and the reference; model names are resolved in 11 . All banks have prompts and responses per prompt.
| model | repository |
|---|---|
| gpt2-imdb | lvwerra/gpt2-imdb |
| gpt2-imdb + DPO | trained by us from lvwerra/gpt2-imdb (recipe above) |
| distilbert-imdb (DPO preference judge) | lvwerra/distilbert-imdb |
| Tulu-3-8B-SFT | allenai/Llama-3.1-Tulu-3-8B-SFT |
| Tulu-3-8B-DPO | allenai/Llama-3.1-Tulu-3-8B-DPO |
| gemma-2-9b-it (SPPO Iter0) | google/gemma-2-9b-it |
Table 11: Model repositories (Hugging Face) for the names used in 10 .
Figure 4: Variance reduction as a function of quantity magnitude for all losses and divergences in the proper-loss experiment. Each panel shows the variance ratio against the estimated magnitude of the quantity, with one colored marker per policy–reference bank and grey points for individual prompts. For KL and the power-score regrets, variance reduction becomes stronger as the policy–reference divergence decreases, with this dependence becoming more pronounced for larger . Entropy and Brier loss instead remain comparatively close to the Monte Carlo baseline across banks.
| vs. | vs. | prompts with | ||||||
|---|---|---|---|---|---|---|---|---|
| quantity | median [IQR] | % | median [IQR] | % | med. | med. [95% CI] | % | |
| 0.204 [0.061, 0.503] | 98.5 | 0.204 [0.061, 0.482] | 99.5 | 63 | 1.134 | 0.974 [0.945, 0.990] | 65.1 | |
| regret | 0.202 [0.063, 0.441] | 99.0 | 0.193 [0.063, 0.436] | 99.8 | 40 | 1.083 | 0.951 [0.932, 0.973] | 77.5 |
| regret | 0.164 [0.063, 0.400] | 99.6 | 0.160 [0.063, 0.396] | 99.9 | 15 | 1.058 | 0.935 [0.891, 1.044] | 66.7 |
| regret | 0.142 [0.062, 0.365] | 99.9 | 0.141 [0.061, 0.363] | 99.9 | 4 | 1.190 | 1.080 [0.925, 1.102] | 25.0 |
| regret | 0.131 [0.060, 0.330] | 99.9 | 0.129 [0.060, 0.329] | 99.9 | 3 | 1.125 | 1.040 [1.009, 1.043] | 0.0 |
Table 12: Per-prompt variance ratios for the losses and divergences experiments, pooling all eight model pairs. We compare the potential estimator and its cross-fitted control-variate variant against Monte Carlo. The rightmost columns restrict attention to prompts for which has higher variance than Monte Carlo and show how often the control variate reduces the ratio below one.
| std. dev. | ||||
| estimator | mean | |||
| 6.69 | 3.32 | 1.49 | 1.05 | |
| HT | 6.69 | 3.32 | 1.49 | 1.05 |
| Schulman CV, | 6.47 | 35.67 | 15.95 | 11.28 |
| Schulman CV, fitted | 6.36 | 11.44 | 5.11 | 3.61 |
| 6.69 | 2.05 | 0.92 | 0.65 | |
Table 13: Reproduction of the KL-estimation experiment of Amini et al. (2025) : mean estimate and empirical standard deviation (nats) at sample sizes , over prompts with samples each. The mean does not depend on . Lowest standard deviation per column in bold.
Figure 5: Gradient estimation along a DPO trajectory of the GPT-2 sentiment policy (checkpoints every steps; the variance of an estimator is over the responses). (a) Variance of relative to the centered Monte Carlo gradient , one series per quantity (regrets darker for larger ; undefined at step ). (b) Variance of the cross-fitted control-variate estimator relative to . (c) The cross-fitted coefficient, mean over prompt folds with the range as band; recovers . Bands in (a) and (b) are prompt-cluster bootstrap intervals; one training run and one seed.
| model | – | – | – | |
|---|---|---|---|---|
| GPT-2 ( ) |
Table 16: Variance of the observed-prefix estimator relative to Monte Carlo at equal sample size ( ): median over the unigram targets of GPT-2, banded by their occurrence count in the samples (columns).
Figure 6: Figure 2 for expected counts. Cost-adjusted variance ratio for -gram count estimation under the per-target cost proxy . (a) Empirical cumulative distribution over targets with bootstrap RSE at most , with approximate log-ratio confidence bounds. (b) Ratio versus estimated expected count per document for the targets with at least expected occurrences in the independent frequency bank. (c) Fraction of targets whose count estimate reaches each RSE threshold, grouped by token length, at the same sample size for both estimators. Dashed parity lines in (a–b) mark equal variance at equal cost under the proxy; in (c), line styles distinguish RSE thresholds.
Figure 7: -gram target inventory and token-length distribution. The -gram task provides a controlled way to study targets of different lengths and, consequently, different marginalization costs, which scale with the number of GPT-2 tokens in the target. (a) Number of targets at each order , together with the minimum, maximum, mean, and median of their canonical realizations. The inventory contains targets in total. (b) Distribution of token lengths within each order. Increasing shifts the inventory toward longer targets, while GPT-2 tokenization also induces within-order variation in . (c) The same distributions on a log scale, showing that each order additionally contains a progressively thinning tail of targets requiring more tokens, and therefore more costly marginalization.
Figure 8: Cost-adjusted variance ratio against target rarity before the occurrence filter. Both panels show the frequency-estimator variance ratio, against estimated -gram frequency (left) and expected count per document (right), before applying the occurrence filter to the independently estimated -axis quantity. In the main figure, we retain only targets whose independent frequency estimate has RSE at most , corresponding to at least expected occurrences in the estimation sample, to restrict the analysis to more reliable estimates. Here, we show all targets with positive estimated frequency and a finite positive variance ratio; this remains a subset of the -target inventory.
| (a) all targets | (b) -gram order | (c) token length | ||||||||
| frequency | # targets | |||||||||
| 354 | 1.5 | 2.0 | 1.5 | 1.5 | 1.3 | – | 1.6 | 1.5 | 1.3 | |
| 701 | 5.8 | 6.5 | 6.6 | 5.3 | 5.1 | – | 6.6 | 5.4 | 5.1 | |
| 725 | 28 | 35 | 28 | 25 | 17 | 42 | 26 | 25 | 17 | |
| 684 | 130 | 160 | 110 | 85 | – | 160 | 110 | 85 | – | |
| 299 | 310 | 330 | 260 | – | – | 330 | 260 | – | – | |
Table 17: Variance reduction is strongest for rare targets, even within fixed -gram order and token length. The targets of Figure 2 b are grouped by their independently estimated frequency . Each entry reports the median cost-adjusted variance ratio within that frequency bin; lower values indicate larger gains over Monte Carlo. (a) The trend across all targets. (b–c) The same comparison within fixed -gram order and canonical token length , showing that the rarity trend persists after controlling for target length and marginalization cost. The bottom rows give the number of targets in each column and Spearman’s correlation between and . Cells with fewer than targets are omitted.
| Expected count | Frequency | |||
|---|---|---|---|---|
| Max. RSE | # targets | 95% CI below | # targets | 95% CI below |
| 5% | 84 | 61.9% | 41 | 70.7% |
| 10% | 317 | 84.2% | 240 | 89.6% |
| 20% | 1,048 | 92.2% | 953 | 95.2% |
| 50% | 2,661 | 93.0% | 2,583 | 94.7% |
Table 18: A majority of targets in each precision subset retain a variance advantage under the cost proxy. Among targets with positive estimated frequency and a finite positive variance ratio, each row retains those with bootstrap RSE below the indicated cutoff . We report the number of retained targets and the percentage satisfying , a sufficient condition for the approximate normal interval to lie below parity. These are conservative shares for the normal interval, rather than the log-ratio intervals in Figures 2 and 6 .
Algorithm 1 Potential selection by nested cross-fitting
Figure 9: Nested cross-fitting, our in-sample selection among candidate potentials ( Algorithm 1 ). Each of outer folds is held out in turn (red); on the remaining folds the candidates are compared by inner cross-validation, the winner is refitted on all remaining data and applied to the held-out fold.
Figure 10: Variance ratio of the selected potential on the event-analysis experiment: chain-level variance of the cross-fitted estimator relative to MBAR, per ARI bin, over the bins – : the bins on which MBAR’s own variance is itself measured to within , plus bin (at ), kept so that the range is contiguous; shading gives the bin probability, the dashed line the median ( ) and the band its interval. All populated bins are in Figure 12 .
Figure 11: The same view as Figure 10 for each single candidate, each refitted without the held-out fold and applied to every bin: chain-level variance relative to MBAR per ARI bin over the same bins ( by the rule plus bin ), with the bin probability as shading and the candidate’s own median over these bins with its interval. The terminal potential and its calibrated version remove , the logistic model , and the mixture model , against for the selection, which picks the mixture model in four folds and the logistic model in one.
Figure 12: Variance reduction across ARI bins. (a) variance of the cross-fitted estimator relative to MBAR for all populated bins. Filled markers denote the bins for which the MBAR variance is estimated with RSE below ; over these bins, the median variance ratio is with a bootstrap interval of , and over all populated bins it is with interval . Ratios in the remaining tail bins are much noisier because the MBAR variance itself is poorly estimated. Whiskers show bootstrap standard deviations from paired chain-bootstrap replicates. (b) RSE of the MBAR variance and mean estimates across bins, showing that uncertainty in the variance estimate grows sharply in the tails. (c) Variance ratios against the RSE of the MBAR variance estimate. Among the reliable bins, the ratio shows no systematic dependence on baseline uncertainty (Spearman , ), indicating that the observed reduction is not driven by noise in the MBAR variance estimate.
| Estimator | Var. ratio vs. MBAR [95% CI] | Sample-size eff. [95% CI] | MBAR-equiv. samples ( ) [95% CI] |
|---|---|---|---|
| Nested selection | 0.88 [0.86, 0.91] | 1.14 [1.10, 1.17] | 7.58 [7.34, 7.80] |
| Terminal | 0.92 [0.90, 0.94] | 1.09 [1.06, 1.11] | 7.29 [7.09, 7.40] |
| Terminal, | 0.92 [0.90, 0.95] | 1.09 [1.06, 1.11] | 7.27 [7.05, 7.40] |
| Logistic | 0.90 [0.86, 0.91] | 1.12 [1.09, 1.16] | 7.46 [7.31, 7.75] |
| Mixture model | 0.87 [0.84, 0.90] | 1.15 [1.11, 1.19] | 7.69 [7.43, 7.92] |
Table 19: Event analysis: chain-level variance of each estimator relative to MBAR, median over all populated ARI bins, with intervals from paired chain-bootstrap replicas (fitted potentials, selections, and coefficients held fixed). Sample-size efficiency: median over bins of ; MBAR-equivalent samples: that factor times the actual sample size states.