Estimating great expectations under autoregressive language models with potentials

Francesco I. Re; Shubhangi Ghosh; Tim Vieira; Ryan Cotterell

cs.AIOct 8, 2026

Estimating great expectations under autoregressive language models with potentials

Authors: Francesco I. Re, Shubhangi Ghosh, Tim Vieira, Ryan Cotterell

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

Panel of Figure 1 — 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 ( $\beta\to 1$ gives $D_{\mathrm{KL}}$ ); 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 $D_{\mathrm{KL}}$ , variance ratio against the estimated $D_{\mathrm{KL}}(p\|q)$ , 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.

Panel of Figure 2 — Figure 2: Cost-adjusted variance ratio $\widehat{\rho}_{c}(w)$ for $n$ -gram frequency estimation. (a) Empirical cumulative distribution over $3{,}700$ targets with bootstrap RSE at most $2$ , with approximate $95\%$ log-ratio confidence bounds. (b) Variance reduction versus estimated frequency for the $2{,}831$ targets with at least $25$ 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 $M=4{,}000$ 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, $\alpha$	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

64

populated ARI bins;

95\%

intervals from

1{,}000

paired bootstrap replicas. Full version in 19 .

Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

$\delta$	$m$	median	$[p_{10},\,p_{90}]$	worst prompt
$1$	$32$	$0.471$	$[0.413,\,0.564]$	$0.907$
	$64$	$0.513$	$[0.460,\,0.597]$	$0.894$
	$128$	$0.561$	$[0.508,\,0.677]$	$0.931$
	$200$	$0.612$	$[0.544,\,0.748]$	$0.965$
$2$	$32$	$0.538$	$[0.471,\,0.631]$	$0.836$
	$64$	$0.584$	$[0.520,\,0.671]$	$0.876$

Appendix

Table 6: Per-prompt variance of

\widehat{\theta}({\vec{\phi}})

against Monte Carlo,

512

prompts.

Panel of Figure 3 — 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 $p$	reference $q$	horizon	$B$
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

Appendix

Table 10: Policy–reference banks of the proper-loss experiment. Labels are those used in the figures;

p

is the sampling policy and

q

the reference; model names are resolved in 11 . All banks have

512

prompts and

B

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

Appendix

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 $\mathrm{Var}_{p}(\widehat{\theta}({\vec{\phi}}))/\mathrm{Var}_{p}(\widehat{\theta}_{\mathrm{MC}})$ 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 $\beta$ . Entropy and Brier loss instead remain comparatively close to the Monte Carlo baseline across banks.

quantity	median [IQR]	% $<1$	median [IQR]	% $<1$	$n$	med. $\widehat{\theta}({\vec{\phi}})$	med. $\widehat{\theta}({\vec{\phi}},\alpha^{\star})$ [95% CI]	% $<1$
	$\widehat{\theta}({\vec{\phi}})$ vs. $\widehat{\theta}_{\mathrm{MC}}$		$\widehat{\theta}({\vec{\phi}},\alpha^{\star})$ vs. $\widehat{\theta}_{\mathrm{MC}}$		prompts with $\mathrm{Var}_{p}(\widehat{\theta}({\vec{\phi}}))>\mathrm{Var}_{p}(\widehat{\theta}_{\mathrm{MC}})$
$D_{\mathrm{KL}}$	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 $\beta{=}1.25$	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 $\beta{=}1.5$	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 $\beta{=}1.75$	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 $\beta{=}2$	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

Appendix

Table 12: Per-prompt variance ratios for the losses and divergences experiments, pooling all eight model pairs. We compare the potential estimator

\widehat{\theta}({\vec{\phi}})

and its cross-fitted control-variate variant

\widehat{\theta}({\vec{\phi}},\alpha^{\star})

against Monte Carlo. The rightmost columns restrict attention to prompts for which

\widehat{\theta}({\vec{\phi}})

has higher variance than Monte Carlo and show how often the control variate reduces the ratio below one.

		std. dev.
estimator	mean	$M{=}1$	$M{=}5$	$M{=}10$
$\widehat{\theta}_{\mathrm{MC}}$	6.69	3.32	1.49	1.05
HT	6.69	3.32	1.49	1.05
Schulman CV, $\alpha=1$	6.47	35.67	15.95	11.28
Schulman CV, fitted $\alpha$	6.36	11.44	5.11	3.61
$\widehat{\theta}({\vec{\phi}})$	6.69	2.05	0.92	0.65

Appendix

Table 13: Reproduction of the KL-estimation experiment of Amini et al. (2025) : mean estimate and empirical standard deviation (nats) at sample sizes

M=1,5,10

, over

512

prompts with

4{,}000

samples each. The mean does not depend on

M

. Lowest standard deviation per column in bold.

Panel of Figure 5 — Figure 5: Gradient estimation along a DPO trajectory of the GPT-2 sentiment policy (checkpoints every $50$ steps; the variance of an estimator $\bm{\widehat{\theta}}$ is $\tr\mathrm{Var}_{p}(\bm{\widehat{\theta}})$ over the $512\times 20$ responses). (a) Variance of $\bm{\widehat{\theta}}(\bm{{\vec{\phi}}}_{\nabla})$ relative to the centered Monte Carlo gradient $\bm{\widehat{\theta}}_{\mathrm{MC}}$ , one series per quantity (regrets darker for larger $\beta$ ; undefined at step $0$ ). (b) Variance of the cross-fitted control-variate estimator $\bm{\widehat{\theta}}(\bm{{\vec{\phi}}}_{\nabla},\alpha^{\star})$ relative to $\bm{\widehat{\theta}}(\bm{{\vec{\phi}}}_{\nabla})$ . (c) The cross-fitted coefficient, mean over $8$ prompt folds with the range as band; $\alpha^{\star}=1$ recovers $\bm{\widehat{\theta}}(\bm{{\vec{\phi}}}_{\nabla})$ . Bands in (a) and (b) are $95\%$ prompt-cluster bootstrap intervals; one training run and one seed.

model	$1$ – $9$	$10$ – $99$	$100$ – $999$	$\geq 1{,}000$
GPT-2 ( $854$ )	$0.997$	$0.996$	$0.993$	$0.997$

Appendix

Table 16: Variance of the observed-prefix estimator relative to Monte Carlo at equal sample size (

M=4{,}000

): median over the unigram targets of GPT-2, banded by their occurrence count in the

4{,}000

samples (columns).

Panel of Figure 6 — Figure 6: Figure 2 for expected counts. Cost-adjusted variance ratio $\widehat{\rho}_{c}(w)$ for $n$ -gram count estimation under the per-target cost proxy $L_{w}+1$ . (a) Empirical cumulative distribution over $4{,}149$ targets with bootstrap RSE at most $2$ , with approximate $95\%$ log-ratio confidence bounds. (b) Ratio versus estimated expected count per document for the $2{,}831$ targets with at least $25$ 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 $M=4{,}000$ 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.

Panel of Figure 7 — Figure 7: $n$ -gram target inventory and token-length distribution. The $n$ -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 $L_{w}$ in the target. (a) Number of targets at each order $n$ , together with the minimum, maximum, mean, and median $L_{w}$ of their canonical realizations. The inventory contains $7{,}647$ targets in total. (b) Distribution of token lengths within each order. Increasing $n$ shifts the inventory toward longer targets, while GPT-2 tokenization also induces within-order variation in $L_{w}$ . (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

n

-gram frequency (left) and expected count per document (right), before applying the occurrence filter to the independently estimated

x

-axis quantity. In the main figure, we retain only targets whose independent frequency estimate has RSE at most

20\%

, corresponding to at least

25

expected occurrences in the estimation sample, to restrict the analysis to more reliable estimates. Here, we show all

3{,}835

targets with positive estimated frequency and a finite positive variance ratio; this remains a subset of the

7{,}647

-target inventory.

	(a) all targets		(b) $n$ -gram order				(c) token length $L_{w}$
frequency $\widehat{\theta}_{\mathrm{FR}}$	# targets	$10^{3}\widehat{\rho}_{c}$	$1$	$2$	$3$	$4$	$1$	$2$	$3$	$4$
$[10^{-8},10^{-7})$	354	1.5	2.0	1.5	1.5	1.3	–	1.6	1.5	1.3
$[10^{-7},10^{-6})$	701	5.8	6.5	6.6	5.3	5.1	–	6.6	5.4	5.1
$[10^{-6},10^{-5})$	725	28	35	28	25	17	42	26	25	17
$[10^{-5},10^{-4})$	684	130	160	110	85	–	160	110	85	–
$[10^{-4},10^{-3})$	299	310	330	260	–	–	330	260	–	–

Appendix

Table 17: Variance reduction is strongest for rare targets, even within fixed

n

-gram order and token length. The

2{,}831

targets of Figure 2 b are grouped by their independently estimated frequency

\widehat{\theta}_{\mathrm{FR}}

. Each entry reports the median cost-adjusted variance ratio

10^{3}\widehat{\rho}_{c}

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

n

-gram order and canonical token length

L_{w}

, 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

r_{s}

between

\widehat{\theta}_{\mathrm{FR}}

and

\widehat{\rho}_{c}

. Cells with fewer than

8

targets are omitted.

	Expected count		Frequency
Max. RSE	# targets	95% CI below $1$	# targets	95% CI below $1$
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%

Appendix

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

s

. We report the number of retained targets and the percentage satisfying

\widehat{\rho}_{c}<1/(1+1.96s)

, a sufficient condition for the approximate normal interval

\widehat{\rho}_{c}[1\pm 1.96\,\mathrm{RSE}]

to lie below parity. These are conservative shares for the normal interval, rather than the log-ratio intervals in Figures 2 and 6 .

Panel of Figure 18 — Algorithm 1 Potential selection by nested cross-fitting

Panel of Figure 9 — Figure 9: Nested cross-fitting, our in-sample selection among $p$ candidate potentials ( Algorithm 1 ). Each of $K=5$ outer folds is held out in turn (red); on the remaining folds the candidates are compared by inner cross-validation, the winner $j^{\star}$ 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 $39$ bins $20$ – $58$ : the $38$ bins on which MBAR’s own variance is itself measured to within $20\%$ , plus bin $57$ (at $25\%$ ), kept so that the range is contiguous; shading gives the bin probability, the dashed line the median ( $0.88$ ) and the band its $95\%$ interval. All $64$ populated bins are in Figure 12 .

Panel of Figure 11 — 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 $39$ bins ( $38$ by the $20\%$ rule plus bin $57$ ), with the bin probability as shading and the candidate’s own median over these bins with its $95\%$ interval. The terminal potential and its calibrated version remove $8\%$ , the logistic model $11\%$ , and the mixture model $13\%$ , against $12\%$ 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 $64$ populated bins. Filled markers denote the $38$ bins for which the MBAR variance is estimated with RSE below $20\%$ ; over these bins, the median variance ratio is $0.88$ with a $95\%$ bootstrap interval of $[0.84,0.90]$ , and over all $64$ populated bins it is $0.88$ with interval $[0.86,0.91]$ . Ratios in the remaining tail bins are much noisier because the MBAR variance itself is poorly estimated. Whiskers show $\pm 1.96$ bootstrap standard deviations from $1{,}000$ 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 $38$ reliable bins, the ratio shows no systematic dependence on baseline uncertainty (Spearman $\rho=0.05$ , $p=0.78$ ), 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 ( $\times 10^{6}$ ) [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, $\alpha$	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]

Appendix

Table 19: Event analysis: chain-level variance of each estimator relative to MBAR, median over all

64

populated ARI bins, with

95\%

intervals from

1{,}000

paired chain-bootstrap replicas (fitted potentials, selections, and coefficients held fixed). Sample-size efficiency: median over bins of

\mathrm{Var}(\mathrm{MBAR})/\mathrm{Var}(\text{estimator})

; MBAR-equivalent samples: that factor times the actual sample size

N=6.68\times 10^{6}

states.

Explore similar work

Jun 29, 2026cs.AI

Using Large Language Models as Low-Cost Statistical Estimators for Human-Response Data

Quantitative research across the social and behavioral sciences depends on human subject experiments that are expensive, slow, and subject to sampling bias. Here we show that pretrained large language models induce risk-equivalent estimators of conditional expectations under squared loss, establishing restricted functional risk equivalence: under squared loss, the LLM induces an estimator whose risk matches the Bayes optimal risk for squared-loss prediction of conditional expectations for any inference that depends on the data only through the conditional mean. We formalize the LLM as a misspecified functional estimator

T(\hat{P}_n)

trained on i.i.d.\ data, decompose the estimation error into representation bias

ε_{\mathrm{rep}}

and optimization error, and prove that under mild regularity conditions the LLM's expected error converges to the irreducible population variance plus the squared representation bias, with the representation bias bounded by the Pinsker inequality. The identifiability error

δ

propagates into the effective bias, inflating the asymptotic risk floor. We establish restricted functional risk equivalence via a bidirectional Le Cam deficiency analysis: the forward deficiency vanishes asymptotically while the reverse deficiency is exactly zero. We provide finite-sample concentration bounds and a calibration protocol with explicit decision rules. The result is a precise, provable statement: a well-calibrated LLM achieves the Bayes-optimal risk for conditional-mean-dependent inference, bounded by explicit scope conditions. In practical applications, this means that under satisfied conditions and well-calibrated models, large language models can be used in many prediction and decision-making tasks that originally relied on human experiments, approximating near-optimal statistical inference at lower cost.

Haobo Yang

Department of Computer Science and Engineering, SUSTech University

Language Model Calibration Statistical Learning Theory

Ask PDF

Jul 20, 2026cs.LG

Estimating Rare Events in Language Models with Proper Evaluation

Quantifying the risk of rare failures in language models, such as those triggered by adversarial distribution shifts or very large-scale deployments, requires estimating probabilities far too small for random sampling. While recent work has formalized Low Probability Estimation, existing pipelines remain fragile in the rarest regimes: estimators can suffer zero-estimate collapse or systematic bias, and standard evaluation losses can become unstable or poorly matched to asymmetric safety costs. In this work, we introduce Gradient Activation Adaptive Multi-Level Splitting (GA-AMLS), which adapts rare-event Monte Carlo methods to the continuous activation space of language models. Specifically, GA-AMLS uses a gradient-based MCMC kernel to navigate activation space, eliminating the zero-estimate collapse of input-space search and replacing the independence assumptions of prior activation-space estimators with conditional sampling under an explicit, heavier-tailed activation prior. We also propose the Shifted-Power Bregman (SPB) Loss, a proper scoring rule that remains finite for zero-estimates and offers tunable asymmetry between underestimation and overestimation penalties. Experiments on small transformer models reveal a bias-variance tradeoff: GA-AMLS achieves the lowest loss under symmetric evaluation, reducing average log-space squared error relative to the strongest baseline across model sizes, while methods with overestimation bias prevail under asymmetric penalties. Our findings highlight that estimator choice should be matched to deployment context. More broadly, our work establishes activation space as a tractable domain for rare-event estimation in language models, circumventing the brittleness of discrete input-space search.

Nikita Y. Parulekar, Anqi Liu

Johns Hopkins University

Language Model Safety Evaluation Proper Scoring Rules LLM Reliability Rare-Event Estimation

Ask PDF

May 22, 2026cs.CL

When Is Next-Token Prediction Useful? Marginalization, Ergodicity, Mixture Identifiability, Local Sufficiency, RAG, Tools, and Programming

Language models trained on observed sequences are often described as learning the conditional distribution of the next token given previous tokens. This description is only conditionally correct. A model trained on realized token trajectories does not observe full conditional laws; it receives sampled continuations. Moreover, real language generation is conditioned not only on previous words but also on non-textual circumstances: facts, events, intentions, goals, beliefs, social context, and task-specific constraints. This paper distinguishes three objects that are often conflated: the full conditional language process conditioned on latent circumstances, the marginal text-only process obtained by integrating those circumstances out, and the model-induced distribution learned from finite observed corpora. The paper argues that interpreting model training as estimating the marginal text-only law requires strong assumptions of stationarity, representativeness, and ergodicity, assumptions that are standard in statistical estimation but problematic when applied to heterogeneous language corpora. Even if these assumptions hold, the marginal text-only law is useful only when the observed prefix is an approximately sufficient statistic for the latent circumstances relevant to continuation. In information-theoretic terms, usefulness requires that the residual conditional mutual information between the next token and the omitted circumstances, given the observed text, be small. The paper then extends this argument to heterogeneous training corpora. Finally, the paper interprets Retrieval Augmented Generation (RAG) and tool use as conditional sufficiency devices.

Francesco Corielli

Autoregressive Language Modeling Next-Token Prediction Conditional Distribution Estimation Sequence Modeling

Ask PDF