PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Mintong Kang, Nezihe Merve Gürel, Ning Yu, Dawn Song, Bo Li

arXiv:2402.03181v5cs.AIcs.CLcs.IR

TL;DR

RAG is used to ground language-model generations, but its generation risks lack theoretical understanding. C-RAG uses constrained generation and conformal analysis to certify those risks, and the paper reports sound, tight guarantees and lower conformal risk for RAG than for a single LLM.

  • Problem

    RAG is proposed to improve generation credibility, but its generation risks remain theoretically unexplored, leaving open whether risks can be guaranteed and reduced.

    Source: §1 · p. 1

  • Method

    C-RAG applies conformal analysis to a constrained RAG generation protocol to provide risk bounds and identify configurations meeting a desired risk level.

    Source: §1 · p. 1, §4.1 · p. 3

  • Results

    Across four datasets and retrieval methods, guarantees are sound and tight with gaps narrowing to 1e-3, and RAG consistently has lower conformal risk than a single LLM.

    Source: §1 · p. 2

  • Takeaways & Limitations

    C-RAG provides both a risk bound for a chosen RAG configuration and a set of configurations that can meet a desired risk level.

    Source: §8 · p. 9

    The theoretical analysis uses a one-layer self-attention transformer and treats each instance approximately as a single token.

    Source: §5.1 · p. 4

Abstract

from arXiv · show

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.

1. Introduction

The paper addresses the gap in theoretical understanding of RAG generation risks and asks whether those risks can be certified and reduced. It presents C-RAG and reports empirical support for sound, tight guarantees and lower conformal risk for RAG than for a single LLM.

Motivation

  • LLM generations can be unreliable, untrustworthy, and risky, motivating certifiable risk control before deployment, especially in safety-critical domains.The passage attributes these concerns to prior work.Source: §1 · p. 1
  • RAG's generation risks remain theoretically unexplored despite its proposed use of retrieved knowledge to enhance generation credibility.The paper asks whether RAG can yield low generation risks, how to guarantee risks for RAG and vanilla LLMs, and what conditions enable reductions.Source: §1 · p. 1

Contributions

  • C-RAG proposes conformal risk analysis for RAG, including a high-probability upper bound on generation risk and a way to identify configurations that control risk below a desired level.The framework also introduces a constrained generation protocol for RAG.Source: §1 · p. 1

Theoretical contributions

  • C-RAG extends conformal generation-risk analysis to general bounded risk functions under test-time distribution shifts and proves RAG has lower conformal risk than a vanilla LLM with or without shifts.The shift guarantee is presented as the first for general bounded risk functions.Source: §1 · p. 2

Empirical results

  • Across four datasets and retrieval methods, C-RAG's risk guarantees are sound and tight, with gaps narrowing to 1e-3, while RAG consistently has lower conformal risk than a single LLM.Text-embedding-ada-002 and supervised fine-tuned embedding models achieve lower conformal risks than the other evaluated retrieval baselines.Source: §1 · p. 2

2. Related work

Prior theoretical analysis of RAG and conformal risk control leaves important scope gaps. This work analyzes RAG for self-attention transformers and extends conformal guarantees under distribution shifts to general bounded risk functions.

RAG analysis

  • Prior RAG analysis considered a constrained function and data class, while this work analyzes RAG generation risks in self-attention transformers for arbitrary data distributions.The paper describes this as the first such theoretical analysis.Source: §2 · p. 2

Conformal risk control

  • Conformal risk control provides high-confidence guarantees under exchangeability, but those guarantees can break under test-time shifts.Existing shift guarantees cited here are limited to monotonic risk functions, an assumption the paper says is not always practical.Source: §2 · p. 2
  • The paper introduces conformal risk guarantees for general bounded risk functions under test-time distribution shifts.This broadens the stated scope beyond shift analyses limited to monotonic risks.Source: §2 · p. 2

3. Preliminaries

The preliminaries formulate conformal risk control as testing whether each parameter configuration exceeds a desired risk level. Valid p-values and family-wise error control provide the basis for selecting configurations.

Conformal risk control

  • For a finite parameter space, the setup defines one null hypothesis per parameter: that its population risk exceeds the desired risk level α.The empirical and population risks are denoted separately, and the hypotheses are indexed across N parameters.Source: §3 · p. 2
  • Given n calibration samples, the cited conformal methods provide valid p-values for the null hypotheses.The p-values are used in the subsequent family-wise error control procedure.Source: §3 · p. 2
  • A family-wise error-controlling algorithm returns a set of indices, denoted as the output of the algorithm.The algorithm is defined to include only false null hypotheses with probability at least 1−δ.Source: §3 · p. 2

4. Conformal generation risks of RAG models

C-RAG combines a constrained RAG generation protocol with conformal risk analysis to provide an upper risk bound for a configuration and a set of configurations meeting a desired risk level. The paper also extends the analysis to distribution shifts and general bounded risk functions.

4. Conformal generation risks of RAG models

  • C-RAG provides a high-probability generation-risk upper bound for a given RAG configuration and a configuration set that can maintain risk below a desired level.The conformal analysis uses calibration-sample statistics.Source: §4 · p. 2

4.1. Problem setup

  • The setup represents a generation protocol as producing a set of candidate outputs and defines risk against a reference text, with 1−ROUGE as one example.The size of the generation set can be adjusted as a risk-control parameter.Source: §4.1 · p. 3

4.2. Constrained generation protocol for RAG models

  • The constrained protocol retrieves relevant examples, samples generations, and rejects outputs that exceed a similarity threshold against previously collected generations.Its configuration comprises the number of retrieved examples, the target number of generations, and the similarity threshold.Source: §4.2 · p. 3

4.3. Conformal generation risks for RAG models

  • For a desired risk level α, Risk Guarantee (2) returns configurations whose generation risk is guaranteed below α.The procedure can use family-wise error control, including Bonferroni correction.Source: §4.3 · p. 4
  • For a given configuration, Risk Guarantee (1) supplies its conformal generation risk as an upper bound on generation risk.The overview depicts this output in the inference stage, based on calibration-set risk statistics.Source: Figure 1 · p. 3, §4.3 · p. 4
  • The two guarantees are dual: one bounds risk for a chosen configuration, while the other selects configurations statistically below the target level.The paper relates this connection to the duality between p-values and confidence intervals.Source: §4.3 · p. 4
  • C-RAG extends conformal generation-risk analysis under test-time shifts to general bounded risk functions, beyond earlier work limited to monotonic risks.The paper also reports that RAG has lower conformal risk than vanilla LLMs with or without shifts.Source: §4.3 · p. 4

5. Theoretical analysis of C-RAG

The analysis proves that RAG can achieve lower conformal generation risk than a single LLM without retrieval, with benefits tied to retrieval and transformer quality.

5.1. Analysis setup

  • The analysis treats each instance approximately as one token and defines positive pairs as examples with identical semantic meanings.The authors note that the analysis can extend to broader definitions of positive pairs.Source: §5.1 · p. 4

5.2. Retrieval quality analysis

  • For sufficiently large external knowledge bases, the positive-example lower bound approaches 0.9Nrag, while long-tailed classes require more samples.The analysis links retrieval quality to the number of positive examples retrieved for in-context demonstrations.Source: §5.2 · p. 5

5.3. RAG achieves provably lower conformal generation risk than a single LLM without retrieval

  • Under non-trivial retrieval and transformer conditions, RAG has lower conformal generation risk than a single LLM with high probability.The probability of reduced risk increases with calibration size, retrieval scope, and transformer quality.Source: §5.3 · pp. 5–6, §5.2 · p. 5

6. C-RAG under distribution shifts

C-RAG extends conformal generation-risk guarantees to test-time distribution shifts and establishes that RAG can have lower conformal risk than a single LLM under stated conditions. The analysis connects shift size and retrieval-model quality to the risk guarantee.

6.2. Conformal generation risk under distribution shifts

  • C-RAG provides generation-risk certification for any bounded risk function under any distribution shift, extending beyond prior guarantees limited to monotonic risks and sample-weight shifts.The shifted test distribution is assumed to have bounded Hellinger distance from the calibration distribution.Source: §6.2 · p. 6
  • The shift-adjusted risk bound uses calibration-set statistics, accounts for distribution drift, and recovers the empirical mean as ρ approaches zero and calibration size grows.The risk adjustment includes a term scaled by ρ and a finite-sample error term.Source: §6.2 · pp. 6–7

6.3. RAG achieves provably lower conformal generation risk than a single LLM under distribution shifts

  • Under bounded Hellinger shifts and the theorem's conditions, RAG has lower conformal generation risk than a single LLM.The theorem specifies conditions on retrieval and transformer quality and on the sizes of the external knowledge base, retrieved set, and generation set.Source: §6.3 · p. 7
  • Distribution shift degrades retrieval quality and raises retrieval uncertainty, but the shift's influence on that uncertainty decreases inversely with Nrag and exponentially with Next.The paper also states that lower-variance retrieval models are more robust to distribution drift.Source: §6.3 · p. 7
  • The theorem's external-knowledge-base condition is described as moderately large, a condition the authors say is often met in practice.The paper additionally notes an analysis of retrieved positive examples under distribution shifts in an appendix.Source: §6.3 · p. 7

7. Evaluation

The evaluation uses four NLP datasets and four retrieval models. Across experiments, C-RAG's guarantees are reported as sound and tight, while higher retrieval counts and selected configurations reduce or control conformal generation risk.

7. Evaluation

  • Across four datasets, C-RAG's guarantees are sound and tight, RAG reduces conformal risk, and certified configurations have empirical risks below the desired level.The evaluation overview summarizes results across retrieval models and experimental sections.Source: §7 · p. 7

7.1. Evaluation setup

  • The experiments use AESLC, CommonGen, DART, and E2E, four retrieval models, Llama-2-7b, and 1 − ROUGE-L as the risk function.The knowledge base comprises 30 public datasets from nine categories and over six million documents.Source: §7.1 · p. 7

7.2. Evaluation of conformal generation risks

  • Conformal risks upper-bound sampled empirical risks, can be nearly attained by some test instances, and decrease as Nrag increases.Across retrieval models, RAG also has lower conformal risk than an LLM without retrieval; Biencoder-SFT is generally lowest.Source: §7.2 · p. 8, Figure 3 · p. 8, Figure 4 · p. 8

7.3. Conformal generation risk under distribution shifts

  • Under distribution shifts, conformal risks are sound and tight across models, rise linearly with Hellinger distance ρ, and remain non-trivial up to ρ = 0.2.Biencoder-SFT and OpenAI/ada show lower risks at varying distances than the other compared retrieval models.Source: §7.3 · p. 8, Figure 5 · p. 9

7.4. C-RAG with multi-dimensional RAG configurations

  • Multi-dimensional configurations retain sound, tight guarantees, and increasing Nrag reduces generation risk more effectively than adjusting λg.The evaluation varies both Nrag and λg.Source: §7.4 · pp. 8–9, Figure 6 · p. 9

7.5. Valid configurations given desired risk levels

  • Empirical risks for certified configurations remain below the target α; higher Nrag and larger λg also contribute to reducing conformal risk.C-RAG computes the valid configuration set using a Bonferroni correction for family-wise error control.Source: §7.5 · p. 9, Figure 7 · p. 9

8. Conclusion

C-RAG provides conformal generation-risk guarantees for RAG models, with both theoretical and empirical support for its risk-reduction claims.

8. Conclusion

  • C-RAG certifies a conformal generation risk for a given RAG configuration.Source: §8 · p. 9
  • For a desired risk level, C-RAG certifies a valid set of configurations.Source: §8 · p. 9
  • The paper theoretically shows that RAG reduces the conformal generation risk of a single LLM.Source: §8 · p. 9
  • Empirical results validate the soundness and tightness of the risk guarantees.Source: §8 · p. 9

Impact Statement

The impact statement frames C-RAG as a way to address ethical and societal concerns through generation-risk control.

Impact Statement

  • The authors say C-RAG could safeguard LLM deployment and applications against ethical and societal concerns.Source: Impact Statement · p. 9
  • C-RAG can define bias or ethics as a risk function and control generation risk below a desired level.Source: Impact Statement · p. 9
  • The authors do not expect negative societal consequences from the work.Source: Impact Statement · p. 9

Appendices

The appendices detail C-RAG’s constrained generation and conformal certification procedures, its theoretical conditions, and practical limitations. They also situate the framework among prior retrieval and conformal-risk work.

A. Discussions of limitations and future work

  • C-RAG’s practical use may require costly real-time calibration, still permits risky generations under a high-confidence bound, and trades knowledge-base scale against inference efficiency.The appendix also identifies conformal risk analysis for time-series data as future work.Source: A. Discussions of limitations and… · p. 15

B. Additional related work

  • Prior risk guarantees under distribution shifts can require monotonic risk, whereas C-RAG provides a conformal risk bound for general bounded risks.The related-work discussion also notes that earlier statistical retrieval analysis does not generalize to self-attention transformers.Source: B. Additional related work · pp. 15–16

C.1. Constrained generation protocol for RAG models

  • The constrained protocol retrieves Nrag examples, augments the prompt, and samples a generation set while rejecting outputs too similar to earlier generations.Its configuration comprises the number of retrieved examples, generation-set size, and a similarity threshold.Source: §C.1 · p. 16

D.1. Risk guarantee (1) (Prop. 1)

  • C-RAG calibrates empirical risk to certify an upper bound at a chosen confidence level, using Hoeffding-Bentkus inequalities.The appendix favors its general bound because a tighter alternative assumes a one-dimensional configuration and monotonic risk, conditions absent from many practical protocols.Source: §D.1 · p. 17

D.2. Risk guarantee (2) (Prop. 2)

  • For multiple configurations, C-RAG tests each risk null with finite-sample p-values and applies an FWER-controlling search to identify certified configurations.The graph-based search dynamically assigns error levels and can yield tighter certification than Bonferroni correction.Source: §D.2 · p. 18

E. Grammian generalization bound

  • The Grammian generalization bound supports certificates for generic models and bounded risks under shifts when the distribution distance is limited.The appendix says computing the certificate requires the risk mean and variance under the source distribution.Source: E. Grammian generalization bound · p. 19

F.1. Detailed remark of Def. 1

  • The Vrag-retrieval condition requires retrieval to distinguish positive from negative examples better than random, with bounded similarity-difference variance.The discussion allows contrastive losses beyond the commonly used formulation, including triplet loss.Source: §F.1 · p. 19

F.2. Detailed remark of Def. 2

  • The transformer-quality discussion separates attention to positive pairs from prediction quality, associating better-aligned predictions with smaller margins and a larger integrated margin CDF.The appendix notes that Softmax attention scores are nonnegative.Source: §F.2 · p. 19

F.3. Proof and detailed remark of Prop. 3

The appendix derives how retrieval quality affects the number of positive examples retrieved, then connects those examples and transformer quality to the probability that RAG has lower conformal generation risk.

F.3. Proof and detailed remark of Prop. 3

  • A larger external knowledge base can raise the lower bound on retrieved positive examples, while long-tailed reference outputs require more examples for comparable retrieval quality.The bound also indicates that lower retrieval-model variance can support generalization to the test distribution and better retrieval quality.Source: §F.3 · p. 20

F.4. Proof and detailed remark of Thm. 1

  • The probability that RAG’s conformal generation risk is lower increases with calibration-set size, retrieved-example count, knowledge-base size, and transformer quality.The appendix also connects the risk reduction to more positive retrieved examples and lower retrieval uncertainty.Source: §F.4 · pp. 21–22

G. Cor. 1: Asymtotic result of Thm. 1

The asymptotic corollary strengthens the theorem’s comparison when the external knowledge base is sufficiently large. Its bound approaches a high-probability guarantee as calibration and retrieval-example counts grow.

G. Cor. 1: Asymtotic result of Thm. 1

  • With a sufficiently large external knowledge base, RAG has lower conformal generation risk than the vanilla model with high probability, asymptotically approaching 1 as calibration and retrieved-example counts grow.The bound no longer depends on the knowledge-base or calibration/test label distributions, and is tighter than the theorem’s bound.Source: G. Cor. 1: Asymtotic result of Thm. 1 · pp. 25–26

H. Prop. 4: Retrieval quality analysis under distribution shifts

This section examines retrieval quality and conformal-risk guarantees under bounded distribution shifts, including conditions for preserving retrieval quality and the evaluation setup. It also reports that C-RAG’s conformal risk bounds are sound and tight against sampled empirical risks, with performance varying across retrieval models.

H. Prop. 4: Retrieval quality analysis under distribution shifts

  • Under distribution shifts, retrieval quality decays linearly with shift radius; maintaining Vrag(ρ) < 1 requires lower Vrag as ρ grows, while a sufficiently large knowledge base removes dependence on varying label proportions.The stated decay factor is m(ρ); the knowledge-base size must offset worst-case long-tail distributions.Source: H. Prop. 4: Retrieval quality analysis… · p. 26

I.1. Proof and detailed remark of Thm. 2

  • The cited passage concerns conformal risk on distribution Q, denoted α̂ρ, and says Q is within bounded Hellinger distance ρ of original distribution D; it introduces a 1 − 3δ/4 probability statement.Source: §I.1 · p. 30

I.2. Proof and detailed remark of Thm. 3

  • Distribution shifts lower retrieval quality Vrag(ρ) and raise retrieval uncertainty pr(ρ), whose shift impact decreases inversely with Nrag and exponentially with Next.The theorem’s remarks also state that sufficiently large external knowledge bases remove dependence on varying label portions r(c), and that low-variance retrieval models are more robust to shifts.Source: §I.2 · p. 30

J.1. Experiment setup

  • The experiments cover AESLC, CommonGen, DART, and E2E, and evaluate four retrieval-model types, including BM25 and BAAI/bge.The datasets span email-subject generation, commonsense descriptions, structured-record-to-text generation, and restaurant meaning-representation generation.Source: §J.1 · pp. 32–33
  • The shift protocol samples test sets from a pool at the original test-set size and samples a weight vector for each set.The protocol reports empirical risk on each shifted test set and its Hellinger distance from the original distribution.Source: §J.1 · p. 34

J.2. Conformal risk bounds evaluation

  • C-RAG’s certified conformal risks exceed sampled empirical risks, with some empirical risks approaching the bounds; bounds decrease as more examples are retrieved.The evaluation samples multiple test sets from a pool and compares their empirical risks with bounds computed using validation-set calibration statistics.Source: §J.2 · p. 34
  • Biencoder-SFT generally has the lowest conformal generation risks, while OpenAI/ada outperforms it on CommonGen.The paper attributes Biencoder-SFT’s general performance to training in the same domain as the test sets.Source: §J.2 · p. 35

J.3. Conformal risk bounds evaluations under distribution shifts

Across distribution shifts and multi-configuration settings, C-RAG’s bounds are reported as sound and tight. The evaluations also show lower risks with more retrieved examples, and a qualitative example illustrates how diversity constraints affect generations.

J.3. Conformal risk bounds evaluations under distribution shifts

  • Under shifts, bounds are sound and tight, rise linearly with Hellinger distance, and remain non-trivial through 0.2; Biencoder-SFT and OpenAI/ada start lower and share the slope.The bounds for these two retrieval models are lower because they have lower initial risk without distribution shifts.Source: §J.3 · pp. 35–36, Figure 14 · p. 36

J.5. Valid configurations identification with specified risk levels

  • For configurations certified below target α, sampled empirical risks remain below α; valid configurations include regions with large Nrag and λg.The passage reports these results on DART and ECE; Figure 15 varies λg and Nrag.Source: Figure 15 · p. 36, §J.5 · p. 37

J.4. Conformal risk bounds with multi-dimensional RAG configurations

  • Multi-dimensional bounds remain sound and tight; increasing Nrag lowers risk more sensitively than increasing λg, while a larger generation set also helps.The passage reports evaluations on DART and E2E and describes the generation-set comparison in Fig. 17.Source: §J.4 · p. 36, Figure 17 · p. 37

J.6. Conformal generation risks with different inference models

  • Conformal risk decreases as Nrag increases across inference models, and Llama-2-13b has lower risk than Llama-2-7b.Table 1 compares these risks on AESLC using OpenAI text-embedding-ada-002 for retrieval.Source: Table 1 · p. 37, §J.6 · p. 38

J.7. Qualitative example

  • In the World Cup prompt, one retrieved example removes the COVID explanation but not the date error; λs=0.5 diversity constraints yield correct identification, unlike λg=3.The example contrasts retrieval alone, a larger generation set, and diversity constraints.Source: §J.7 · p. 38, Table 2 · p. 38
Loading 2402.03181v5…