REST : Retrieval-Based Speculative Decoding (2024)

Abstract: We introduce Retrieval-Based Speculative Decoding (REST), a novel algorithm designed to speed up language model generation. The key insight driving the development of REST is the observation that the process of text generation often includes certain common phases and patterns. Unlike previous methods that rely on a draft language model for speculative decoding, REST harnesses the power of retrieval to generate draft tokens. This method draws from the reservoir of existing knowledge, retrieving and employing relevant tokens based on the current context. Its plug-and-play nature allows for seamless integration and acceleration of any language model, all without necessitating additional training. When benchmarked on 7B and 13B language models in a single-batch setting, REST achieves a significant speedup of 1.62 \times to 2.36 \times on code or text generation. The source code of REST is available at https://github.com/FasterDecoding/REST.

## introduction transformer-based large language models (llms) have emerged as a foundation model in natural language processing @xcite @xcite @xcite @xcite . while they achieve impressive performance across various tasks, the inference cost is huge in practical scenarios. during inference, the model autoregressively uses the preceding context to generate the next token. each iteration requires reloading the billion-parameter llm from the high-bandwidth memory (hbm) to the on-chip cache of modern accelerators like gpus, making the whole generation inefficient and time-consuming. a recent direction in accelerating the llm generation is to reduce the number of forward processes with llms while guaranteeing the quality of the output sequence simultaneously. speculative decoding @xcite @xcite is one of the typical approaches in this direction. intuitively, speculative decoding methods leverage a small lm to generate tokens with less computational cost. during inference, the method first uses the small lm to create a draft token sequence and then uses the llm for verification. if the predictions from both models are consistent, we can accept the draft and return it to the user. here, the actual token generation is carried out using the small lm, and the large lm is only used to validate the draft, which can be performed in parallel and requires reloading the memory only once. consequently, the entire framework of speculative decoding reduces the overall inference cost. however, obtaining a high-quality draft model remains challenging: it must balance small size and strong predictive power while matching the vocabulary of the base model; also, it should integrate well into a distributed system for serving. therefore, people often need to train a draft model specifically for their model and use cases @xcite @xcite . in this study, rather than relying on an additional small lm, we investigate using a data corpus directly to construct the draft token sequence in speculative decoding. we develop a retrieve-based approach, called retrieval-based speculative decoding (rest) (figure 1 ). compared to previous approaches, our retrieval-based system replaces the parametric draft model with a non-parametric retrieval datastore, which can easily port to any llm and accelerate its inference. to use rest, the first step is constructing the datastore. in this paper, we leverage either the pretraining data corpus or the instruction-tuning data corpus to build our datastore, which serves as the source for the draft token sequence. during each inference step, we first use previous tokens (pre-generated tokens or prompt tokens) as queries to identify exact matches in the datastore. the subsequent tokens from these exact matches are considered generation candidates. a trie is constructed using these candidates. the nodes with the highest frequencies are selected as the draft tokens. this sequence then undergoes verification by the llm through a single forward pass, aided by a meticulously designed attention mask known as tree attention @xcite @xcite . as many subsequences during generation likely appear in the datastore, rest can frequently generate multiple correct tokens per step. we conduct extensive experiments to test the efficiency and effectiveness of rest in different scenarios. for the code domain, we use a portion of python pretraining code (2.7m samples) from the stack @xcite as the datastore and accelerate codellama @xcite 7b and 13b respectively. the results show on hu-maneval @xcite rest achieves 2.12× to 2.36× speedup. for the general domain, we construct a datastore using ultrachat @xcite , containing around 774k conversations. the results show on mt-bench @xcite rest accelerates 7b and 13b vicuna @xcite by 1.62× to 1.77× respectively. ## related work improving the efficiency of llm inference has been an emergent research direction in recent years. broadly, previous attempts can be divided into two categories: lossless acceleration and lossy acceleration. lossy acceleration approaches aim to learn efficient models that can execute faster and act similarly to a target llm. these methods include pruning @xcite @xcite , quantization @xcite @xcite @xcite and knowledge distillation @xcite . lossless acceleration strategies focus on directly accelerating the target llm from different perspectives, such as memory and io optimization @xcite @xcite , and ways to reduce the function calls of llm during decoding, e.g., specu-lative decoding @xcite @xcite @xcite . this work falls within the second branch. speculative decoding @xcite @xcite leverages a smaller model to generate a draft and use llm to verify the draft tokens with a single forward pass. in this framework, blockwise parallel decoding @xcite and medusa @xcite train multiple heads based on the llm for draft token generation. our method diverges from these approaches by retrieving draft tokens from a datastore, presenting a novel avenue for efficiency improvement in large language model generation. while there is a similar study, llma @xcite , that employs retrieval to accelerate generation, our work distinguishes itself in two primary ways: (1) the llma approach is tailored towards scenarios where referred contexts (as in retrieval-augmented generation and cache-assisted generation) are provided during generation, and it only retrieves from these referred contexts. in contrast, our method retrieves draft tokens from a comprehensive datastore, thereby not being confined to a small context. ## retrieval-based speculative decoding in this section, we first provide notations and a background overview of speculative decoding and then introduce our proposed rest framework. ## ablation study to gain a deeper understanding of our method, we conduct a series of ablation studies and analyses focused on each individual component. more ablation studies can be found in appendix b. effect of the datastore size increasing the size of the datastore is an effective strategy for enhancing the accuracy of retrieved draft tokens in the trie, which in turn can significantly boost generation speed. in ## conclusion in this work, we propose rest: retrieval-based speculative decoding. instead of requiring a small lm, rest employs a datastore for retrieving and employing draft tokens. we construct a trie to select the most probable draft tokens. rest is not only straightforward to implement but also easily integrates into the generation processes of any existing language models without necessitating ad-ditional training. we would like to explore largescale retrieval in the next step. for situations where disk storage is limited, we will also explore methods of minimizing the size of the datastore without compromising performance.