Dist2ill: Distributional Distillation for One-Pass Uncertainty Estimation in Large Language Models
Abstract
Large Language Models (LLMs) often exhibit misalignment between the quality of their generated responses and the confidence estimates they assign to them. Bayesian treatments, such as marginalizing over a reliable weight posterior or over the space of reasoning traces, provide an effective remedy, but incur substantial computational overhead due to repeated sampling at test time. To enable accurate uncertainty estimation in a single forward pass, we propose a novel distributional distillation framework (Dist2ill) that trains an LLM to produce multiple diverse reasoning paths within one inference pass, while using a lightweight parametric module to approximate empirical confidence scores derived from the sampling distribution. Extensive experiments demonstrate that Dist2ill preserves reasoning diversity and achieves state-of-the-art uncertainty estimation, substantially improving Expected Calibration Error (ECE) and Negative Log-Likelihood (NLL), while remaining computationally efficient.
Keywords
Cite
@article{arxiv.2505.11731,
title = {Dist2ill: Distributional Distillation for One-Pass Uncertainty Estimation in Large Language Models},
author = {Yicong Zhao and King Yeung Tsang and Harshil Vejendla and Haizhou Shi and Zhuohang Li and Zhigang Hua and Qi Xu and Tunyu Zhang and Yi Wang and Ligong Han and Bradley A. Malin and Hao Wang},
journal= {arXiv preprint arXiv:2505.11731},
year = {2026}
}
Comments
Preprint; work in progress. Update Log: 05/2025 (v1&v2): Introduced Dist2ill (previously named EUD) for efficient uncertainty estimation, focusing on discriminative reasoning tasks. 02/2026 (v3): Extended Dist2ill to a unified framework supporting both discriminative and generative reasoning