CAELUM: Matchability-Aware Conformal Prediction for Open-Ended Language Model Generation

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_λ(x) = 1 + U_ale(x) + λ U_epi(x) to control for the contribution of the epistemic uncertainty of each instance and fit a coefficient λ^⋆ to minimize prediction set size. When the objective is nonflat and selection is stable across splits, λ^⋆ could be a diagnostic, where large values may indicate epistemic dominance (retrieve, abstain) and λ^⋆ → 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.

1Sampled support and matchability

For each question, CAELUM samples answers from the model and groups answers with similar meanings into semantic clusters. A prediction set is a subset of these clusters. Matchability records whether the reference answer matches at least one sampled cluster under a fixed evaluation rule.

If the reference answer matches no sampled cluster, even the collection of all clusters will miss it. Coverage over a population is therefore at most the matchability rate π, no matter what threshold is used in the conformal step.

Figure 1: Qwen3-8B on TriviaQA, question 39

“The Museum of Bags and Purses is in which European city?”

Reference: Amsterdam
Sampled: 50 answers in two clusters, represented by Antwerp (72%) and Madrid (28%). Amsterdam is not among them.

Cutoff η = 0.35

At η = 0.35 the set is {Antwerp}. It does not cover Amsterdam.

A cluster joins the set when the mass of the clusters before it is at most η, so Madrid joins at η = 0.72. The cutoff here is moved by hand. It does not fit λ or recalibrate a conformal threshold. Every setting misses Amsterdam.

2Building prediction sets

In our experiments, the model produces five clarifications for each question and ten answers to each clarification. Each question therefore gets 50 answers.

SampleThe model produces n clarifications of each question. For factual questions, these are rewordings; for ambiguous ones, they are intended to disambiguate the question. It then produces m answers for each clarification.
ClusterWe cluster the nm answers by the similarity of their sentence embeddings. The cluster frequencies p_k define an empirical distribution over candidate meanings.
Measure uncertaintyThe kernel von Neumann entropy of all pooled answers gives U_total. Averaging the entropy within each clarification gives U_epi. The residual is U_ale = max(U_total − U_epi, 0).
CalibrateWe divide the cumulative cluster score by w_λ(x) = 1 + U_ale(x) + λ U_epi(x). One split selects λ; a separate split fits the conformal threshold τ̂.
PredictFor a new question, consider clusters in decreasing order of p_k. Include each cluster when the mass of all preceding clusters is at most η = min{τ̂ w, 1}.

Figure 2 gives the overview from our paper, with the two calibration constants presented as controls. The first question is ambiguous and has three valid answers. The second and third are knowledge gaps: for the second, the reference was sampled; for the third, it never was.

Figure 2: Three example queries (illustrative), recomputed as you move the controls

Epistemic weight λ = 2.0 Threshold τ̂ = 0.39

Constructed illustration with n = 3 clarifications, m = 10 answers and cluster Shannon entropy, as in Figure 1 of our paper. The experiments use n = 5 and kernel entropy. Each row of dots is one clarification, and the color of a dot is its semantic cluster. Here τ̂ stays fixed while λ moves. CAELUM refits τ̂ for every λ, which offsets much of the change, so this view overstates how much λ matters. The museum column uses a constructed 18/12 split, where the real run in Figure 1 gave 36/14.

Within a clarification

The epistemic component captures the variability in answers that remains once we condition on a clarification. It gives a modest empirical signal for detecting failures of the sampled support.

Between clarifications

The aleatoric residual captures variability introduced by the clarification ensemble. It depends on the clarifications themselves and is not a verified measure of ambiguity in the original question.

Normalization makes the calibration threshold depend on λ. Without normalization, the calibration scores are independent of λ. Set size can only increase as λ grows, so the search returns λ^⋆ = 0 under our rule of choosing the smallest weight in a tie. Even with normalization, discrete sets can stay the same across much of the λ grid. A coefficient fitted to a flat objective cannot identify an uncertainty regime. Even when the objective distinguishes weights, an interpretation requires stability across splits.

3Coverage bounds

A logistic classifier, m̂, takes features of a question and predicts whether it is matchable, prior to generating an answer. This classifier is trained on one data split and then fixed. Mondrian calibration uses only the calibration questions that the classifier predicts will be matchable.

Let ρ be the proportion of questions in this subset for which there is no corresponding reference in the support. Assuming exchangeability of the calibration and test examples, including their generations and any score randomization, and fixed matching pipeline and score, Theorem 1 gives us:

1 − α − ρ ≤ ℙ(y* ∈ 𝒞(x) | m̂(x) = 1) ≤ 1 − ρ.

The lower bound allows for false matchability and conformal error. The upper bound is the proportion of matchable questions in this subset. In its group, a classifier with ρ = 0.20 caps coverage at 0.80 at any given conformal threshold.

Figure 3: Twenty questions (illustrative), ordered by the classifier's score

Questions the classifier calls matchable: 14

False matchability ρ1/14 = 0.07

Bound in the stratum[0.83, 0.93]

Share of questions selected

Population lower bound0.58

With 14 questions selected, one of them is unmatchable.

Constructed population, the twenty questions in the lower panel of Figure 1 in our paper, of which 15 are matchable. Here α = 0.10. The population bound is (1 − α) ℙ(m̂ = 1) − ℙ(M = 0, m̂ = 1) from the same theorem. Selecting fewer questions usually lowers ρ and tightens the bound, but the guarantee then speaks for fewer questions. The bound is not monotone in the cut. It drops each time the selection takes in an unmatchable question, as at 10 and 16, and recovers slowly as matchable ones follow. The classifier's precision and the share it selects have to be reported together.

How does the support-conditional proposition differ?

Proposition 1 produces calibration of matchable questions and guarantees at least 1 − α coverage, given matchability. This is calibration over matched examples, different from the calibration in our main experiment. In general, the stratum considered in Theorem 1 can include some unmatched questions. So in that theorem's bound we include ρ. If our filter has ρ = 0, we could guarantee nominal coverage of questions in its group. And if in addition it included all matchable questions, we would recover Proposition 1.

What must be fixed before calibration?

To get the guarantee in Theorem 1, we would fix the scoring rule, λ, classifier, encoder, clustering threshold and reference-matching rule independently of the calibration and test points. They could be built on a separate data set and then fixed. Taken across a stratum, this gives us a guarantee. It is not a guarantee for every question.

4Results

Six QA datasets and three generators give us 18 cells, with 2,104 questions per generator. We use the same sampling budget across Gemma-4-E4B-it, Ministral-3-8B-Instruct-2512 and Qwen3-8B. Below, we follow the format of our main table. There, we use deterministic APS with seed 42, 40/30/30 split.

Figure 4: Test coverage against test matchability

MODELS

DATASETS

Each mark is one model–dataset cell. The diagonal is the exact support ceiling for those test items, and the shaded half above it cannot be reached.

Sixteen of the 18 cells sit on the diagonal, so on those cells every matchable test reference was covered. The two AmbigInst cells below it are the only ones where a sampled reference was left out of the set.

Most cells stay well below the 0.90 target, and the distance to the target is set by matchability.

Click a model or a dataset to show or hide its cells. Pointing at a dataset highlights its cells. Both axes use the same test split. Marks that share a value overlap, and hovering or focusing a mark lists every cell at that point. The table also gives each cell's match rate over all its questions, which uses a different denominator from the test-split values in the plot.

The 18 cells as a table

Deterministic APS at α = 0.10 on a 40/30/30 split. Test π and coverage use the same held-out questions. Full-cell π uses all questions in each model–dataset cell.
Generator	Dataset	Full-cell π	Test π	Test coverage

Matchability prediction

Including Qwen3-32B and Qwen3-14B on these six datasets means we have five generators and 30 cells in total. Across four seeds of the split, mean coverage in the predicted-matchable stratum of our 78 learned classifiers was 0.74 ± 0.15. Across all test questions in the same rows, it was 0.65 ± 0.17. On these rows, the calibration split AUC of the classifier was 0.67 ± 0.08. Rows in which the training or calibration data had only one class or in which its AUC was under 0.55 are shown separately. These are oracle-fallback reference results. Neither coverage reaches our target of 0.90.

Taken across six generator comparisons, the classifier on the questions of each generator achieves a mean AUC of 0.78 on the target labels and 0.80 on the source labels. Each generator sees the same questions. Where a question appears in more than one dataset, its label can be different in each.

Set size and uncertainty scores

Across the four split seeds, the aleatoric-only baseline and CAELUM each produce about 2.21 clusters per set. Differences in set size between them within each cell are smaller than variation across seeds.

For ranking failures of the sampled support, U_epi beats U_total on 15 of 17 non-degenerate cells across seeds. The simpler U_epi signal also beats R(x) on 16 of 18 cells in the separate routing analysis.

Cross-generator transfer of the classifier, all six directions

2,104 aligned question records per pair, reused for fitting and scoring. AUC_src uses the source model's labels and AUC_tgt scores the same frozen classifier against the target model's labels. Independent labels would agree 0.50–0.53 of the time.
Source → target	AUC_src	AUC_tgt	Δ	Agreement
Qwen3 → Ministral	0.793	0.771	−0.022	0.831
Qwen3 → Gemma	0.793	0.797	+0.004	0.817
Ministral → Qwen3	0.788	0.771	−0.017	0.831
Ministral → Gemma	0.788	0.782	−0.006	0.796
Gemma → Qwen3	0.821	0.782	−0.039	0.817
Gemma → Ministral	0.821	0.766	−0.055	0.796
Mean	0.801	0.778	−0.023	0.815

5Model size and matchability

We compare Qwen3-8B and Qwen3-32B on the same ordered questions and references. Both use FP16, generation seed 42 and five clarifications × ten answers. The figure reports heuristic match rates over all questions. It does not report calibrated coverage.

Figure 5: Share of questions whose reference was sampled

Click a model above the plot to show or hide it. Both generators use the same sampling budget. The larger model adds support on five of six datasets, and SimpleQA and PopQA stay at 36.6% and 57.4%. This audit does not test classifier transfer across model scales or an advantage in epistemic ranking. The 30-cell analyses use separate answer samples and protocols.

Matched items per dataset

Matched items over dataset size, with the match rate in parentheses.
Dataset	Qwen3-8B	Qwen3-32B

6Limitations

The learned classifier is weak on some cells with low matchability, including PopQA. Its precision, recall, selection rate and induced ρ must be evaluated at the intended operating point. AUC alone does not establish a useful coverage stratum.
Numerical matchability depends on semantic clustering and on the reference matcher. The judge audit examines heuristic-unmatched cases and identifies additional matches. That audit tests sensitivity to missed matches, and it leaves the heuristic-positive cases unchecked.
The uncertainty decomposition measures the clarification pipeline. It does not establish why a model missed an answer, and an unstable λ or R(x) cannot prescribe retrieval, abstention or clarification.
Abstention pilots add an outcome to the label space. Their augmented coverage must be reported alongside ordinary semantic coverage and the cost of abstaining. They do not repair the pre-generation classifier, which reads only the question text.
Our experiments concern English short-answer QA with references that can be checked. Longer text with several claims, and generation over several turns, would need a separate formulation of coverage.

Citation

@inproceedings{
azizi2026matchabilityaware,
title={Matchability-Aware Conformal Prediction for Open-Ended Language Model Generation},
author={Ilia Azizi},
booktitle={The Fortieth Annual Conference on Neural Information Processing Systems},
year={2026},
url={https://openreview.net/forum?id=eY9UXFg5E0}
}