PaperlayerView exact revision on arXiv

Source-linked AI summary

Certifiably Robust RAG against Retrieval Corruption

Chong Xiang, Tong Wu, Zexuan Zhong, David Wagner, Danqi Chen, Prateek Mittal

arXiv:2405.15556v2cs.LGcs.CLcs.CR

TL;DR

Retrieval corruption can cause inaccurate RAG responses when malicious passages enter retrieved context. RobustRAG isolates passages into disjoint groups and securely aggregates their LLM responses, achieving certifiable robustness across evaluated tasks, datasets, and models while maintaining benign performance.

  • Problem

    RAG is vulnerable to inaccurate responses when malicious actors compromise retrieved passages.

  • Method

    RobustRAG uses isolate-then-aggregate processing with secure keyword-based and decoding-based aggregation for unstructured text responses.

  • Results

    RobustRAG demonstrates certifiable and empirical robustness across three datasets and three LLMs while maintaining high clean performance.

  • Takeaways & Limitations

    The framework enables worst-case robustness analysis for RAG text generation against bounded adaptive passage corruption.

  • Takeaways & Limitations

    RobustRAG operates in the generation phase and is most effective when the number of corrupted passages is small.

Abstract

from arXiv · show

Retrieval-augmented generation (RAG) is susceptible to retrieval corruption attacks, where malicious passages injected into retrieval results can lead to inaccurate model responses. We propose RobustRAG, the first defense framework with certifiable robustness against retrieval corruption attacks. The key insight of RobustRAG is an isolate-then-aggregate strategy: we isolate passages into disjoint groups, generate LLM responses based on the concatenated passages from each isolated group, and then securely aggregate these responses for a robust output. To instantiate RobustRAG, we design keyword-based and decoding-based algorithms for securely aggregating unstructured text responses. Notably, RobustRAG achieves certifiable robustness: for certain queries in our evaluation datasets, we can formally certify non-trivial lower bounds on response quality -- even against an adaptive attacker with full knowledge of the defense and the ability to arbitrarily inject a bounded number of malicious passages. We evaluate RobustRAG on the tasks of open-domain question-answering and free-form long text generation and demonstrate its effectiveness across three datasets and three LLMs.

I. INTRODUCTION

RAG improves access to external knowledge but remains vulnerable when retrieved passages are corrupted. RobustRAG addresses this problem by isolating passages, securely aggregating responses, and certifying robustness across tasks and models.

  • Problem: Retrieval corruption can make RAG fragile when attackers compromise even a small fraction of retrieved passages.Examples include poisoned knowledge bases and indirect prompt injections that induce incorrect answers or override instructions.
  • Approach: RobustRAG isolates retrieved passages into disjoint groups, generates one response per group, and securely aggregates them.Isolation limits a small number of corrupted passages to a small fraction of isolated responses, supporting robust aggregation.
  • Approach: Secure keyword and decoding aggregation methods address the challenge of combining unstructured responses while limiting malicious influence.The keyword method counts each unique keyword once per isolated response and uses high-count keywords to prompt a final response.
  • Robustness: RobustRAG provides certifiable robustness for unstructured text generation against adaptive attackers with bounded passage corruption.The framework formally certifies non-trivial lower bounds on response quality for certain evaluation queries.
  • Evaluation: 71% clean accuracy and 38% certifiable robust accuracy are reported on RealtimeQA, compared with 69% clean accuracy and no certifiable accuracy for vanilla RAG.The evaluation spans three datasets and three LLMs across open-domain question answering and free-form long text generation.

C. Robustness Formulation

RobustRAG formulates robustness as a worst-case lower bound on response quality under arbitrary retrieval corruption, then develops certification to compute that bound despite adaptive attackers. The formulation is meaningful only when useful benign passages outnumber malicious ones, and evaluation emphasizes small attack budgets.

  • Robustness objective: Robustness targets accurate generation when k′ retrieved passages may be maliciously corrupted and their ranking positions arbitrarily manipulated.
  • Robustness definition: τ-robustness is the lower bound on LLM response quality against every attacker injecting k′ arbitrary-content passages with full defense knowledge.
  • Robustness limits: τ = 0 is trivial, while non-zero robustness can be fundamentally impossible when the number of corrupted passages is extremely large.
  • Feasibility: Robust generation is tractable only when useful benign passages outnumber malicious ones, because arbitrary malicious content can defeat even human responders otherwise.
  • Evaluation focus: With k = 10 retrieved passages, evaluation primarily studies the practical small-budget setting k′ = 1, while also considering larger budgets.
  • Robustness certification: Certification computes correct τ values with provable guarantees rather than trusting performance under a potentially suboptimal empirical attack.
  • Certification challenges: Certification is difficult because arbitrary injected content creates infinitely many corrupted retrievals and responses, so RobustRAG limits attacker influence to make analysis tractable.
  • Certification interpretation: For binary question answering, certifiable robust accuracy averages per-query binary τ values and provides a guaranteed accuracy lower bound under corruption.

III. ROBUSTRAG: A DEFENSE FRAMEWORK

RobustRAG isolates retrieved passages into disjoint groups, generates separate LLM responses, and securely aggregates them to limit corrupted passages’ influence. It instantiates this design with keyword aggregation for unstructured responses, while group size trades off response quality against certifiable robustness.

  • Overview: RobustRAG isolates passages into disjoint groups, generates one response per group, and securely aggregates those responses for the final output.Isolation limits corrupted passages to the groups containing them, leaving other benign groups available for robust aggregation.
  • Passage Isolation: The number of corrupted groups satisfies m′ ≤ k′, so RobustRAG relies on the remaining m − m′ benign groups.Each passage appears in only one group, making the number of affected groups no larger than the number of corrupted passages.
  • Passage Isolation: Larger group size ω can improve isolated-response quality but reduces the number of groups and can make certifiable robustness impossible.With ω = k, RobustRAG becomes vanilla RAG with one group, so even one corrupted passage can manipulate outputs.
  • Secure Keyword Aggregation: Keyword aggregation addresses free-form text because responses such as “Mount Everest” and “The highest mountain is named Everest” need not match exactly.The method distills information across differently worded responses before generating the final answer.
  • Secure Keyword Aggregation: Secure keyword aggregation extracts keywords from each isolated response, counts each unique keyword once per response, and prompts the LLM using high-frequency keywords.The final prompt excludes retrieved passages and uses keywords whose counts meet a threshold.

C. Secure Decoding Aggregation

Secure decoding aggregation operates during generation by combining next-token probability vectors from isolated passage groups. It is designed for settings where keyword extraction may lose information, with decoding choices controlled by probability gaps and no-retrieval predictions.

  • Decoding Aggregation: Secure decoding aggregation combines next-token probability vectors from isolated passage groups at each decoding step and selects tokens from the aggregated vector.Repeating this process produces the full response.
  • Motivation: Decoding aggregation addresses the information loss of keyword extraction on longer responses by operating during the decoding phase.For long-form generation, no-retrieval tokens can help because sentence coherence supplies information beyond retrieved passages.
  • Inference: The inference procedure filters isolated groups unlikely to produce “I don’t know,” sums their probability vectors, and compares the top two token probabilities using η.When the gap is insufficient, it uses a no-retrieval token instead of the top-ranked aggregated token.
  • Certification: The certification procedure analyzes all possible responses induced by bounded passage corruption and returns the lowest evaluation score among them as τ.This exhaustive minimum is the certified robustness value when the response set is tractable.
  • Certification: The keyword certification theorem states that Algorithm 3 correctly returns τ for every response produced by the keyword inference procedure under the specified settings.The guarantee is expressed as M(r, a) ≥ τ for all responses in the possible-response set.

C. Certifying Decoding Aggregation

Decoding certification bounds RobustRAG’s worst-case response quality by enumerating possible token sequences under corruption. Probability-gap cases determine whether the next token is fixed, branches between two choices, or causes certification to return zero.

  • Probability-Gap Cases: If A − B > η + k′, certification fixes the next token to the aggregated top-1 token; if η − k′ ≥ A − B > 0, it fixes the no-retrieval token.Here A and B are the top-two aggregated probability sums.
  • Probability-Gap Cases: If η + k′ ≥ A − B > |η − k′|, the next token may be either the top-1 token or the no-retrieval token.The certification tree branches on these two attacker-relevant possibilities.
  • Certification Outcome: In all other probability-gap cases, the response set becomes intractable and certification aborts with τ = 0.If no such case occurs, the algorithm returns the worst evaluation score over the tractable response set.
  • Implementation Details: When the possible response set becomes very large, some experiments approximate τ by evaluating a random subset of 100 responses because LLM-as-a-judge evaluation is costly.These approximated robustness values are marked with ‡.
  • Certification Procedure: Decoding certification explores a finite response set by treating each possible next-token choice as a branch in a binary tree.Branches terminate at EOS or the maximum token length, after which τ is the lowest evaluation score across completed responses.

D. Remark: Certifiable Evaluation

The certification algorithms evaluate the robustness of inference algorithms by computing τ-based lower bounds over possible responses. Certification can use any evaluation function, including LLM-as-a-judge or human evaluation.

  • Certification algorithms analyze the response set R to compute τ-certifiable robustness for a query q and reference answer a.
  • The evaluation averages per-query τ values across query-answer pairs as a robustness metric.
  • Certification algorithms require benign passages and the reference answer, and can be computationally expensive because they reason over all possible responses.
  • Robustness certification is compatible with any evaluation function M, including alternative or improved evaluators such as human evaluation.

V. EVALUATION

The evaluation measures RobustRAG across multiple QA and generation datasets, models, aggregation settings, attacks, and efficiency dimensions. It finds substantial certified and empirical robustness, generally strong benign performance, parameter trade-offs, and moderate runtime overhead.

  • Experimental setup: Four datasets cover multiple-choice QA, short-answer QA, and free-form biography generation, with 100 sampled queries per dataset except 50 for Bio.
  • Certifiable robustness: RobustRAG achieves 69.0–71.0% certifiable robust accuracy on RQA-MC, 31.0–38.0% on RQA, 26.0–36.0% on NQ, and 38.8–51.2 certifiable LLM-judge score on Bio.
  • Certifiable robustness: 71.0% certifiable accuracy means the response is always correct for 71.0% of queries against an adaptive attacker corrupting one retrieved passage.
  • Benign performance: RobustRAG maintains high benign performance, with a 7% RQA drop reduced to 0% at group size ω = 3 and long-form drops as small as 1.2%.
  • Empirical attacks: Attack success rates are ≤15% for RobustRAG across prompt-injection and poisoning cases, whereas prompt injection can exceed 90% for vanilla RAG.
  • Parameter analysis: Increasing group size generally improves benign performance while reducing certifiable robustness, making ω a trade-off parameter.
  • Parameter analysis: Larger corruption sizes reduce certifiable robustness, reaching zero when 5 of 10 passages are corrupted because benign passages cannot outnumber malicious ones.
  • Runtime: RobustRAG is 1.16–3.65× slower than vanilla RAG because it makes k isolated calls, each processing only one passage.

VI. DISCUSSION

The discussion identifies scope and efficiency boundaries for RobustRAG. Future work targets stronger retrieval-phase attacks, computation overhead, benign performance, and multi-hop questions.

  • Limitations: RobustRAG operates in the generation phase and is effective when the number of corrupted passages k′ is small.
  • Future work: Defending against attackers corrupting many passages requires extending protection to the retrieval phase before robust generation.
  • Future work: RobustRAG has small-to-moderate computation overhead, motivating efficiency improvements such as removing repeated ICL exemplars or exploiting shared prefixes.
  • Future work: Benign performance drops are minimal for QA but small-to-moderate for long-form generation, especially summarization.
  • Future work: The paper focuses on single-hop RAG tasks; extending RobustRAG to multi-hop questions remains an open direction.
  • Conclusion: The paper concludes that isolate-then-aggregate enables worst-case analysis and strong certifiable robustness with negligible-to-small benign-performance impact.

A. Proof of lemmas

The lemmas characterize how bounded malicious probability contributions affect decoding outcomes. By comparing top-token probability gaps with the threshold η and corruption bound m′, the analysis determines when the algorithm selects the top token, a no-retrieval token, or excludes alternatives.

  • Lemma 1: If A − B > η + m′, the decoding algorithm always predicts the top-1 token ta.
  • Lemma 2: When η + m′ ≥ A − B > |η − m′|, the algorithm may predict ta or the no-retrieval token, but not another token.
  • Proof analysis: The probability-gap analysis shows that the top-2 token tb cannot be selected when its gap remains below η.
  • Lemma 3: If η − m′ ≥ A − B > 0, the algorithm always predicts a no-retrieval token.
  • Proof setup: The proof models attacker effects as bounded additions x, y ∈ [0, m′] to the top-two token probability sums A and B.
  • Certification construction: Certification enumerates possible corrupted passage placements, isolates groups, removes corrupted groups, and deduplicates the remaining benign group cases into H.

B. Generalizing to larger group sizes ω

For group sizes ω > 1, certification enumerates possible benign passage-group configurations that remain after corruption and returns the lowest certified robustness value across them.

  • Generalizing to larger group sizes ω: When ω > 1, injected passages can overlap m′ of m = ⌈k/ω⌉ isolated groups, so certification analyzes the remaining m−m′ benign groups.For each injection-position case, RobustRAG simulates isolation and denotes the unaffected groups by ¯Gm,m′.
  • Certification procedure: The certification workflow first enumerates all possible ¯Gm,m′ configurations with BENIGNGROUPCASES(Pk, ω, k′).The resulting set H contains the distinct benign-group configurations considered by certification.
  • Generalizing to larger group sizes ω: Meaningful non-trivial certification requires m−m′ > 0; otherwise, the procedure cannot compute a non-trivial τ value.The group size ω must therefore be chosen to avoid this failure case.
  • Certification procedure: For each configuration, CERTIFYONECASE computes a τ value using the certification procedures for the underlying query, model, reference answer, and attack parameters.CERTIFYONECASE can use either Algorithm 3 or Algorithm 4.
  • Certification procedure: The final certification outcome τ* is the lowest τ computed across all possible benign-group configurations.Taking the minimum accounts for the worst case among enumerated configurations.

C. Generalizing to Passage Modification

RobustRAG extends from passage injection to passage modification by modeling modification as removal followed by injection and changing only the benign-case enumeration procedure.

  • Generalizing to Passage Modification: Passage modification removes k′ original passages and injects k′ malicious passages at arbitrary locations.The rest of RobustRAG’s inference and certification algorithms can remain unchanged under this decomposition.
  • Generalizing to Passage Modification: For each possible removal and injection case, the procedure applies passage isolation and removes groups containing corrupted passages.The resulting distinct benign groups form the cases used for certification.
  • Experiments: RobustRAG achieves good certifiable robustness against both passage modification and passage injection for k′ = 1.The experiment uses Mistral-7B-Instruct with top-10 retrieved passages and reports certifiable robust accuracy in Table V.

APPENDIX B ADDITIONAL DETAILS OF IMPLEMENTATION

The implementation extracts informative words and keyphrases for keyword aggregation, reduces long-form keyword-set size, and defines empirical attack evaluation through attack success rate.

  • Implementation of keyword extraction: Keyword extraction retains lemmatized informative words and keyphrases formed from consecutive informative words.Informative words use POS tags including adjectives, adverbs, nouns, numerals, proper nouns, symbols, and other tags.
  • Implementation of keyword extraction: For long-form generation, the implementation prompts for short phrases and retains only keyphrases longer than two words to reduce certification cost.This addresses keyword sets that can become too large for computationally feasible robustness certification.
  • Detailed setup of empirical attacks: Empirical attacks use injected prompt instructions or GPT-4-generated fake texts supporting attacker-chosen incorrect answers.Malicious content is repeated within a passage for QA and long-generation experiments.
  • Detailed setup of empirical attacks: Attack success rate is the ratio of model responses containing the malicious target texts.The metric complements performance under attack as an empirical robustness measure.

APPENDIX C ADDITIONAL EXPERIMENT RESULTS AND ANALYSES

Additional experiments show that more retrieved passages improve both benign performance and certifiable robustness, while larger corruption increases difficulty and tuning robustness often trades off against clean performance.

  • Experiments with GPT Models: With GPT-3.5, RobustRAG reaches 69.6% certified accuracy on RQA-MC and 37.8% on RQA, while limiting prompt-injection attack success to 5.0% versus over 80% for vanilla RAG.Decoding aggregation was not implemented for GPT-3.5 because of excessive API-call requirements.
  • Impact of retrieved passages k: As the number of retrieved passages increases, both certifiable robustness and benign performance improve across the evaluated datasets and models.The analysis covers RealtimeQA, Natural Questions, and Biography Generation with Llama7B and Mistral-7B.
  • Impact of corruption size k′: Certifiable robust accuracy remains above 50% on RealtimeQA-MC when corruption reaches 3 of 10 passages.This result is reported for Figure 12(a).
  • Impact of keyword filtering thresholds α, β: Larger keyword-filtering thresholds α and β improve robustness to multi-passage corruption but slightly reduce benign performance at k′ = 0.The trade-off is reported for keyword aggregation.
  • Impact of decoding probability threshold η: Increasing the decoding threshold η lowers benign accuracy because it produces more no-retrieval tokens, while robustness effects differ between Natural Questions and RealtimeQA.The paper attributes the dataset difference partly to Mistral’s no-retrieval performance being higher on Natural Questions than RealtimeQA.
  • Impact of decoding probability threshold η: For biography generation, increasing η lowers benign performance but permits tolerance of larger corruption sizes because certification Case 4 cannot occur when η−m′ ≥ 0.The mechanism involves producing more nonretrieved tokens.
Loading 2405.15556v2…