Matchability-Aware Conformal Prediction for Open-Ended Language Model Generation
NeurIPS
Published
Abstract
Large language models (LLMs) often do not produce calibrated uncertainty estimates for their answers. We propose CAELUM, a framework that builds semantic prediction sets while separating aleatoric from epistemic uncertainty. CAELUM samples answers using an ensemble of clarification prompts, clusters them by semantic equivalence, and estimates the aleatoric and epistemic components via kernel-based von Neumann entropy before calibrating nonconformity scores with split conformal prediction. We use a weight \(w_\lambda(x)=1+U_{\mathrm{ale}}(x)+\lambda\,U_{\mathrm{epi}}(x)\) to control for the contribution of the epistemic uncertainty of each instance and fit a coefficient \(\lambda^\star\) to minimize prediction set size. When the objective is nonflat and selection is stable across splits, \(\lambda^\star\) could be a diagnostic, where large values may indicate epistemic dominance (retrieve, abstain) and \(\lambda^\star\to 0\) aleatoric dominance (clarify). For individual queries, a log ratio \(R(x)\) compares the weighted epistemic component with the aleatoric component. Theoretically, we demonstrate that the coverage is capped by matchability, the fraction of items whose reference answer appears in any sampled cluster. Empirically, this matchability can be predicted from the original question features using a lightweight classifier. We use those predictions to define a stratum and derive a conformal coverage bound within it. Across 30 LLM and dataset combinations (4–32B), mean coverage is 0.74 in the learned stratum versus 0.65 overall, averaged over the 78 evaluations using the learned classifier.