Recurrent Drafter for Fast Speculative Decoding in Large Language Models

Authors: Yunfei Cheng, Aonan Zhang, Xuanyu Zhang, Chong Wang, Yi Wang
Primary area: foundation or frontier models, including LLMs
Keywords: speculative decoding, beam search

Abstract:
We present Recurrent Drafter (ReDrafter), an advanced speculative decoding approach that achieves state-of-the-art speedup for large language models (LLMs) inference. The performance gains are driven by three key aspects: (1) leveraging a recurrent neural network (RNN) as the draft model conditioning on LLM's hidden states, (2) applying a dynamic tree attention algorithm over beam search results to eliminate duplicated prefixes in candidate sequences, and (3) training through knowledge distillation from the LLM. ReDrafter accelerates Vicuna inference in MT-Bench by up to 3.5x with a PyTorch implementation on Nvidia H100 GPUs. To demonstrate its practicality in production environments, we integrate ReDrafter into TensorRT-LLM, reaching up to 2.5x speedup on H100 GPUs. We also validated its effectiveness for on-device applications by implementing the approach in MLX and benchmarking performance on Metal GPUs in Apple Silicon chips, achieving up to 2.3x speedup.

Decision: reject
Reviewer scores: 3, 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730175196750

thread_type:
official_review

<thread_object>
creator: reviewer_uws2

<content>
summary:
This paper introduces ReDrafter, which achieves state-of-the-art speed improvements for LLM inference in speculative decoding through the use of RNNs, dynamic tree attention, and knowledge distillation from the LLM. ReDrafter demonstrates significant speed enhancements compared to baseline methods across various frameworks and chips.

strengths:
1. This paper tackles a significant issue in LLM inference by promoting speculative decoding, demonstrating considerable speed improvements compared to state-of-the-art methods like Medusa and Eagle.
2. The method has been extensively tested across various chips and frameworks, including H100, Apple’s M2, TensorRT-LLM, MLX, and PyTorch.
3. The paper is well written and effectively conveys the main design of the algorithm.

weaknesses:
1. Could you include a comparison with Eagle-2? It's listed in your references but isn’t included in the evaluation.

questions:
See weakness. Thanks!
</content>
</thread_object>

</thread>
<thread>
creation date:
1730182990202

thread_type:
official_review

<thread_object>
creator: reviewer_swoj

<content>
summary:
This paper aims to accelerate the inference of LLMs. The proposed architecture utilized speculative decoding technology and applied a draft model with an RNN structure. They also use dynamic tree attention to further accelerate inference.

strengths:
1.The paper is well-written and flows very smoothly. The framework and workflow of the entire method are clearly articulated.
2.This paper conduct extensive experiments to demonstrate the effectiveness of their approach, covering various models and different devices.

weaknesses:
1. Although there is textual description of ReDrafter, providing a structural diagram would better assist readers in understanding the structure of the ReDrafter Model.
2. I believe the paper lacks sufficient explanation for the idea of using an RNN-based draft model. For example, I think it would be beneficial to include an experiment that substitutes the RNN with a single-layer decoder model to observe the resulting changes. 
3.Combining embeddings and hidden states as input is not a novel approach.[1]
4. The article mentions that tree attention is dynamic and states that it "relies on individual beam search results," but it does not explicitly explain the relationship. Alternatively, what does the author consider to be “static” tree attention? If it does indeed exist, I believe it should be compared in the ablation experiment section between dynamic TA, static TA, and without TA.

[1] EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty, ICML 2024

questions:
1. Although there is textual description of ReDrafter, providing a structural diagram would better assist readers in understanding the structure of the ReDrafter Model.
2. I believe the paper lacks sufficient explanation for the idea of using an RNN-based draft model. For example, I think it would be beneficial to include an experiment that substitutes the RNN with a single-layer decoder model to observe the resulting changes. 
3.Combining embeddings and hidden states as input is not a novel approach.[1]
4. The article mentions that tree attention is dynamic and states that it "relies on individual beam search results," but it does not explicitly explain the relationship. Alternatively, what does the author consider to be “static” tree attention? If it does indeed exist, I believe it should be compared in the ablation experiment section between dynamic TA, static TA, and without TA.

[1] EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty, ICML 2024
</content>
</thread_object>

</thread>
<thread>
creation date:
1730673728280

thread_type:
official_review

<thread_object>
creator: reviewer_euga

<content>
summary:
The paper proposes a new speculative decoding algorithm (ReDrafter), that uses the final hidden state of the LLM as input to an RNN. The speedup is further enhanced by using dynamic tree attention over beam search results for the draft model. ReDrafter provides speedups over other decoding algorithms such as EAGLE and Medusa.  The algorithm is implemented and compared with other decoding techniques in multiple frameworks (Pytorch, TensorRT-LLM,  MLX on Apple M2 Ultra Metal GPU). Finally, a new training objective for the draft model is proposed based on knowledge distillation to better align the draft model with the target LLM.

strengths:
Several ablation studies are done to show the effectiveness of all the main contributions, including ablations for the knowledge distillation training objective and the ablation for the 
dynamic tree attention. This makes the contributions more clear.

The algorithm is implemented in several frameworks, and performance improvements are shown in all frameworks, making the practical significance of the algorithm higher. 

The work is quite original while building on previous work on knowledge distillation and using tree structures to save computation. 

The ReDrafter approach gives significant speedups over EAGLE and Medusa on several frameworks and model sizes.

weaknesses:
The paper does not seem to include details on the compute resource/time requirements to train the draft model. This is important to understand for others wanting to train ReDrafter on their local LLM. 

For many LLMs, there exist smaller LLMs from the same family that can be used as draft models. It would have been good to compare against approaches using separate draft models. 

Very minor comments: 
There is a parenthesis missing in equation 1. 
In the first sentence of the abstract, I would use LLM rather than LLMs.

questions:
Will the code be open sourced?

Can you clarify why you are only comparing against the EAGLE and Medusa papers? 

How is the performance of ReDrafter compared to approaches that use a separate draft model?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730784837786

thread_type:
official_review

<thread_object>
creator: reviewer_zzts

<content>
summary:
This paper proposes "ReDrafter", a speculative decoding algorithm that uses an RNN as the draft model. The RNN takes as input the outputs of the LLMs last hidden layer, and then uses beam search to construct a set of possible continuations, which are then verified using an efficient "dynamic tree attention" algorithm by the target model. The RNN is trained using knowledge distillation, to maximize alignment with the target model. Empirically, ReDrafter shows large speedups (on both H100s in PyTorch/TensorRT-LLM, and on Apple Silicon in MLX), generally attaining meaningfully larger speedups than Eagle and Medusa.

**Update**: After reviewing the authors’ responses, I have lowered my score to a 3, as I remain unconvinced that beam search can be used by the draft model during stochastic decoding while maintaining the output distribution of the target model.

strengths:
- Using an RNN as a draft model, taking as input the outputs of the target model's last hidden layer, is a nice idea. This allows introducing dependence between the draft model's speculated tokens (unlike Medusa), while also leveraging the strong representations from the target model (like Eagle), while not requiring storing a KV cache for the draft model.
- The empirical speedups attained by ReDrafter appear quite large.

weaknesses:
- **Most importantly**: It's unclear to me how beam search can be used by the draft model while maintaining the guarantee that the output distribution of the target model is unchanged, when stochastic decoding is used. Can you please clarify?
- It seems that a large percentage of the ReDrafter speedups may come from the large beams used during decoding, as opposed to from an improved draft model architecture. I would have appreciated more careful ablations to clarify this. For example, could one use a transformer draft model (e.g., Llama-3.2-1B for Llama-3.1-70B) with beam search, and attain similar/larger speedups than the RNN draft model? Could you show the relative speeds and acceptance rates of ReDrafter vs. standalone draft models / Medusa / Eagle, to understand the relative merits of these different draft model architectures? How do these different draft model architectures compare when a simple chain of tokens is speculated (like in Leviathan et al), instead of a tree of tokens?
- Given that much of the speedups seem to come from the beam search process, it's important to compare with methods like Sequoia and Eagle2, which also perform speculative decoding using large trees of tokens (instead of simply token sequences).
- It would have been nice to see experiments with Llama 3 models.

questions:
- In the equation in line 156, it doesn't seem $h$ or $g_t$ are used in the recurrence equation for $s_t$. Can you please clarify?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732562785730

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Overall Response

comment:
We deeply appreciate the insightful feedback and thoughtful recommendations provided by all the reviewers.

We appreciate their recognition of the significance of ReDrafter, and our contributions in:
- Innovative and effective approaches in model and algorithm development (zZts, EUGA, Uws2)
- Strong empirical performance compared to existing methods, supported by comprehensive ablation studies (zZts, EUGA, SwoJ, Uws2)
- Implementation across multiple platforms (zZts, EUGA, SwoJ, Uws2)
- Clear and well-structured writing (SwoJ, Uws2)

We hope that our response will help address the following concerns:
- Guarantee for ReDrafter's output distribution (zZts):
  - We provide a detailed example, along with relevant references, to clarify that ReDrafter preserves the output distribution. 
- Training cost of ReDrafter (EUGA):
  - Additional details about the training cost are included to address this concern.
- Comparison with existing methods (zZts, Uws2):
  - We include comparisons with Sequoia and EAGLE2 for a comprehensive evaluation.
- Ablation studies to justify design choices:
  - We include ablations using standalone draft models (zZts, EUGA), chain-based speculative methods (zZts), dynamic vs. static tree attention (SwoJ), and RNN vs. single-layer decoder (SwoJ).
- Open-source code (EUGA):
  - An anonymized version of the code is included in the supplementary material.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a novel method called the Recurrent Drafter for accelerating speculative decoding in LLMs. The primary claim is that the proposed approach improves decoding efficiency by leveraging a recurrent mechanism to generate multiple tokens in parallel, reducing computational overhead compared to traditional autoregressive decoding. The authors present empirical results demonstrating speedups over baseline methods while maintaining comparable output quality. Additionally, the paper discusses theoretical justifications for the proposed method and its compatibility with existing LLM architectures.

The paper presents a novel approach to speculative decoding, addressing an important problem in improving the efficiency of LLMs. However, the evaluation is limited in scope, with experiments conducted on a narrow set of datasets and tasks, which raises concerns about the generalizability of the findings. Baseline comparisons are not comprehensive, as stronger state-of-the-art baselines are missing, making it difficult to contextualize the improvements. Furthermore, reviewers raised concerns about whether this tradeoff is acceptable across different contexts. With these concerns, I recommend rejection for this submission. I hope the authors can improve the paper and submit it to a future venue.