Domain-Shift-Aware Conformal Prediction for Large Language Models

Technical Summary: Domain-Shift-Aware Conformal Prediction for Large Language Models

Problem Statement
Large Language Models (LLMs) frequently produce overconfident yet factually incorrect outputs (hallucinations), posing significant risks in high-stakes domains like healthcare and finance. While Conformal Prediction (CP) offers a principled framework for uncertainty quantification (UQ) with finite-sample, distribution-free coverage guarantees, standard CP relies on the exchangeability assumption. This assumption requires that calibration and test data be drawn from the same distribution. In real-world deployments, domain shift is pervasive; test prompts often differ significantly from historical calibration data due to evolving user behavior or input distributions. Under such shifts, standard CP frequently suffers from under-coverage, producing prediction sets that are too narrow to contain the true answer, thereby failing to provide reliable uncertainty guarantees.

Methodology: DS-CP Framework
The authors propose Domain-Shift-Aware Conformal Prediction (DS-CP), a framework designed to adapt CP to LLMs under domain shift without requiring labeled target data. The method addresses the infeasibility of estimating density ratios in high-dimensional, unstructured prompt spaces through the following steps:

1. Semantic Embedding: Instead of operating in the raw token space, the framework uses a pre-trained embedding model $h: \mathcal{X} \to \mathcal{Z}$ to project prompts into a lower-dimensional semantic space. This captures cross-domain similarity, allowing the domain shift to be characterized by the shift in the embedding distribution.
2. Density Ratio Estimation: In the embedding space, a domain classifier is trained to distinguish between source (calibration) and target (test) prompts. This classifier estimates the density ratio $\tilde{r}(z) = \frac{dP'_Z}{dP_Z}(z)$, which serves as a weight for calibration samples based on their proximity to the target domain.
3. Regularized Weighting: To prevent the prediction sets from degenerating into the entire output space (a risk when density ratios are highly unbalanced due to large shifts or small sample sizes), the authors introduce a regularization step. The weight assigned to the test point is replaced by a regularized parameter $\lambda$ that depends only on the calibration data.
   • The empirical score distribution is constructed as:
     $$\hat{\mu}_n = \sum_{i=1}^n \frac{\hat{w}_i}{\sum \hat{w}_j + \lambda} \delta_{S_i} + \frac{\lambda}{\sum \hat{w}_j + \lambda} \delta_{\infty}$$
     where $\hat{w}_i$ are the estimated density ratios for calibration samples, and $\lambda$ is chosen (typically as $\max_i \hat{w}_i$) to balance validity and adaptivity.
4. Prediction Set Construction: The final prediction set for a new prompt is defined by the $(1-\alpha)$-quantile of this weighted empirical distribution.

Theoretical Guarantees
The paper establishes finite-sample coverage bounds for DS-CP:

• Lower Bound: The coverage probability is guaranteed to be at least $1 - \alpha$ minus a gap term. This gap depends on the normalized weights and the total variation distance between the score distributions of the calibration and test points in the embedding space.
• Smoothness Condition: The authors show that if the nonconformity score varies smoothly with respect to the semantic embedding, the coverage gap remains small even if the raw prompt distributions are distinct.
• Conservativeness: An upper bound is provided, showing that the regularization term $\lambda$ introduces a controlled degree of conservativeness (over-coverage) when the effective sample size of relevant calibration data is low.
• Recovery of Standard CP: Under the exchangeability condition (no domain shift), the method reduces exactly to standard CP, ensuring no loss of validity in the ideal setting.

Experimental Results
The framework was evaluated on the MMLU benchmark, treating different subject categories as distinct domains to simulate domain shift. Experiments were conducted across 16 open-source LLMs from 9 model families.

• Coverage: Standard CP frequently failed to achieve the target 90% coverage, often dropping significantly below this level when calibration and test domains diverged. DS-CP consistently improved coverage, maintaining it closer to the target level across all models and domain pairs.
• Set Size: While DS-CP produced slightly larger prediction sets than standard CP to compensate for the shift, the increase was modest. This demonstrates a successful trade-off between restoring validity and maintaining efficiency.
• Adaptivity: Analysis showed that DS-CP selectively corrects under-coverage in difficult cases (where standard CP fails) while making minimal adjustments when standard CP is already valid, proving it is not merely a uniform set enlargement.

Significance and Claims
The paper claims to provide a practical step toward trustworthy uncertainty quantification for LLMs in real-world deployment. Its primary significance lies in:

1. Extending CP beyond exchangeability: It offers a systematic, theoretically grounded method to handle domain shift in LLMs, a scenario where standard CP breaks down.
2. Semantic Adaptation: By leveraging semantic embeddings rather than raw text, it makes density ratio estimation feasible for high-dimensional LLM inputs.
3. Balanced Reliability: It achieves a balance between statistical validity (preserving coverage guarantees) and adaptivity (tailoring prediction sets to the specific target domain), addressing the critical issue of unreliable predictions in safety-critical applications.

The authors note limitations, including dependence on the quality of the embedding and density ratio estimators, the requirement for unlabeled target prompts, and the current focus on multiple-choice tasks rather than open-ended generation. However, they posit that DS-CP represents a principled adaptation of conformal prediction specifically tailored for the high-dimensional, unstructured nature of LLM prompts under distribution shift.