C-RAG: Certified Generation Risks for Retrieval-Augmented Language Models

TL;DR

C-RAG introduces a principled framework to certify generation risks for retrieval-augmented language models using conformal risk control. It formalizes a constrained RAG generation protocol and derives two main guarantees: a per-configuration conformal risk upper bound and a valid-configuration set ensuring risk below a target level, with extensions to test-time distribution shifts. The authors prove theoretically that RAG can provably reduce conformal generation risk relative to a vanilla LLM, with guarantees enhanced by higher retrieval quality and larger external knowledge bases, and they validate these results empirically across four NLP datasets and multiple retrievers. The work provides a practical pathway for risk-controlled deployment of RAG systems, including strategies to select configurations that meet predefined risk targets and to understand robustness under distribution shifts.

Abstract

Despite the impressive capabilities of large language models (LLMs) across diverse applications, they still suffer from trustworthiness issues, such as hallucinations and misalignments. Retrieval-augmented language models (RAG) have been proposed to enhance the credibility of generations by grounding external knowledge, but the theoretical understandings of their generation risks remains unexplored. In this paper, we answer: 1) whether RAG can indeed lead to low generation risks, 2) how to provide provable guarantees on the generation risks of RAG and vanilla LLMs, and 3) what sufficient conditions enable RAG models to reduce generation risks. We propose C-RAG, the first framework to certify generation risks for RAG models. Specifically, we provide conformal risk analysis for RAG models and certify an upper confidence bound of generation risks, which we refer to as conformal generation risk. We also provide theoretical guarantees on conformal generation risks for general bounded risk functions under test distribution shifts. We prove that RAG achieves a lower conformal generation risk than that of a single LLM when the quality of the retrieval model and transformer is non-trivial. Our intensive empirical results demonstrate the soundness and tightness of our conformal generation risk guarantees across four widely-used NLP datasets on four state-of-the-art retrieval models.

Figures (18)

• Figure 1: Overview of C-RAG. In the estimation stage (upper row), the conformal risk controller computes conformal generation risks for different RAG configurations (\ref{['risk_guarantee_1']}), and valid configuration sets for different risk levels (\ref{['risk_guarantee_2']}), both based on risk statistics on the calibration set. In the inference stage (lower row), for any configuration $\hbox{${\bm{\lambda}}$}$ and any desired risk level $\hbox{$\alpha$}$ provided by users, the conformal risk controller outputs the conformal generation risk $\hbox{$\hat{\alpha}_{\bm{\lambda}}$}$ with Risk Guarantee (1) and the configuration set $\hbox{$\hat{\Lambda}_\alpha$}$ with Risk Guarantee (2).
• Figure 2: Certification framework of C-RAG. We provide theoretical results with the data exchangeability assumption in \ref{['sec:rag_conf']} (upper row) and extend the results to more complex scenarios under test-time distribution shifts in \ref{['sec:distribution_shift']} (lower row).
• Figure 3: Conformal generation risk $\hbox{$\hat{\alpha}_{\text{rag}}$}$ and empirical risk based on retrieval model OpenAI/ada taking different $\hbox{$N_{\text{rag}}$}$ ($\hbox{$\lambda_g=1,\lambda_s=1.0$}$). We observe that our conformal generation risk (\ref{['risk_guarantee_1']}) is valid and tight; larger $\hbox{$N_{\text{rag}}$}$ reduces risk $\hbox{$\hat{\alpha}_{\text{rag}}$}$ (empirically validating \ref{['thm:gene_rag']}).
• Figure 4: Conformal generation risk $\hbox{$\hat{\alpha}_{\text{rag}}$}$ with different $\hbox{$N_{\text{rag}}$}$ using different retrieval models ($\hbox{$\lambda_g=1,\lambda_s=1.0$}$). We observe that large $\hbox{$N_{\text{rag}}$}$ effectively reduces $\hbox{$\hat{\alpha}_{\text{rag}}$}$ for different models; the trained Biencoder-SFT usually leads to the lowest conformal generation risk.
• Figure 5: Conformal generation risk $\hbox{$\hat{\alpha}_{\text{rag}}(\rho)$}$ and empirical risks based on retrieval model OpenAI/ada under distribution shifts ($\hbox{$N_{\text{rag}}=15, \lambda_g=1, \lambda_s=1.0$}$). We observe that our distribution-drift conformal generation risk (\ref{['pro:conf_shf']}) is empirically valid and tight.

Theorems & Definitions (19)

• Proposition 1: Risk Guarantee (1), adaptation of bates2021distribution to constrained RAG generation
• Proposition 2: Risk Guarantee (2), adaptation of angelopoulos2021learn to constrained RAG generation
• Definition 1: $\hbox{$V_{\text{rag}}$}$-retrieval model
• Proposition 3: lower bound of the number of retrieved positive examples
• Definition 2: $\hbox{$(d^+,\Phi_M)$}$-transformer
• Theorem 1: RAG reduces the conformal generation risk
• Theorem 2: Conformal generation risk under distribution shifts
• Theorem 3: RAG reduces conformal generation risk even under distribution shifts
• proof : Proof of \ref{['risk_guarantee_1']}
• proof : Proof of \ref{['risk_guarantee_2']}