Domain-Shift-Aware Conformal Prediction for Large Language Models
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.
Lay Summary
AI chatbots can answer many kinds of questions, but they can also sound confident when they are wrong. This is especially risky when people use them in new situations that differ from the examples used to check them, such as a different subject area or a different group of users. This paper introduces a way to help an AI system say, in effect, “here are the answers I think are plausible” rather than giving only one answer with misplaced confidence. The method compares new questions with previously checked questions and gives more influence to past examples that are most similar to the new setting. In tests on a broad question-answering benchmark, this approach made the system’s answer sets more reliable than a standard method, especially when the new questions came from different topics. The goal is to make AI systems safer to use when conditions change.