Set-Valued Prediction for Large Language Models with Feasibility-Aware Coverage Guarantees
Authors: Ye Li, Anqi Hu, Yuanchang Ye, Shiyan Tong, Zhiyuan Wang, Bo Fu
Organizations: University of Electronic Science and Technology of China, Chengdu, Sichuan, China · Zhejiang University of Finance and Economics, Hangzhou, Zhejiang, China · Southeast University, Nanjing, Jiangsu, China
Large language models (LLMs) inherently operate over a large generation space, yet conventional usage typically reports the most likely generation (MLG) as a point prediction, which underestimates the model's capability: although the top-ranked response can be incorrect, valid answers may still exist within the broader output space and can potentially be discovered through repeated sampling. This observation motivates moving from point prediction to set-valued prediction, where the model produces a set of candidate responses rather than a single MLG. In this paper, we propose a principled framework for set-valued prediction, which provides feasibility-aware coverage guarantees. We show that, given the finite-sampling nature of LLM generation, coverage is not always achievable: even with multiple samplings, LLMs may fail to yield an acceptable response for certain questions within the sampled candidate set. To address this, we establish a minimum achievable risk level (MRL), below which statistical coverage guarantees cannot be satisfied. Building on this insight, we then develop a data-driven calibration procedure that constructs prediction sets from sampled responses by estimating a rigorous threshold, ensuring that the resulting set contains a correct answer with a desired probability whenever the target risk level is feasible. Extensive experiments on six language generation tasks with five LLMs demonstrate both the statistical validity and the predictive efficiency of our framework.
Figures & tables
Figure 1: Overview of our feasibility-aware calibration framework with statistically rigorous coverage guarantees.
Figure 2: (a-c) Comparison between point-prediction accuracy and set-valued attainability under different semantic matching thresholds. The consistent gap shows that point prediction systematically under-utilizes admissible answers already present in the sampled candidate space, and (d) that this gap widens under stricter semantic evaluation.
Figure 3: Coverage Guarantees on six NLG benchmarks utilizing five LLMs. The threshold of sentence similarity is fixed at 0.7.
Figure 4: Coverage Guarantees on six NLG benchmarks utilizing five LLMs. The threshold of sentence similarity is fixed at 0.5.
Figure 5: Coverage Guarantees on six NLG benchmarks utilizing five LLMs. The threshold of sentence similarity is fixed at 0.6.
Figure 6: Sampling budget shifts the feasibility boundary. We plot the attainable upper bound 1−αl(K) as a function of the sampling budget K under different semantic admission thresholds. A larger K lowers the MRL αl(K) , thereby enlarging the feasible region for coverage guarantees. The improvement is consistent across settings but gradually saturates.
Figure 7: APSS vs. risk level. APSS decreases monotonically as α increases, reflecting the trade-off between prediction efficiency and coverage guarantee. Stricter admission thresholds generally require larger prediction sets and induce a higher MRL αl .
Figure 8: Effect of semantic deduplication on prediction efficiency. We merge responses with sentence similarity above 0.9 within each prediction set.
Figure 9: Ablation of reliability-score components in the confidence evaluation function F .
Figure 10: Sensitivity to the split ratio on TriviaQA. We compare the default 50% calibration split with a reduced 10% split. Evaluation is restricted to the common feasible regime determined by the larger MRL. The similar coverage and APSS curves indicate that the proposed framework remains robust and calibration-efficient even with substantially fewer calibration samples.
Large language models (LLMs) are prone to generating factually incorrect outputs. Recent work has applied conformal prediction to provide uncertainty estimates and statistical guarantees for the factuality of LLM generations. However, existing approaches are typically not prompt-adaptive, limiting their ability to capture input-dependent variability. As a result, they may filter out too few items (leading to over-coverage) or too many (under-coverage) for a given task or prompt. We propose an adaptive conformal prediction approach that extends conformal score transformation methods to LLMs, with applications to long-form generation and multiple-choice question answering. This enables prompt-dependent calibration, retaining marginal coverage guarantees while improving conditional coverage. In addition, the approach naturally supports selective prediction, allowing unreliable claims or answer choices to be filtered out in downstream applications. We evaluate our approach on multiple white-box models across diverse domains and show that it significantly outperforms existing baselines in terms of conditional coverage.
Aleksandr Rubashevskii, Dzianis Piatrashyn, Preslav Nakov +1
Mohamed bin Zayed University of Artificial Intelligence (MBZUAI), Abu Dhabi, UAE
Prediction sets provide a theoretically grounded framework for quantifying uncertainty in machine learning models. Adapting them to structured generation tasks, in particular, large language model (LLM) based code generation, remains a challenging problem. An existing attempt proposes PAC prediction sets but is limited by its strong monotonicity assumption on risk and single-label classification framework, which severely limits the space of candidate programs and cannot accommodate the multiple valid outputs inherent to code generation. To address these limitations, we propose an approach RisCoSet that leverages multiple hypothesis testing to construct risk-controlling predictions for LLM-based code generation. Given a trained code generation model, we produce a prediction set represented by a partial program, which is guaranteed to contain a correct solution with high confidence. Extensive experiments on three LLMs demonstrate the effectiveness of the proposed method. For instance, compared with the state-of-the-art, our method can significantly reduce the code removal by up to 24.5%, at the same level of risk.
Senrong Xu, Yuhao Tan, Yanke Zhou +6
State Key Lab of Novel Software Technology, Nanjing University · ETH Zürich · Birkbeck, University of London
Accurate confidence estimation is critical for reliable use of large language models (LLMs). Prior work on LLM calibration largely focuses on response-level confidence, which estimates the correctness of a single generated output. However, this formulation is misaligned with many practical settings where the central question is how likely a model is to solve a query overall. We show that this mismatch results from the stochastic nature of modern LLM decoding, under which single-response correctness fails to reflect underlying model capability. To address this issue, we introduce capability calibration, a new evaluation framework for measuring how well query-level confidence aligns with a model's expected accuracy on individual queries. We formally distinguish capability calibration (CC) from response calibration (RC) and show that the two differ both theoretically and empirically. We further show that CC is better suited than RC to applications like pass@k prediction and inference budget allocation. Finally, we evaluate common confidence estimation methods to understand the practical feasibility of CC.