Speculative RAG: Enhancing Retrieval Augmented Generation through Drafting

TL;DR

Speculative RAG tackles latency and context-bias in retrieval-augmented generation by splitting the task into drafting by a smaller specialist LM and verification by a larger generalist LM. Drafts are produced in parallel from diverse document subsets, reducing per-draft input length and enabling better coverage of evidence. A small RAG drafter plus a generalist verifier yields state-of-the-art results on TriviaQA, MuSiQue, PopQA, PubHealth, and ARC-Challenge, with substantial latency reductions. The approach demonstrates the value of task decomposition and parallel drafting for efficient, robust knowledge-intensive NLP systems.

Abstract

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.

Figures (10)

• Figure 1: Illustration of different RAG approaches. Given a knowledge-intensive query $Q$ and retrieved documents, (a) Standard RAG incorporates all documents into the prompt, increasing input length and slowing inference; (b) Self-Reflective RAG asai2023self requires specialized instruction-tuning of the general-purpose language model (LM) to generate specific tags for self-reflection; (c) Corrective RAG yan2024corrective employs an external retrieval evaluator to refine document quality, focusing solely on contextual information without enhancing reasoning capabilities; (d) In contrast, our proposed Speculative RAG leverages a larger generalist LM to efficiently verify multiple RAG drafts produced in parallel by a smaller, specialized LM. Each draft is generated from a distinct subset of retrieved documents, providing diverse perspectives on the evidence while minimizing the number of input tokens per draft.
• Figure 2: Average number of tokens in the generated rationale and the retrieved documents in TriviaQA, MuSiQue, PopQA, PubHealth, and ARC-Challenge. The generated rationale is of much shorter length than the original retrieved documents.
• Figure 3: Latency analysis of Standard RAG, Self-RAG, and Speculative RAG on TriviaQA, MuSiQue, PopQA, PubHealth, and ARC-Challenge. The latency difference between Standard RAG/Self-RAG and Speculative RAG is highlighted in red (+$x$%). The latency varies across different datasets due to different retrieved document lengths. Speculative RAG encodes the retrieved documents in parallel and generates answer drafts with a smaller RAG drafter. This significantly improves the efficiency.
• Figure 4: Performance analysis of Speculative RAG with (a) different numbers of drafts, and (b) different supporting document subset size on TriviaQA and PubHealth.
• Figure 5: Prompt of Rationale Generation for Gemini-Ultra