Domain-Shift-Aware Conformal Prediction for Large Language Models
Zhexiao Lin, Yuanyuan Li, Neeraj Sarna, Yuanyuan Gao, Michael von Gablenz
TL;DR
This work tackles uncertainty quantification for large language models under domain shift by extending conformal prediction with domain-shift awareness. DS-CP embeds prompts into a semantic space, estimates a density-ratio in that space, and uses a regularized weighting scheme to reweight calibration scores, preserving coverage while improving adaptivity. The authors provide finite-sample guarantees and demonstrate, on the MMLU benchmark, that DS-CP consistently improves coverage over standard CP under substantial domain shifts with only modestly larger prediction sets. The approach offers a practical, theoretically grounded path toward trustworthy uncertainty quantification for LLMs in real-world deployment.
Abstract
Large language models have achieved impressive performance across diverse tasks. However, their tendency to produce overconfident and factually incorrect outputs, known as hallucinations, poses risks in real world applications. Conformal prediction provides finite-sample, distribution-free coverage guarantees, but standard conformal prediction breaks down under domain shift, often leading to under-coverage and unreliable prediction sets. We propose a new framework called Domain-Shift-Aware Conformal Prediction (DS-CP). Our framework adapts conformal prediction to large language models under domain shift, by systematically reweighting calibration samples based on their proximity to the test prompt, thereby preserving validity while enhancing adaptivity. Our theoretical analysis and experiments on the MMLU benchmark demonstrate that the proposed method delivers more reliable coverage than standard conformal prediction, especially under substantial distribution shifts, while maintaining efficiency. This provides a practical step toward trustworthy uncertainty quantification for large language models in real-world deployment.