Recurrent Drafter for Fast Speculative Decoding in Large Language Models

Thanks for your comments.

Compare with Eagle2.

We located Eagle2's [1] the source code and checkpoint on the authors' website and successfully reproduced their results on MT-Bench and Alpaca within our experimental environment with a single H100 GPU. This approach ensures consistency with ReDrafter's setup. The reproduced results are presented below.

Inference Speed-up VS. AR ↑| Tokens/Step↑

In the results presented above, ReDrafter surpasses EAGLE2 at a temperature setting of 1, highlighting the advantage of beam search in managing uncertainty within the distribution. On the other hand, EAGLE2 shows superior performance at a temperature of 0, owing to its use of a greedy search approach for token tree construction and additional optimizations tailored for zero-temperature conditions.

[1] Y. Li et al. EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees.

Thanks for your comments.

To make it more clear for readers, we begin by presenting a structured workflow for ReDrafter:

---

1. Initialize - Start with the previously generated tokens and the last hidden state.

2. While stopping criteria are not met:
   2.1. Draft Model Generation - Use beam search to generate a beam (set of candidate sequences).
   2.2. Dynamic Tree Attention - Flatten and compress the beam into a packed beam. Then formulate an attention mask.
   2.3. LLM Verification - Perform a forward pass through the LLM to compute log probabilities for all proposed tokens.
   2.4. Candidate Selection - Select the best candidate with the longest accepted prefix.
   2.5. Update State - Append accepted tokens to the prefix and update hidden states.

3. Return The final generated sequence.

---

Next, we present additional ablation results to address your questions.

Substitutes the RNN with a single-layer decoder.

Thank you for the suggestion. We have included this analysis and provided the results to assess the impact of the RNN component on ReDrafter's performance on MT-Bench when temperature is zero. Our findings show that the RNN is more effective at accurately predicting the next few tokens compared to a one-layer decoder:

Inference Speed-up VS. AR ↑ | Tokens/Step ↑

Compare dynamic TA, static TA, and without TA.

By "static tree attention," we refer to the draft model speculating a token tree without context or dependency consideration. Both Medusa and EAGLE use this approach, selecting high-probability tree tokens offline instead of using beam search during runtime. Our comparison with them highlights the benefits of dynamic tree attention, but we're happy to provide additional results by implementing static tree attention in ReDrafter.

Inference Speed-up VS. AR ↑ | Tokens/Step ↑

In the paper, we compare experimental results with and without tree attention (see Figure 6). Tree attention reduces memory and compute usage by pruning duplicate tokens, enabling a larger number of requests to be served with limited resources. It does not affect tokens per step. The best way to evaluate the difference is by measuring throughput. In our experimental setup, we use batch_size * tokens/second as a proxy for throughput.

Max Batch Size x Tokens/Step ↑

Thanks for your comments.

Compute resource/time requirements to train the draft model.

There are two steps to train a draft model: (i) training (ii) (optional) data distillation. We analyze their cost separately.

As our draft model is attached to the LLM, training can be done more efficiently conditioning on LLM's hidden states. In that case, we can skip pretraining stage of the draft model and directly start with SFT without losing accuracy. The time cost of training draft models on ShareGPT dataset for 2 epochs based on Vicuna 7B, 13B, 33B with 8 H100 GPUs is as follows. In practice, ReDrafter-V-7B converges after training for 2 epochs. ReDrafter for larger draft model requires less training epochs.

We perform data distillation as an optional preliminary step before model training. The ShareGPT dataset is divided into 121 shards and processed using 8 H100 GPUs. The costs are outlined below. Since data distillation for each LLM is a one-time offline process, users can weigh the benefits (see Table 6 in the paper) against the costs to decide whether to use it.

Smaller LLMs as standalone draft models.

We used Vicuna-7B as the draft model for Vicuna-33B and compare with ReDrafter's result as follows. ReDrafter with Vicuna-33B achieves higher accepted tokens per step while using only 1B parameters, highlighting its superior parameter efficiency.

Accept Tokens Per Step ↑

Open source code.

We commit to open-sourcing our code upon acceptance. In the meantime, we provide a preliminary zipped code package (see supplementary material) to help reproduce our results while maintaining anonymity.

Why only comparing against EAGLE and Medusa?

We recognize the broad range of research related to ReDrafter. For our comparisons, we chose EAGLE and Medusa as both utilize a draft-head-attached-to-LLM framework, allowing us to highlight ReDrafter's key strengths: (i) superior draft model design (compared to Medusa) and (ii) dynamic tree attention (compared to EAGLE). Furthermore, we extended our comparisons to include Sequoia and EAGLE2. ReDrafter achieves significant speedups over Sequoia on Vicuna-33B (2.61x vs. 1.82x) using the MT-Bench dataset and outperforms EAGLE2 in speedup performance at a temperature of 1 on Vicuna-7B (3.50x vs. 2.74x) and 13B (3.51x vs. 2.86x). For additional details, please refer to our responses to Reviewers zZts and Uws2.

Below, we present additional ablation studies on ReDrafter. We compare it against standalone methods to highlight the effectiveness of incorporating attached draft models. Additionally, we demonstrate the advantages of utilizing a tree structure for speculative tokens over a linear chain. We also include a comparison with recently proposed speculative decoding methods.

ReDrafter vs. standalone draft models.

We appreciate your interest in comparing draft models attached to LLMs with standalone draft models. Since the training data for Llama 3.2 is not open-sourced, we were unable to train a ReDrafter for comparison. Instead, we used Vicuna-7B as the draft model for Vicuna-33B, achieving 3.71 accepted tokens per step. In comparison, ReDrafter with Vicuna-33B achieves 3.87 accepted tokens per step while using only 1B parameters, highlighting its superior parameter efficiency.

We must note that using standalone draft models results in restrictions and cost that cannot be accounted for in the quantitative results, such as:

1. A standalone draft model must share the same tokenizer as the LLM.
2. Standalone draft models often require intricate alignment of pre-filling and KV-cache, complicating deployment.

ReDrafter avoids these issues, ensuring broader compatibility. Also, ReDrafter can quickly adapt to the LLM during training. It does not need to train from scratch. In our experiments we only train the ReDrafter on ShareGPT for a few epochs, which only cost less than 5 hours on 8 H100 GPUs.

Using a chain of speculative tokens.

When ReDrafter operates with a beam width of 1, it simplifies to a chain of speculative tokens. In this case, ReDrafter-V-7B achieves a 1.79x speedup on MT-Bench at t=0, compared to the 2.80x speedup observed with beam search. This highlights the effectiveness of allocating additional compute during inference time through beam search.

Compare with Sequoia and Eagle2.

Sequoia [2] and Eagle2 [3] authors provide their codebase and checkpoints in their repositories, allowing us to reproduce some of their results. We made adjustments to align the code with ReDrafter's experimental setup, as detailed below.

• Sequoia experiments: Sequoia reports results on Vicuna 33B. Utilizing the authors' codebase, we adapted it to the Mt-Bench dataset and executed their implementation on a single H100 GPU in offloading mode. Without offloading, we encountered an out-of-memory (OOM) issue.

  Inference Speed-up VS. AR ↑ | Tokens/Step ↑

  Model | Temperature	t=0	t=1
  Sequoia-V-33B	1.82x | 1.84	1.45x | 1.46
  ReDrafter-V-33B	2.61x | 3.87	3.27x | 4.69

• EAGLE2 experiments: EAGLE2 authors reports results on Vicuna 7B and 13B in their preliminary draft. Below, we reproduce these results in our experimental setup on Mt-Bench using a single H100 GPU.

  Inference Speed-up VS. AR ↑ | Tokens/Step ↑

  Model | Temperature	t=0	t=1
  EAGLE2-V-7B	3.00x | 4.63	2.74x | 4.10
  ReDrafter-V-7B	2.80x | 4.20	3.50x | 5.31
  EAGLE2-V-13B	3.13x | 4.61	2.86x | 4.18
  ReDrafter-V-13B	2.80x | 4.21	3.51x | 5.29

  In the results above, ReDrafter outperforms EAGLE2 when the temperature is set to 1, demonstrating the effectiveness of beam search in handling uncertainty within the distribution. Conversely, EAGLE2 performs better when the temperature is set to 0. This is because EAGLE2 employs a greedy search strategy to construct the token tree and incorporates additional acceleration optimizations specifically for zero-temperature scenarios.

Equation in line 156.

We concatenate $h$ on $s_t$ as $g_t$ to forecast token at position $t$. This is a special case of RNN with $h_t=h$, that is, an identity transition matrix on $h$. This could help enhance longer range forcasting and get rid of short-sighted errors. See [4] for more intuitions.

[2] Z. Chen et al. Sequoia: Scalable and Robust Speculative Decoding. NeurIPS 2024.
[3] Y. Li et al. EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees.
[4] G. Bachmann and V. Nagarajan. The Pitfalls of Next-Token Prediction. ICML 2024.

Thanks for your comments.

Below, we use an example to thoroughly explain the theoretical guarantee of ReDrafter in producing the same distribution as the LLM. Regarding your question:

Why is the output distribution unchanged?

This is guaranteed through the verification stage in speculative decoding, either with greedy decoding (when t=0) or speculative sampling (when t>0).

• A greedy decode example: if the drafter proposes I like tea and the LLM outputs I am bob, only the first token I is accepted to ensure a match.

• Speculative sampling: When we sample tokens from the drafter, the beam search method returns their log probabilities log Q(v), in addition to v. To verify v, we pass the token before v to the LLM, which returns log P(v). Given v, log Q(v), and log P(v), we follow the routine:

  • Step 1: accept v with probability min(1,P(v)/Q(v)).
  • Step 2.1: if v is accepted, continue to the next token proposed by the drafter.
  • Step 2.2: if v is rejected, sample v from an un-normalized distribution R(v)=max(0, P(v)-Q(v)).

  Consider a simple vocabulary with only two tokens: A and B. Assume that Q(A)=Q(B)=1/2, P(A)=1/4, and P(B)=3/4.

                                P(B)=3/4
                                ▄▄
  Q(A)=1/2             Q(B)=1/2 ██
       ┌─┐ P(A)=1/4         ┌─┐ ██
       │ │ ▄▄               │ │ ██
       └─┘ ▀▀               └─┘ ▀▀
  ----------------------------------------
        A                     B
  

  Assume that sampling from Q(v) produces a token A. Since the true probability of A, P(A), is lower than Q(A), we accept A with a probability of P(A)/Q(A)=1/2​ and reject it with a probability of 1-P(A)/Q(A)=1/2. Now, assume that sampling from Q(v) produces a token B. Since the true probability P(B) is higher than Q(B), we always accept B. This procedure results in the following outcomes:

  • Acceptance of v=A with probability 1/4.
  • Acceptance of v=B with probability 1/2.
  • Rejection with probability 1/4.

  If the token is rejected, we resample from a tilted, unnormalized distribution R(v)=max(0,P(v)−Q(v)).

                    R(B)=1/4   
                    ▄▄        
        R(A)=0      ▀▀         
  -------------------------
        A           B
  

  In this case, R(A)=0, R(B)=1/4, meaning sampling from R yields B almost surely. The overall probably of selecting B is 1/2+1/4=3/4, which matches the LLM's probability.

  Append A.1 of the paper at [1] shows that the above intuition is exactly what we need to do to alter drafter samples such that they look like LLM samples.

[1] Y. Leviathan et al. Fast Inference from Transformers via Speculative Decoding. ICML 2023.

I thank the authors for their thoughtful responses.

I remain concerned about the correctness of the proposed algorithm, during stochastic decoding. My concern is that in beam search, even if we know the true probabilities (according to the draft model) of all of the tokens in the beam, these tokens were not actually sampled with those probabilities from the draft model (because beam search is a deterministic algorithm). Knowing the probabilities from which the drafted tokens were actually sampled is necessary to guarantee that the output distribution of the target model is unchanged by the speculative decoding process. Therefore, I do not believe that the proposed algorithm is correct.

As a result of this, for now I have lowered my score to a 3.