RASD : Retrieval-Augmented Speculative Decoding (2025)

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.

## introduction transformer-based large language models (llms) @xcite exhibit remarkable capabilities and are extensively applied across diverse domains @xcite . however, autoregressive generation in llms produces tokens sequentially, resulting in slow inference speeds @xcite . to address 0.0 0.5 1.0 1.5 2.0 2.5 3.0 speedup rate 2.84 2.23 1.65 2.29 2.98 2.31 2.10 2.43 2.72 2.05 2.46 2.86 2.56 2.57 (pld retrieval) (rest retrieval) rasd (ours) eagle2 this issue, an innovative approach called speculative decoding has been introduced @xcite . speculative decoding modifies the inference process by dividing the llm's task into two phases: a cost-efficient draft phase and a parallel verification phase. this strategy significantly improves computational parallelism in llm inference. by enabling llms to generate multiple tokens simultaneously and minimizing the time spent in the drafting and verification phases, speculative decoding reduces the overall inference time @xcite . in speculative decoding, if the majority of draft tokens are correct, the overall decoding steps can be significantly reduced. therefore, obtaining draft tokens with a high acceptance rate is essential. some researchers employ small draft models to predict draft tokens @xcite @xcite . these draft models are trained to match the distribution of the target model and then generate draft tokens. other researchers have proposed using a parameter-efficient model structure to generate the next k candidate tokens in a single forward pass of the target model @xcite . this approach also requires training the additional structure. an alternative approach is retrieval-based speculative decoding, which is a train-free method @xcite . here, a retrieval library pre-defines tokens to follow the suffix of the current content as draft tokens. this approach is fast and eliminates the need for additional model training. although it is generally less effective than generation-based methods, it performs well in specific knowledge-intensive scenarios. in draft model-based speculative decoding, the acceptance rate of draft tokens hinges on the draft model's capabilities. however, these models typically have simple structures with limited parameters, which restricts their ability to effectively retain training data. typically, a draft model is trained on a general dataset to perform well in common scenarios, such as the sharegpt dataset. as a result, in specific scenarios, the draft model's token acceptance rate tends to be low. this issue arises because the draft model lacks capabilities for outof-domain scenarios. there has been work on improving the generalizability of the draft model through feature-level supervision, which closes the representation gap between the draft model and the target model without requiring large amounts of data @xcite . however, we have evaluated these methods on our domain datasets and discover that the draft model still does not perform well in the specific downstream domain after our test of business data. moreover, generating the draft token sequence requires multiple forward passes through the draft model, making the process time-intensive. consequently, previous methods limit the length of candidate draft tokens. this restriction leads to a low upper limit on the acceptance length during a single verification step. in this paper, we propose the rasd (retrieval-augmented speculative decoding) method to address the draft model's shortcomings mentioned above. first, we design an algorithm to select suitable retrieval results and then combine the draft tokens using a tree-fusion approach. this enables the generation of draft tokens that incorporate information from the language model and retrieval. finally, the target model verifies the draft tokens to achieve acceleration. we tested two retrieval methods pld @xcite and rest @xcite to enhance model-based speculative decoding. as shown in figure 1 , the pld method performs well in tasks where the output includes input content. therefore, we conducted comprehensive experiments using rasd (pld) on datasets with these characteristics across four tasks: code generation, document question answering, summarization, and retrieval-augmented generation. we have conducted experiments on humaneval @xcite , cnn/daily mail @xcite , mfqa @xcite , and dpr @xcite . rasd (pld) achieved the best results in all four tasks. to extend the applicability of the rasd method, we also tested the rest retrieval method. rest is theoretically beneficial for any dataset when the database's data distribution closely matches the test set, excelling in knowledge-intensive tasks. we experimented with rasd (rest) on three tasks: humaneval, medqa foot_0 , and nl2sql. the results demonstrate that rasd (rest) effectively improves the acceleration ratio of speculative decoding. ## retrieval-augmented speculative decoding in this section, we first present the background of speculative decoding and then introduce our proposed rasd framework. ## conclusion in this work, we propose rasd: retrieval-augmented speculative decoding. we use retrieval methods to improve the quality of candidate sequences from draft models. concretely, rasd improves the speedup of speculative sampling on the out-of-domain datasets that are difficult for the draft model to handle and improves the maximum output length of draft models. we develop a pruning method to select appropriate retrieval sequences and we fuse the sequence tree from the draft model with the retrieval tree, creating a final combined tree for verification. experiments have proven that our method is effective and exhibits strong scalability.