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 399 papers

Reinforcement Learning over Predictive Distributions for LLM Regression
May 20, 2026Jungsoo Park, Hyungjoo Chae, Ethan Mendes +4Distributional RL RL for Language Models
Georgia Institute of Technology · Allen Institute for AI
UnpredictaBench: A Benchmark for Evaluating Distributional Randomness in LLMs
Jun 4, 2026Amirhossein Abaskohi, Amirhossein Dabiriaghdam, Liang Luo +4LLM Evaluation Language Model Generation Evaluation
University of British Columbia · 2Independent Researcher
Enhancing Numerical Prediction in LLMs via Smooth MMD Alignment
Jun 26, 2026Zhuo Zuo, Li Yue, Wenhao Zheng +2Numerical Reasoning in Language Models Maximum Mean Discrepancy
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

May 20, 2026J/K move · Enter open · S save

Reinforcement Learning over Predictive Distributions for LLM Regression

Jungsoo Park, Hyungjoo Chae, Ethan Mendes, Jay DeYoung+3

Georgia Institute of Technology · Allen Institute for AI

Distributional RL RL for Language Models Language Model Calibration

Open paper AskPDF arXiv

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.

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