Eliciting Numerical Predictive Distributions of LLMs Without Autoregression

Julianna Piskorz; Katarzyna Kobalczyk; Mihaela van der Schaar

cs.LGMar 3, 2026

Eliciting Numerical Predictive Distributions of LLMs Without Autoregression

Authors: Julianna Piskorz, Katarzyna Kobalczyk, Mihaela van der Schaar

Organizations: Department of Applied Mathematics and Theoretical Physics, University of Cambridge

Abstract

Large Language Models (LLMs) have recently been successfully applied to regression tasks -- such as time series forecasting and tabular prediction -- by leveraging their in-context learning abilities. However, their autoregressive decoding process may be ill-suited to continuous-valued outputs, where obtaining predictive distributions over numerical targets requires repeated sampling, leading to high computational cost and inference time. In this work, we investigate whether distributional properties of LLM predictions can be recovered without explicit autoregressive generation. To this end, we study a set of regression probes trained to predict statistical functionals (e.g., mean, median, quantiles) of the LLM's numerical output distribution directly from its internal representations. Our results suggest that LLM embeddings carry informative signals about summary statistics of their predictive distributions, including the numerical uncertainty. This investigation opens up new questions about how LLMs internally encode uncertainty in numerical tasks, and about the feasibility of lightweight alternatives to sampling-based approaches for uncertainty-aware numerical predictions.

Figures & tables

Panel of Figure 1 — Figure 1: Illustration of this paper’s goals and methodology.

$y_{i}^{*}$ , target	$\hat{y}_{i}^{*}$ (ours)	$\bar{\mathbf{x}}$	$\bar{\mathbf{x}}_{i}$	$x_{i,n}$
$*=$ mean	0.006	0.256	0.035	0.085
$*=$ median	0.006	0.260	0.041	0.087
$*=$ greedy	0.015	0.273	0.065	0.109

Table 1: MSE obtained when predicting the statistics of the LLM’s predictive distribution.

Figure 2: Predicted vs. true values of mean, median and greedy prediction, presented on $\log_{10}$ scale. The probing model accurately recovers the number that the LLM intends to predict, indicating that the internal representations encode the order of magnitude of prediction.

Figure 3: Predicted vs. sample-based IQR (both median-normalised). The model accurately tracks the variability of the LLM’s output distribution.

$\alpha$	50%	90%	95%
1.0	52.0 ± 0.4	90.9 ± 0.3	95.5 ± 0.2
10.0	52.7 ± 0.5	91.3 ± 0.3	96.1 ± 0.2
1000.0	51.4 ± 0.3	90.7 ± 0.3	95.7 ± 0.2
10000.0	48.2 ± 0.3	90.5 ± 0.2	95.4 ± 0.2

Table 2: Coverage of the predicted confidence intervals. Values denote empirical coverage (%) ± standard error.

Table 3: MSE of predicted values vs. the ground truth

x_{i,n+1}

Figure 4: The probe achieves comparable error to using 20-25 LLM samples on the one step ahead prediction task.

Figure 5: Generalisation to unseen context lengths. A probe trained on a restricted context length range (Restricted) exhibits greater deviation in empirical coverage outside its training range.

$\alpha$	50%	90%	95%
Model
Real (all)	48.8 ± 0.1	88.5 ± 0.1	94.3 ± 0.1
Real (5 fold)	43.4 ± 0.2	82.1 ± 0.2	89.4 ± 0.2
Synth	30.2 ± 0.3	67.7 ± 0.4	77.3 ± 0.4

Table 4: Coverage of the CI intervals on previously unseen testing inputs.

Figure 6: Absolute error on the median across different sub-dataset. Comparison of generalisation across models trained on different data.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Specification	NC24rs_v3	NC80adis_H100_v5
vCPUs	24	80
System Memory (GiB)	448	640
GPU Model	4× NVIDIA Tesla V100	2× NVIDIA H100 NVL
GPU Memory (per GPU)	16 GiB	94 GiB
Total GPU Memory	64 GiB	188 GiB
GPU Architecture	Volta	Hopper

Appendix

Table 5: Azure Virtual Machine Specifications

Function name	Formula	$a$ -range
sin	$\sin(x)$	[0.5, 6.0]
linear_sin	$0.2\cdot\sin(x)+\frac{x}{450}$	[0.5, 6.0]
sinc	$\text{sinc}(x)$	[0.05, 0.2]
xsine	$\frac{x-30}{50}\cdot\sin(x-30)$	[0.5, 1.3]
beat	$\sin(x)\cdot\sin\left(\frac{x}{2}\right)$	[0.1, 6.0]
gaussian_wave	$e^{-\frac{(x-2)^{2}}{2}}\cdot\cos(10\pi(x-2))$	[0.01, 0.1]

Appendix

Table 6: Functions used to generate time series data, their mathematical forms, and the range of the time-scaling parameter

a

Hyperparameter	Description	Default Value
min_mag	Minimum exponent for base-10 magnitude scaling (as used by $f_{\text{order}})$	$-3$
max_mag	Maximum exponent for base-10 magnitude scaling	$\log_{10}(D_{\text{scale}}$ )
alpha	Weight for classification loss component	100.0
beta	Weight for regression loss component	50.0
K	Top- $K$ exponents taken into consideration (see Equation 4)	3
hidden_layers	Number of hidden layers in feature extractor	1

Appendix

Table 7: Model-specific hyperparameters for the magnitude-factorised models.

Hyperparameter	Description	Default Value
learning_rate	Learning rate for the optimiser	$10^{-5}$
weight_decay	L2 regularization weight	1.0
scheduler_step_size	Learning rate scheduler step size	100
scheduler_gamma	Learning rate scheduler gamma	0.5
batch_size	Number of samples per training batch	1024
max_epochs	Number of training epochs	600

Appendix

Table 8: Optimiser and training-related hyperparameters.

predicted value	MFP probe (ours)	MLP probe (log-scaling)	MLP probe (no log-scaling)
greedy	0.0150	0.0400	0.0255
mean	0.0061	0.0091	0.0092
median	0.0058	0.0101	0.0100

Appendix

Table 9: MSE of the values predicted with the magnitude-factorised (MFP) model vs. a vanilla MLP.

Panel of Figure 7 — Figure 7: Predicted vs. sample mean, median and greedy prediction (on $\log_{10}$ scale). Results analogous to Figure 2 .

Panel of Figure 8 — Figure 8: Predicted vs. sample mean, median and greedy prediction on the dataset with scale $D_{\text{scale}}=1.0$ .

Table 10: MSE for the predictions on the dataset with scale

D_{\text{scale}}=1.0

, reported for all models. Results analogous to Table 1 .

Table 11: Coverage of the CI for all models. Results analogous to Table 2 .

Panel of Figure 9 — Figure 9: Sample efficiency on the task of one step ahead predictions. Results analogous to Figure 4 . The shaded regions mark 95% confidence sets, computed using 100 bootstrap samples. We use $D_{\text{scale}}=1.0$ .

Panel of Figure 10 — Figure 10: Predicted IQR vs. Sample IQR (median adjusted). Results analogous to Figure 3 .

$n$ samples	1	5	10	20	50	100
time (s)	1.59 $\pm$ 0.08	1.62 $\pm$ 0.07	1.71 $\pm$ 0.07	1.83 $\pm$ 0.07	2.35 $\pm$ 0.08	3.28 $\pm$ 0.08

Appendix

Table 12: Average time to generate

n

LLM samples.

Panel of Figure 11 — Figure 11: Layer ablation for Llama-2-7B on the probing model from section 2 .

Figure 12: Layer ablation for Llama-2-7B on the quantile probing model from section 3 . MSE on the median (left, lower is better) and the Pearson R between predicted IQR and sample IQR (right, higher is better).

quantile level $\tau^{s}$	0.025	0.05	0.25	0.5	0.75	0.95	0.97
dataset
1.0	0.58	0.40	0.16	0.05	0.15	0.40	0.60
10.0	3.12	2.16	0.45	0.28	0.49	1.91	3.26
1000.0	111.16	61.54	14.83	12.10	14.30	46.04	115.71
10000.0	843.11	380.06	119.80	100.41	118.90	483.33	1172.27

Appendix

Table 13: Quantile regression – mean absolute error per quantile . Values represent the mean absolute error between the empirical quantile

\tilde{q}^{s}

and the predicted quantile

\hat{q}^{s}

for all quantile levels.

LLM Uncertainty Estimation

Large language model uncertainty estimation measures when outputs may be wrong, ambiguous, or unreliable to guide safer decisions.

See all 398 papers

May 20, 2026

Reinforcement Learning over Predictive Distributions for LLM Regression

Large language models (LLMs) have emerged as flexible regressors capable of predicting real-valued quantities from heterogeneous inputs. Yet most LLM regression objectives optimize predictions independently, often yielding poor calibration. We introduce Distribution-Aware Reward (DAR), an on-policy reinforcement learning objective that instead jointly evaluates the empirical predictive distribution formed by multiple predictions for the same input. To translate this distribution-level objective into rollout-level rewards, we assign each prediction credit based on its leave-one-out contribution to the quality of the overall predictive distribution. This encourages predictions that are well-centered and appropriately dispersed around the target. We evaluate on three regression settings: a synthetic task probing interpolation and extrapolation, and two real-world scientific tasks involving code and molecular data. Across tasks, DAR produces better-calibrated uncertainty estimates while consistently reducing prediction error and improving ranking quality over supervised fine-tuning and pointwise reinforcement learning. Together, these results highlight the benefits of distribution-aware training for LLM regression.

Jungsoo Park, Hyungjoo Chae, Ethan Mendes +4

Georgia Institute of Technology · Allen Institute for AI

Distributional RL RL for Language Models Language Model Calibration

Ask PDF

Jun 4, 2026

UnpredictaBench: A Benchmark for Evaluating Distributional Randomness in LLMs

We introduce UnpredictaBench, an evaluation that tests the ability of large language models (LLMs) to capture true underlying distributions. As LLMs are increasingly used as substitutes for other entities (e.g., for humans in economic simulations), the tendency of many models to collapse towards a single plausible answer means a failure to capture the unpredictability of real systems. Recent work on improving output diversity is insufficient for this setting: simulation requires samples that are calibrated to a target distribution, not merely varied outputs. UnpredictaBench isolates a simplified but fundamental version of this problem: sampling outcomes from individual target distributions, including canonical statistical distributions, distributions induced by stochastic programs, and natural-language scenarios that describe random processes. We introduce 448 such problems together with KS@N, a general-purpose evaluation metric that quantifies how well a model outputs approximate black-box target distributions via the Kolmogorov-Smirnov statistical test. This is the rate at which we fail to reject model samples of size N against ground-truth samples, with larger N indicating greater difficulty. Tested across open and proprietary models, we find a large spread in distributional capabilities. For instance, when models generate samples of size 100 (KS@100, our standard metric), scores range from near 0 to over 20%. No model is able to achieve over 40% at KS@100, showing significant headroom in distributional sampling as a capability. Although adding reasoning can somewhat increase scores, we find no immediate solution for this issue. UnpredictaBench shows that even simple distributional simulation remains challenging, making it a necessary first step toward using LLMs as stand-ins for complex systems. Project website and resources are available at https://unpredictabenchmark.github.io/.

Amirhossein Abaskohi, Amirhossein Dabiriaghdam, Liang Luo +4

University of British Columbia · 2Independent Researcher

LLM Evaluation Language Model Generation Evaluation Synthetic Benchmark Generation

Ask PDF

Jun 26, 2026

Enhancing Numerical Prediction in LLMs via Smooth MMD Alignment

Despite their strong general capabilities, large language models (LLMs) often remain unreliable when outputs must be numerically precise. A key reason is the training objective: standard cross-entropy treats numeric tokens as unstructured categories and ignores the metric structure of their values. We address this mismatch with Smooth Maximum Mean Discrepancy (SMMD), which builds on the classic MMD by incorporating value-distance kernels over numeric tokens and graph-based smoothness. With this kernel defined over a numeric sub-vocabulary, SMMD aligns the predicted numeric distribution to the target via kernel matching and smooths the prediction-target residual over the induced kernel graph to encourage local consistency. We evaluate SMMD on four numeric-target tasks: mathematical reasoning, arithmetic calculation, clock-time recognition, and chart question answering, across multiple open-weight LLM and VLM backbones. SMMD consistently improves accuracy over both cross-entropy and recent numeric-target losses; analyses show complementary effects between MMD and smoothness and underscore the importance of distance-based kernel design. Code is available at https://github.com/Zuozhuo/smmd-loss.

Zhuo Zuo, Li Yue, Wenhao Zheng +2

College of Computer Science, Sichuan University, Chengdu, China · Dongfang Electric (Chengdu) Academy of Science and Technology Co., Ltd., Chengdu, China · Institute of Medical Information & Library, CAMS & PUMC, Beijing, China

Numerical Reasoning in Language Models Maximum Mean Discrepancy Language Modeling LLM Training

Ask PDF

Abstract

LLM Uncertainty Estimation

Explore similar work

Reinforcement Learning over Predictive Distributions for LLM Regression

UnpredictaBench: A Benchmark for Evaluating Distributional Randomness in LLMs

Enhancing Numerical Prediction in LLMs via Smooth MMD Alignment