RASD: Retrieval-Augmented Speculative Decoding

2025-03-05

AI Summary

Existing speculative decoding methods suffer from low draft quality, poor cross-domain generalization, and limited acceptance length during verification—particularly in out-of-distribution scenarios. To address these issues, we propose Retrieval-Augmented Speculative Decoding (RASD), the first framework to integrate external retrieval into draft generation. RASD employs a probability-driven retrieval tree pruning strategy coupled with a longest-prefix-matching tree fusion mechanism, enabling robust, high-quality draft construction and efficient verification. Crucially, RASD operates without modifying the base language model and is compatible with diverse speculative decoding paradigms. Extensive experiments across document-level question answering, summarization, code generation, and domain-specific QA demonstrate that RASD achieves state-of-the-art speedup ratios while significantly improving cross-domain robustness and average acceptance length.

Technology Category

• Data Mining & Knowledge Management: Conversational Systems for Recommendation & Retrieval
• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Generation

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Speculative decoding accelerates inference in large language models (LLMs) by generating draft tokens for target model verification. Current approaches for obtaining draft tokens rely on lightweight draft models or additional model structures to generate draft tokens and retrieve context from databases. Due to the draft model's small size and limited training data, model-based speculative decoding frequently becomes less effective in out-of-domain scenarios. Additionally, the time cost of the drafting phase results in a low upper limit on acceptance length during the verification step, limiting overall efficiency. This paper proposes RASD (Retrieval-Augmented Speculative Decoding), which adopts retrieval methods to enhance model-based speculative decoding. We introduce tree pruning and tree fusion to achieve this. Specifically, we develop a pruning method based on the draft model's probability distribution to construct the optimal retrieval tree. Second, we employ the longest prefix matching algorithm to merge the tree generated by the draft model with the retrieval tree, resulting in a unified tree for verification. Experimental results demonstrate that RASD achieves state-of-the-art inference acceleration across tasks such as DocQA, Summary, Code, and In-Domain QA. Moreover, RASD exhibits strong scalability, seamlessly integrating with various speculative decoding approaches, including both generation-based and retrieval-based methods.

Problem

• Enhances speculative decoding in large language models.
• Improves efficiency in out-of-domain scenarios.
• Integrates retrieval methods with model-based decoding.

Innovation

• Retrieval-augmented speculative decoding enhances LLM inference.
• Tree pruning optimizes retrieval tree construction.
• Longest prefix matching merges draft and retrieval trees.