Source-linked AI summary

Speculative RAG: Enhancing Retrieval Augmented Generation through Drafting

Zilong Wang, Zifeng Wang, Long Le, Huaixiu Steven Zheng, Swaroop Mishra, Vincent Perot, Yuwei Zhang, Anush Mattapalli, Ankur Taly, Jingbo Shang, Chen-Yu Lee, Tomas Pfister

arXiv:2407.08223v2cs.CLcs.AI

TL;DR

Knowledge-intensive RAG must balance grounding against the latency of processing retrieved documents, while recent refinement approaches may require extra iterations or tuning. Speculative RAG has a smaller specialist draft answers from diverse document subsets and a larger generalist verify them; the paper reports accuracy gains up to 12.97% and latency reductions up to 50.83% versus conventional RAG systems.

  • Problem

    RAG can improve grounding with retrieved documents, but processing multiple documents increases input length and latency, while refinement methods may add iterations or tuning.

    Source: §1 · pp. 1–2
  • Method

    A smaller specialist generates parallel answer drafts and rationales from distinct retrieved-document subsets, and a larger generalist verifies and selects among them.

    Source: §1 · p. 2, §3.1 · p. 4
  • Results

    Accuracy gains reach 12.97% and latency reductions reach 50.83% compared with conventional RAG systems.

    Source: §5 · p. 11
  • Takeaways & Limitations

    The paper presents collaborative drafting and verification as a way to improve RAG answer quality and speed while reducing input tokens per draft.

    Source: §1 · p. 2, §5 · p. 11

Abstract

from arXiv · show

Retrieval augmented generation (RAG) combines the generative abilities of large language models (LLMs) with external knowledge sources to provide more accurate and up-to-date responses. Recent RAG advancements focus on improving retrieval outcomes through iterative LLM refinement or self-critique capabilities acquired through additional instruction tuning of LLMs. In this work, we introduce Speculative RAG - a framework that leverages a larger generalist LM to efficiently verify multiple RAG drafts produced in parallel by a smaller, distilled specialist LM. Each draft is generated from a distinct subset of retrieved documents, offering diverse perspectives on the evidence while reducing input token counts per draft. This approach enhances comprehension of each subset and mitigates potential position bias over long context. Our method accelerates RAG by delegating drafting to the smaller specialist LM, with the larger generalist LM performing a single verification pass over the drafts. Extensive experiments demonstrate that Speculative RAG achieves state-of-the-art performance with reduced latency on TriviaQA, MuSiQue, PopQA, PubHealth, and ARC-Challenge benchmarks. It notably enhances accuracy by up to 12.97% while reducing latency by 50.83% compared to conventional RAG systems on PubHealth.

1 INTRODUCTION

Knowledge-intensive questions can prompt factual errors from LLMs, while retrieving multiple documents improves grounding but increases input length and latency. SPECULATIVE RAG splits drafting across a smaller specialist LM and verification across a larger generalist LM.

1 INTRODUCTION

  • Knowledge-intensive questions requiring current information or obscure facts can expose LLM factual inaccuracies and hallucinations.LLMs also produce plausible responses using parametric memory, which may not suffice for these questions.Source: §1 · p. 1
  • Retrieving multiple documents helps include information needed for complex questions, but longer inputs increase encoding latency and make grounded reasoning over long contexts challenging.The introduction frames balancing RAG efficiency and effectiveness as a central research question.Source: §1 · p. 1
  • Prior RAG refinements often use multiple iterations or customized instruction-tuning, adding latency or training requirements for generic LMs.These practical costs motivate the paper’s alternative division of work.Source: §1 · p. 2
  • The approach adapts speculative decoding’s small-model drafting and larger-model verification to RAG answer drafts.Each draft comes from a distinct document subset, rather than processing all retrieved documents as one prompt.Source: §1 · p. 2, Figure 1 · p. 2
  • SPECULATIVE RAG has a smaller specialist LM generate parallel drafts from distinct retrieved-document subsets, then a larger generalist LM verifies them.The subsets provide diverse evidence perspectives while reducing input tokens per draft.Source: §1 · p. 2
  • The generalist LM verifies and integrates promising drafts in one pass, without additional tuning, while the specialist handles drafting.The paper reports that this approach enhances subset comprehension and mitigates potential lost-in-the-middle effects.Source: §1 · p. 3

2 RELATED WORKS

Prior RAG work has focused on retrieval decisions and context use, while some approaches require additional tuning or lack high-level reasoning. Speculative RAG extends draft-then-verify from token-level decoding to answer-level drafts.

Retrieval Augmented Generation

  • SAIL and Self-Reflective RAG require additional instruction-tuning, which may be resource-intensive and risk forgetting or over-fitting (Luo et al., 2023b).The related-work discussion also notes that long retrieved contexts can bring computational inefficiency and position bias.Source: §2 · p. 3
  • Corrective RAG uses a lightweight retrieval evaluator but lacks high-level reasoning, while Speculative RAG has a smaller drafter generate drafts for generalist LMs to verify and integrate.Other cited approaches focus on deciding what to retrieve or better utilizing context.Source: §2 · p. 3

Speculative Decoding

  • Speculative RAG extends speculative decoding from token-level drafting to answer-level drafts, using language-modeling objectives to assess whole draft confidence.Traditional speculative decoding accepts or rejects tokens based on generation probabilities.Source: §2 · p. 3

3 SPECULATIVE RETRIEVAL AUGMENTED GENERATION THROUGH DRAFTING

The method treats knowledge-intensive tasks as answering a query from retrieved documents, then divides work between a specialist drafter and a generalist verifier. Diverse document subsets yield parallel drafts and rationales, which the verifier scores to select a final answer.

3 SPECULATIVE RETRIEVAL AUGMENTED GENERATION THROUGH DRAFTING

  • Knowledge-intensive task entries pair a query or statement Q, retrieved documents D, and expected answer A; outputs may be free-form answers, truth values, or multiple-choice selections.The RAG objective is to generate or select the expected answer using retrieved supporting documents.Source: §3 · p. 3

3.1 OVERVIEW

  • SPECULATIVE RAG clusters retrieved documents into perspectives, samples one document from each cluster per subset, and sends subsets to parallel drafter instances.The smaller specialist generates answer drafts and rationales; the generalist verifier scores them and integrates the highest-confidence draft.Source: §3.1 · pp. 3–4

3.2 SPECIALIST RAG DRAFTER

  • MDrafter is a smaller specialist trained to answer from supporting documents and generate both an answer draft and rationale.Training augments query-answer-document examples with concise rationales, then fine-tunes the LM to generate the response and rationale from the query and document.Source: §3.2 · pp. 4–5
  • Question-aware clustering and one-document-per-cluster sampling create subsets with diverse content and multiple perspectives for parallel drafting.The embedding model clusters documents by topic; k controls the number of clusters, and m subsets are constructed.Source: §3.2 · p. 5
  • For each subset, MDrafter generates a draft and rationale, and its conditional generation probability measures rationale reliability and draft confidence.The prompt includes the question and subset documents, and drafting runs in parallel across subsets.Source: §3.2 · p. 5

3.3 GENERALIST RAG VERIFIER

  • The generalist verifier ranks draft-rationale pairs without reprocessing the retrieved documents, using its pretrained language-modeling ability.The verifier can be an off-the-shelf pretrained LM and filters less reliable drafts.Source: §3.3 · p. 5
  • Verifier self-consistency is the conditional probability of a draft-rationale pair given the question; self-reflection scores the probability of a positive reliability judgment.Both scores can be computed in one forward pass of the verifier.Source: §3.3 · p. 5
  • The final score multiplies the drafter, self-consistency, and self-reflection scores, and the highest-scoring draft supplies the final answer.The selection is made across answer drafts for the question.Source: §3.3 · p. 6

4 EXPERIMENTS

Across five RAG benchmarks, SPECULATIVE RAG outperforms comparison systems, with a 12.97% gain over Mixtral-Instruct8x7B on PubHealth. It reduces latency while retaining high performance, and its results support diverse document subsets, concise rationales, and draft verification.

4 EXPERIMENTS

  • The evaluation covers five RAG benchmarks: three open-domain question-answering datasets and two closed-set generation datasets.TriviaQA and PopQA typically require one accurate evidence item, while MuSiQue requires multi-document reasoning; PubHealth verifies medical claims and ARC-Challenge tests science questions.Source: §4 · p. 6

4.1 BASELINES

  • Standard RAG uses all retrieved documents as context, while Self-RAG and Corrective RAG refine retrieval context; comparisons use the same backbone LMs where applicable.Standard RAG baselines include Mistral, Mixtral, Alpaca, Toolformer, and SAIL; Corrective RAG uses an external evaluator, while Self-RAG uses self-reflection tags.Source: §4.1 · p. 6

4.2 EXPERIMENT SETTINGS

  • The experiments use Mistral7B as the drafter base and untuned Mistral7B or Mixtral8x7B as verifiers, with five drafts from two-document subsets on four datasets and ten drafts from six-document subsets on MuSiQue.Top-10 documents are retrieved for TriviaQA, PopQA, PubHealth, and ARC-Challenge; MuSiQue uses top 15.Source: §4.2 · pp. 6–7

4.3 MAIN RESULTS

  • SPECULATIVE RAG consistently outperforms baselines on five benchmarks; MVerifier-8x7B + MDrafter-7B surpasses Mixtral-Instruct8x7B by 0.33% on TriviaQA, 2.15% on MuSiQue, 3.86% on PopQA, 12.97% on PubHealth, and 2.14% on ARC-Challenge.MVerifier-7B + MDrafter-7B also outperforms Self-Reflective and Corrective RAG methods with a comparable number of instruction-tuned parameters.Source: §4.3 · p. 7, Table 1 · p. 7

4.4 EFFECTS OF GENERATED RATIONALE FOR VERIFICATION

  • Generated rationales are significantly shorter than retrieved documents, and adding longer retrieved documents during verification does not consistently improve performance and tends to increase latency.The comparisons use the MVerifier-8x7B + MDrafter-7B configuration on TriviaQA and PubHealth.Source: §4.4 · p. 8, Figure 2 · p. 8, Table 2 · p. 8

4.5 LATENCY ANALYSIS WITH BASELINES

  • SPECULATIVE RAG achieves the lowest latency across the five datasets, reducing latency versus the competitive Standard RAG baseline by up to 11.90% on TriviaQA, 15.07% on MuSiQue, 44.31% on PopQA, 50.83% on PubHealth, and 22.77% on ARC-Challenge.The latency analysis samples 100 cases per dataset and processes them individually without batching; parallel encoding and drafting are cited as efficiency mechanisms.Source: §4.5 · pp. 8–9, Figure 3 · p. 9

4.6 ABLATION STUDIES

  • Multi-perspective document sampling performs best, improving accuracy by up to 1.88% on TriviaQA and 2.23% on PubHealth; removing verification causes declines of 5.69% and 5.37%, respectively.Removing either self-containment or self-reflection scores also reduces performance, with declines around 2.0% on TriviaQA and around 0.8% on PubHealth.Source: §4.6 · pp. 9–10, Table 3 · p. 10

4.7 EFFECTS OF DRAFT NUMBER AND DOCUMENT SUBSET SIZE

  • More parallel drafts can improve performance without adding latency, while increasing documents per draft does not consistently improve performance on TriviaQA and PubHealth.With one supporting document per draft, MVerifier-7B + MDrafter-7B can surpass Mistral-Instruct7B; with two or more, it can surpass Mixtral-Instruct8x7B.Source: §4.7 · p. 10, Figure 4 · p. 10

5 CONCLUSION

Speculative RAG separates drafting from verification, using a small specialist LM to draft from diverse document subsets and a large generalist LM to verify. The authors report improved quality and speed.

5 CONCLUSION

  • Accuracy gains up to 12.97% and latency reductions of 50.83% are reported versus conventional RAG systems.Source: §5 · p. 11
  • The framework delegates drafting to a small specialized RAG drafter and verification to a large generalist LM.Source: §5 · p. 11
  • Parallel drafts from diverse document subsets reduce input tokens and potential position bias over long contexts while providing answer candidates.Source: §5 · p. 11

A INSTRUCTION-TUNING SETTINGS

The appendix tests alternative drafter-data curation and model choices, and isolates the effects of instruction tuning. Its reported comparisons favor Speculative RAG over baselines and show instruction tuning improves drafter document understanding.

A INSTRUCTION-TUNING SETTINGS

B EFFECTS OF INSTRUCTION TUNING

C EFFECTS OF SELF-REFLECTION STATEMENT

The authors vary the verifier's self-reflection statement and report little performance change across alternatives.

C EFFECTS OF SELF-REFLECTION STATEMENT

  • Performance changes little across alternative self-reflection statements, which the authors interpret as stable generalist-LM verification capability.Source: C EFFECTS OF SELF-REFLECTION STATEMENT · p. 17
  • Table 7 analyzes SPECULATIVE RAG performance with different statements used to compute the self-reflection score.Source: Table 7 · p. 17

D EFFECTS OF RATIONALE GENERATION

The study examines rationale generation in both drafting and verification. It reports that rationale-based drafting and verification improve performance, while rationale generation may add inference cost.

D EFFECTS OF RATIONALE GENERATION

E EFFECTS OF DIFFERENT VOLUME OF TRAINING DATA

The section evaluates how instruction-tuning data volume relates to SPECULATIVE RAG performance. Accuracy rises with more training instances, reaching its highest observed level at 40,059 instances.

E EFFECTS OF DIFFERENT VOLUME OF TRAINING DATA

  • The scaling evaluation compared drafter-model training on 10,000, 20,000, 30,000, and 40,059 instances.The primary experiment used 40,059 instances; Table 9 reports the performance analysis across training-data volumes.Source: E EFFECTS OF DIFFERENT VOLUME OF… · p. 18, Table 9 · p. 18
  • Table 9 presents a performance analysis of SPECULATIVE RAG across different instruction-tuning data volumes.Source: Table 9 · p. 18
  • Accuracy increased as training data grew, with the highest performance at 40,059 instances.The evaluation also used incremental subsets of 10,000, 20,000, and 30,000 instances.Source: E EFFECTS OF DIFFERENT VOLUME OF… · p. 18

F EFFICACY OF SPECULATIVE RAG IN MULTI-HOP REASONING

On multi-hop reasoning benchmarks, SPECULATIVE RAG achieves the best reported performance, while its multi-perspective document sampling is reported to improve results on MuSiQue and HotpotQA.

F EFFICACY OF SPECULATIVE RAG IN MULTI-HOP REASONING

F.1 PERFORMANCE BREAKDOWN ON HOTPOTQA

On HotpotQA, multi-perspective sampling is reported to perform better on comparison-type questions, while it is less helpful for bridge-type questions. The appendix also includes prompts for rationale generation, RAG drafting, and standard RAG.

F.1 PERFORMANCE BREAKDOWN ON HOTPOTQA

  • Multi-perspective sampling performs better on comparison-type questions, but is less helpful for bridge-type questions that require identifying a bridge entity.The authors note that their current work is not explicitly designed to pinpoint that entity.Source: §F.1 · p. 19, Table 12 · p. 19

G PROMPT OF RATIONALE GENERATION

  • The rationale-generation prompt for Gemini-Ultra is shown in Figure 5.Source: Figure 5 · p. 20

H PROMPT OF RAG DRAFTING

I PROMPT OF STANDARD RAG

J CASE STUDY

The case studies present representative TriviaQA and PubHealth examples in which multiple drafts are generated for the same question. The authors report that the verifier filters unreliable drafts.

J CASE STUDY

  • The TriviaQA and PubHealth case studies show two drafts for the same question, with the verifier filtering unreliable drafts.The authors describe the first draft as receiving relatively low scores in both figures.Source: J CASE STUDY · p. 22, Figure 9 · p. 22, Figure 10 · p. 23
  • Figure 9 is a representative SPECULATIVE RAG case from TriviaQA, where Dolly Parton is the correct answer.Source: Figure 9 · p. 22
  • Figure 10 is a representative SPECULATIVE RAG case from PubHealth, where False is the correct answer.Source: Figure 10 · p. 23
Loading 2407.08223v2…