RASD: Retrieval-Augmented Speculative Decoding

TL;DR

RASD tackles the inefficiency of draft-model-based speculative decoding by introducing retrieval-augmented drafting, supported by pruning and longest-prefix tree fusion to merge draft and retrieved sequences. The approach enables a unified, tree-structured draft for recursive verification, achieving state-of-the-art speedups across diverse tasks and demonstrating strong scalability with both PLD and REST retrievals. By leveraging suffix-based retrieval and efficient pruning, RASD mitigates out-of-domain weaknesses of draft models and expands the feasible draft length, expanding practical acceleration. The work offers a flexible, plug-and-play framework that can complement existing speculative decoding methods to significantly reduce inference latency in large language models.

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.

Figures (3)

• Figure 1: Speedup Performance of EAGLE2 vs. RASD with PLD and REST Retrieval Methods on Qwen2.5-14B.
• Figure 2: An overview of RASD. We obtain the draft token tree and retrieval results through the draft model generation and retrieval methods, respectively. In the next step, we construct and prune the retrieval tree. Then, we fuse the two trees, resulting in the retrieval-augmented draft token tree. Finally, the retrieval-augmented draft token tree is verified recursively. In the figure, green tokens denote $y_0$ in the current turn, and red tokens are accepted by the target model.
• Figure 3: RASD performance with different lengths of retrieval candidates compared with the baselines