One Trajectory Can Be Enough with Wise Step Selection for Semantic Uncertainty Quantification in Large Language Diffusion Models
Abstract
This paper deals with the high inference cost of sampling-based semantic un- certainty quantification (UQ) for Large Language Diffusion Models (LLDMs). We ask whether multiple independent generations can be replaced by informa- tive states from a single denoising trajectory. To this end, we introduce DIMS (Distribution-Matched Step Selection), which learns which diffusion steps best reproduce the semantic variability of independent generations. We show that the single-trajectory approximation error decomposes into distribution mismatch and within-trajectory redundancy, leading to a simple convex MMD-based objective. DIMS requires only a small unlabeled calibration set and no correctness labels. Across three LLDMs, six datasets, and eight semantic UQ estimators, DIMS is best or tied for best among single-trajectory strategies in 34/40 settings. Overall, our results show that sampling itself can be optimized for efficient UQ. Code will be released upon publication.
est. 32% chance this paper gets accepted at ICLR 2027.
What do you think this paper will get?
All positions stay anonymous.