Source-linked AI summary
REST: Retrieval-Based Speculative Decoding
Zhenyu He, Zexuan Zhong, Tianle Cai, Jason D. Lee, Di He
TL;DR
LLM generation is expensive, and conventional speculative decoding requires a suitable small draft model. REST instead retrieves draft tokens from a datastore, organizes them with a Trie, and verifies them with the target LLM, achieving 1.62× to 2.36× speedups across code and text benchmarks.
Problem
Conventional speculative decoding requires a draft model that is small, predictive, vocabulary-compatible, and easy to integrate, making high-quality draft-model acquisition challenging.
Method
REST is a training-free speculative-decoding method that retrieves draft tokens from a datastore, selects them with a Trie, and verifies them using tree attention.
Results
REST achieves 1.62× to 2.36× speedups across HumanEval and MT-Bench evaluations of 7B and 13B language models.
Takeaways & Limitations
REST provides a straightforward, plug-and-play approach for accelerating inference across existing language models without additional training.
Takeaways & Limitations
REST performance depends directly on the accuracy and completeness of its datastore, and retrieval may struggle with context-dependent cases such as personalized variable names.
Abstract
from arXiv · showhide
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 models, 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.62X to 2.36X on code or text generation. The code of REST is available at https://github.com/FasterDecoding/REST.
1 Introduction
REST addresses the high inference cost of autoregressive LLM generation by replacing the difficult-to-train draft model in speculative decoding with retrieval from a datastore. It constructs and verifies draft candidates from matching corpus continuations, achieving substantial speedups on code and general text benchmarks.
- Autoregressive LLM inference is inefficient because each token requires reloading the billion-parameter model from HBM to accelerator cache.
- Speculative decoding reduces inference cost by drafting tokens with a smaller language model and verifying them with the target LLM.
- REST replaces the parametric draft model with a non-parametric retrieval datastore that can integrate with and accelerate any LLM.
- REST builds a Trie from retrieved continuation candidates, selects high-frequency draft tokens, and verifies them with tree attention in one LLM forward pass.
- 2.12× to 2.36× speedup is achieved on HumanEval for 7B and 13B CodeLlama models, while MT-Bench yields 1.62× to 1.77× for Vicuna.
2 Related Work
REST belongs to lossless inference acceleration but differs from conventional speculative decoding by retrieving draft tokens from a broad datastore rather than generating them with a smaller model. Compared with LLMA, it supports comprehensive datastores and many retrieved instances.
- Speculative decoding reduces target-LLM executions by generating a draft with a smaller model and verifying it in one forward pass.
- REST retrieves draft tokens from a datastore instead of using a separately trained smaller language model.
- Unlike LLMA, REST retrieves from a comprehensive datastore rather than only from referred contexts supplied during generation.
- REST is designed to handle many retrieved instances, whereas LLMA typically uses one or a handful.
3 Retrieval-Based Speculative Decoding
REST constructs speculative drafts by exact-match retrieval, Trie-based frequency selection, and tree-attention verification. Its training-free design avoids the size, prediction, vocabulary, and integration challenges of conventional draft models while sharing computation across common prefixes.
- Draft verification: The LLM accepts draft tokens sequentially from the start and rejects tokens after the first verification mistake.
- Background: Speculative Decoding: Classic speculative decoding uses a small language model to generate multiple draft tokens before target-LLM verification.
- Our Approach: REST: REST addresses the challenge of finding a lightweight, accurate, compatible draft model through a training-free retrieval-based approach.
- Datastore construction: REST builds a datastore from context-continuation pairs extracted from a text or code corpus.
- Retrieving from the datastore: At inference, REST uses the current context to retrieve continuation candidates by exact matching against the longest available suffix.
- Retrieving from the datastore: The retrieval algorithm starts at nmax, decreases the suffix length until a match is found, and incurs less than 6% retrieval overhead in experiments.
- Draft construction: A Trie assigns frequency weights to retrieved prefixes, enabling selection of high-frequency draft sequences from the candidate set.
- Draft verification: Tree attention packs drafts with shared prefixes into a pseudo sequence so shared-prefix computation occurs only once during LLM verification.
4 Experiments
REST is evaluated against standard autoregressive decoding and speculative decoding on code and dialogue generation, using retrieval-based drafts under greedy and nucleus sampling. Across these settings, REST achieves substantial speedups, with performance depending on domain and sampling strategy.
- Experimental Setup: REST uses greedy and nucleus sampling, accepting draft tokens only when they match the LLM's sampled tokens.Nucleus sampling selects from the most probable tokens whose cumulative probability reaches threshold p.
- Experimental Setup: Experiments evaluate REST on HumanEval and MT-Bench with CodeLlama and Vicuna 7B and 13B models.HumanEval contains Python programming problems, while MT-Bench contains multi-turn dialogue questions.
- Main Results: 2.16× to 2.36× speedup is achieved for CodeLlama on HumanEval compared with standard autoregressive and speculative decoding.The comparison uses the best speculative-decoding configuration across draft-token counts and small draft models.
- Main Results: 1.62× to 1.77× speedup is achieved for Vicuna on MT-Bench.REST's speedup is lower with nucleus sampling than with greedy sampling, which the authors associate with sampling randomness.
- Main Results: REST's speed improvements are domain-dependent, with substantially greater speedup on HumanEval than on MT-Bench.This pattern is also reported for speculative decoding and Medusa.
- Main Results: Less than 1 ms per token is required on average for retrieval and Trie construction.The authors treat this retrieval overhead as negligible for practical purposes.
5 Ablation Study
REST’s ablations examine datastore size, draft-token selection, and maximum suffix length. Larger datastores and Trie-based selection improve performance, while nmax values above 6 maintain consistently high generation speed.
- Ablation Study: The ablation study evaluates REST’s efficiency and effectiveness through component-focused experiments.The study includes comparisons involving datastore size, draft-token selection, and maximum suffix length.
- Effect of the datastore size: Larger datastores improve retrieved draft-token accuracy and increase generation speed, although speedup grows less than Mean Generated Length because retrieval adds overhead.The authors expect larger datastores and more CPU cores for retrieval in industry applications.
- Effect of draft token selecting strategies: Trie-based draft-token selection enhances performance compared with randomly sampling retrieved continuation candidates.The comparison samples at most eight sequences, each truncated to length 8, yielding up to 64 draft tokens.
- Effect of the choice of the maximum suffix length: nmax values below 6 substantially increase generation time, whereas values above 6 keep generation speed consistently high and largely unchanged.The results suggest that precisely optimizing nmax is unnecessary in practice.
6 Conclusion
REST replaces a small draft language model with datastore retrieval and Trie-based draft-token selection. It integrates into existing language-model generation without additional training, while larger-scale retrieval and datastore compression remain future directions.
- Conclusion: REST uses a datastore to retrieve draft tokens and a Trie to select the most probable candidates.The approach is designed as retrieval-based speculative decoding rather than small-language-model drafting.
- Conclusion: REST integrates into existing language-model generation without requiring additional training.The paper describes the method as straightforward to implement and easy to integrate.
- Conclusion: The authors identify large-scale retrieval and datastore-size reduction under limited disk storage as future work.These directions target broader retrieval and lower storage requirements without compromising performance.
Limitations
REST’s limitations concern dependence on datastore quality and limited in-context ability, including challenges with personalized code variables.
- REST’s performance is directly influenced by the accuracy and completeness of its datastore.The paper suggests constructing datastores from content generated by the LLM to improve alignment.
- REST lacks in-context abilities for tasks such as retrieving personalized variable names in code generation.The paper identifies handling such context-dependent complexities as an open question for retrieval methods.
- The experiments compare REST with small draft LMs across model sizes and draft-token counts, using reproduced implementations and torch.compile.Tested draft LMs include Llama 68M, Llama 160M, TinyLlama 1.1B, and TinyLlama-Chat 1.1B, with 1 to 15 draft tokens.
A.1 Results on CodeLlama 7B
On CodeLlama 7B, generation speed varies with sampling strategy, draft model, and draft-token count, with the best reported settings reaching 15.90 and 18.83 ms/token.
- Greedy Sampling: 15.90 ms/token is achieved with Llama 68M and 4 draft tokens under greedy sampling.This is the best setting reported for CodeLlama 7B with greedy sampling on HumanEval.
- Nucleus Sampling: 18.83 ms/token is achieved with TinyLlama 1.1B and 4 draft tokens under nucleus sampling.This is the best setting reported for CodeLlama 7B with nucleus sampling on HumanEval.
- Greedy Sampling: 19.39 ms/token is achieved with TinyLlama 1.1B and 10 draft tokens under greedy sampling for CodeLlama 13B.The result is reported as the best setting for the CodeLlama 13B greedy-sampling experiment.
- Nucleus Sampling: 22.68 ms/token is achieved with TinyLlama 1.1B and 6 draft tokens under nucleus sampling for CodeLlama 13B.The result is reported as the best setting for the CodeLlama 13B nucleus-sampling experiment.
A.3 Results on Vicuna 7B
On Vicuna 7B, the reported best generation speeds are 19.44 ms/token with greedy sampling and 20.65 ms/token with nucleus sampling, both using Llama 68M.
- Greedy Sampling: 19.44 ms/token is achieved with Llama 68M and 3 draft tokens under greedy sampling.This setting is reported as the best for Vicuna 7B with greedy sampling.
- Nucleus Sampling: 20.65 ms/token is achieved with Llama 68M and 3 draft tokens under nucleus sampling.This setting is reported as the best for Vicuna 7B with nucleus sampling.
A.4 Results on Vicuna 13B
On Vicuna 13B, the reported best generation speeds are 29.80 ms/token with greedy sampling and 31.78 ms/token with nucleus sampling, both using Llama 68M and 3 draft tokens.
- Greedy Sampling: 29.80 ms/token is achieved with Llama 68M and 3 draft tokens under greedy sampling.This setting is reported as the best for Vicuna 13B with greedy sampling.
- Nucleus Sampling: 31.78 ms/token is achieved with Llama 68M and 3 draft tokens under nucleus sampling.This setting is reported as the best for Vicuna 13B with nucleus sampling.
B Additional Ablation Studies
The ablations show that REST benefits from larger datastores, while draft-token limits and matched-suffix distributions reveal practical efficiency boundaries. Most matched suffixes are short, and excessive draft-token limits can reduce speed.
- Datastore size: Larger datastores produce better speedup rates, while REST still achieves notable speedup with a 465 MB datastore.The comparison uses a 465 MB ShareGPT-derived datastore and a 12 GB UltraChat-derived datastore on MT-Bench.
- Draft-token limit: Speed initially increases with the maximum number of draft tokens, stabilizes at approximately 11.75 ms per token beyond 48 tokens, and slows when the count exceeds 200.The slowdown beyond 200 tokens makes smaller limits more efficient by avoiding unnecessary GPU strain.
- Matched suffix length: 85% of matched cases have suffix lengths from 2 to 9 tokens, while lengths from 10 to 16 account for 15%.The distribution indicates that shorter matched suffixes constitute the clear majority of cases.
- Evaluation settings: The datastore-size ablations evaluate generation speed under both greedy and nucleus sampling on MT-Bench.The supplied table captions identify separate datastore-size evaluations for greedy and nucleus sampling.