TRAQ : Trustworthy Retrieval Augmented Question Answering via Conformal Prediction (2024)

Abstract: When applied to open-domain question answering, large language models (LLMs) frequently generate incorrect responses based on made-up facts, which are called hallucinations . Retrieval augmented generation (RAG) is a promising strategy to avoid hallucinations, but it does not provide guarantees on its correctness. To address this challenge, we propose the Trustworthy Retrieval Augmented Question Answering, or *TRAQ*, which provides the first end-to-end statistical correctness guarantee for RAG. TRAQ uses conformal prediction, a statistical technique for constructing prediction sets that are guaranteed to contain the semantically correct response with high probability. Additionally, TRAQ leverages Bayesian optimization to minimize the size of the constructed sets. In an extensive experimental evaluation, we demonstrate that TRAQ provides the desired correctness guarantee while reducing prediction set size by 16.2% on average compared to an ablation. The implementation is available: [https://github.com/shuoli90/TRAQ](https://github.com/shuoli90/TRAQ).

## introduction large language models (llms) have achieved state-of-the-art (sota) results on many question answering (qa) tasks @xcite . however, in open-domain qa tasks where candidate answers are not provided, llms have also been shown to confidently generate incorrect responses, called hallucinations @xcite . hallucinations have already led to real-world consequences when end users rely on the correctness of the generated text. as a consequence, there is an urgent need for techniques to reduce hallucinations. we propose a novel framework, trustworthy retrieval augmented question answering (traq), summarized in figure 1 , that combines retrieval augmented generation (rag) @xcite with conformal prediction @xcite figure 1 : comparison of the standard rag pipeline with traq on a practical illustration reveals a significant difference. with the standard retrieval augmented generation (rag) approach, there is a possibility that the retrieved passage may lack relevance in addressing the given question. on the contrary, traq leverages conformal prediction to ensure that the retrieved set includes the relevant passage with a high probability and that the llm set contains a semantically correct answer with a high probability. through the aggregation of these prediction sets, traq provides a guarantee that a semantically correct answer is contained in its set of answers with a high probability. @xcite @xcite to provide theoretical guarantees on question answering performance. rag reduces hallucinations by retrieving passages from a knowledge base such as wikipedia and then using an llm to answer the question. if the retrieved passages are relevant to the question, the llm can use this information to generate correct answers. however, rag can fail for two reasons: either the retrieved passage is not relevant to the question, or the llm generates the incorrect answer despite being given a relevant passage. to avoid these issues, traq uses conformal prediction, an uncertainty quantification technique that modifies the underlying model to predict sets of outputs rather than a single output. these prediction sets are guaranteed to contain the true output at a user-specified rate, e.g., at least 90% of the time. in particular, traq applies conformal prediction separately to the retrieval model (to obtain sets of retrieved passages guaranteed to contain the relevant passage with high probability) and the generator (to obtain sets of answers that contain the true answer with high probability, assuming the relevant passage is given). then, traq aggregates the two sets for the rag task, as demonstrated in figure 2a . by a union bound, retriever sets contain relevant passages, and generator prediction sets contain true answers with high probability, establishing that the aggregated set by traq contains the ground truth answer with high probability. a major challenge to this basic pipeline is that there may be many different ways of expressing the correct answer in natural language. for example, the responses deep learning is a subset of machine learning and machine learning is a superset of deep learning are different ways of expressing the same meaning @xcite . this diversity of possible responses also makes prediction probabilities less reliable since if an answer can be expressed in many different but equivalent ways, then the probabilities may be divided across these different responses, making them all smaller even if the model is confident it knows the correct answer. traq addresses this challenge by modifying the notion of ground-truth coverage in conformal prediction to focus on semantic notions of uncertainty. in particular, traq aggregates semantically equivalent answers across a large number of samples from the llm and uses the number of clusters of non-equivalent answers as a measure of uncertainty. this measure is used as a nonconformity measure to construct prediction sets. finally, the prediction sets are over clusters of equivalent answers rather than individual answers. this strategy also enables traq to work on black-box apis such as gpt-3.5-turbo, where the predicted probabilities for individual tokens are not available. a second challenge is that the prediction sets can become very large since we are aggregating uncertainty across multiple components. this complexity introduces hyperparameters into traq; while traq guarantees correctness regardless of the choice of these hyperparameters, they can affect the performance of traq in terms of the aver-age prediction set size. to address this challenge, traq uses bayesian optimization to minimize the average size of the prediction sets it generates. we evaluate traq in conjunction with several generative llms, including both gpt-3.5-turbo-0613 @xcite and llama-2-7b @xcite ; and on four datasets, including a biomedical question answering dataset. our experiments demonstrate that traq empirically satisfies the coverage guarantee (i.e. the prediction sets outputs contain semantically correct answers with the desired probability), while reducing the average prediction set size compared to an ablation by 16.2%. thus, traq is an effective strategy for avoiding hallucinations in applications of llms to open domain question answering. contributions. we offer the first conformal prediction guarantees for retrieval augmented generation (rag) targeted question answering. our framework, traq, introduces a novel nonconformity measure that estimates the uncertainty for each semantically distinct meaning and obtains a coverage guarantee at the semantic level. furthermore, traq leverages bayesian optimization to minimize the average size of the generated prediction sets. finally, our experiments demonstrate that traq is effective at avoiding hallucinations in open-domain question answering. ## background retrieval for open-domain qa. a two-stage approach is often used for open-domain question answering (qa): first, a retriever is used to obtain informative passages; and second, a generator produces answers based on the retrieved passages. a popular choice for the retriever is the dense passage retriever (dpr) @xcite , which measures similarity by taking the inner product of the bert @xcite embeddings of the question and passage @xcite . other works @xcite @xcite have improved the performance of dpr and extended it to more diverse settings. retrieval augmented generation (rag) @xcite proposes to jointly fine-tune the retriever and the generator for qa tasks. ## the traq framework traq is composed of two steps. the first is the prediction set construction step, where a question q is used to create a retrieval set c ret (q) for the retriever and a llm set c llm (q, p) for each pair (question q, passage p). these sets are aggregated into an aggregation set c agg (q). the second step is the performance improvement step, where promising error budgets α ret and α llm are sampled from a bayesian model. using these budgets, the prediction sets are constructed on the optimization set and evaluated for their performance. this process is repeated t times, and the final output is the error budgets α ret and α llm with the highest performance. the chosen hyperparameters are used to construct prediction sets as in the first step using a separate held-out calibration set. the traq framework is summarized in figure 2 . ## conclusion we propose an algorithm, called trustworthy retrieval augmented question answering (traq), which applies conformal prediction to construct prediction sets for retrieval augmented generation (rag). traq first constructs prediction sets for the retriever and generator and then aggregates these sets. traq guarantees that for each question, a semantically correct answer is included in the prediction set it outputs with high probability. to the best of our knowledge, this guarantee is the first conformal guarantee for retrieval augmented generation. additionally, to minimize prediction set size, traq leverages bayesian optimization to identify optimal hyperparameters. in our comprehensive experiments, we demonstrate that traq provides an overall semantic level coverage guarantee across different tasks, and that bayesian optimization effectively reduces prediction set size.