DySpec: Faster Speculative Decoding with Dynamic Token Tree Structure

Authors: Yunfan Xiong, Ruoyu Zhang, Yanzeng Li, Tianhao Wu, Lei Zou
Primary area: infrastructure, software libraries, hardware, systems, etc.
Keywords: inference methods, efficient inference, speculative decoding

Abstract:
While speculative decoding has recently appeared as a promising direction for accelerating the inference of large language models (LLMs), the speedup and scalability are strongly bounded by the token acceptance rate.
Prevalent methods usually organize predicted tokens as independent chains or fixed token trees, which fails to generalize to diverse query distributions. 
In this paper, we propose \textsc{DySpec}, a faster speculative decoding algorithm with a novel dynamic token tree structure. 
We begin by bridging the draft distribution and acceptance rate from 
intuitive and empirical clues, and successfully show that the two variables are strongly correlated. Based on this, we employ a greedy strategy to dynamically expand the token tree at run time. Theoretically, we show that our method can achieve optimal results under mild assumptions. Empirically, \textsc{DySpec} yields a higher acceptance rate and speedup than fixed trees. \textsc{DySpec} can drastically improve the throughput and reduce the latency of token generation across various data distribution and model sizes, which significantly outperforms strong competitors, including Specinfer and Sequoia. Under low temperature setting, \textsc{DySpec} can improve the throughput up to 9.10x and reduce the latency up to 9.4x on Llama2-70B. Under high temperature setting, \textsc{DySpec} can also improve the throughput up to 6.21x, despite the increasing difficulty of speculating more than one token per step for draft model.

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

Reviews and discussion:
<thread>
creation date:
1730053100074

thread_type:
official_review

<thread_object>
creator: reviewer_jh6n

<content>
summary:
This paper proposes a draft token selection method in speculative decoding aimed at improving the token acceptance rate. The core idea is to use the draft model's prediction score as evidence to infer the token acceptance rate and use this information to select more promising tokens. The authors provide both theoretical and empirical analyses to support their approach.

strengths:
The proposed token selection method demonstrates improvement over existing methods such as Sequoia and Specinfer.

weaknesses:
1. The paper lacks discussion on existing works that share very similar ideas (see Questions 1 and 2). 
1. The experimental results do not sufficiently validate the method's effectiveness from various perspectives (see Questions 3 and 4).

questions:
1. There are concurrent works using similar ideas. For instance, EAGLE 2 (https://arxiv.org/abs/2406.16858) attempts to maximize the acceptance rate based on a similar intuition (see Figure 6 in their paper). The authors should mention these works. What are the differences between your method and EAGLE 2?
1. Dynamic tree-based selection methods have appeared in prior literature. For example, Recurrent Drafter (https://arxiv.org/abs/2403.09919) proposes an RNN draft model and uses beam search for adaptive token selection. Did you modify the draft model architecture in your approach?
1. In practice, small draft models may be overconfident about incorrect answers in reasoning tasks. Have you studied your Hypothesis 1 across different tasks and analyzed the failure examples?
1. Does your method support batch inference? Have you experimented with large batch sizes to push your algorithm to its computational limits and test its throughput?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730390112776

thread_type:
official_review

<thread_object>
creator: reviewer_dmmc

<content>
summary:
Summary:

The authors propose dyspec to generate an optimal tree. The core idea is based on the idea that the target distribution and draft distribution should match.

The authors based on the idea propose an algorithm to come up with optimal tree. They overcome several implementation challenges and achieve reasonable speedups compared to existing baselines.

“DYSPEC achieves a 9.1× throughput improvement and 9.4× reduction in latency.” ([“DySpec”, p. 2]-> Should compare to SpecInfr or Medusa

“Figure 2: Connection between acceptance rate/target distribution and draft distribution on CNN DailyMail.The density of each block is normalized by column.” ([“DySpec”, p. 3] -> Can you please explain this figure. I think it’s the main motivation for your method, however it is not clearly explained, what this figure is showing.

“Specinfer-Baseline” ([“DySpec”, p. 7] -> Specifinfer requires training the models. What have authors done here.

“DYSPEC leverages CUDA Graph to capture 129 different input lengths ranging from 128 to 258” ([“DySpec”, p. 7-> What is the significance of CUDA Graph here, my understanding is CUDA graph is a mechanism to launch multiple kernels at the same time on the GPU to minimize the overhead of kernel dispatch

“We selected Llama2-7B as the draft model” ([“DySpec”, p. 7] -> What happens when you use the 68M model ?

“Set the maximum draft token tree size to 64, DYSPEC achieves up to a 9.1x improvement in throughput and a 9.4x reduction in latency compared to auto-regressive generation” ([“DySpec”, p. 8]-> Unfair comparison, compare to baselines

Additional experiments -

I would really like to experimentally understand the optimality of your tree. Will it be possible to construct an optimal tree based on some post-hoc data and compare it with the tree generated by Dyspec.

strengths:
- The proposed idea is quite useful and can provide significant speedups for tree based speculative decoding methods. 

- The overheads reported our negligible.

weaknesses:
- See the summary section please.

questions:
- See the summary section 

Please answer some of the question there. I am happy to bump up the score to 6.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730677441294

thread_type:
official_review

<thread_object>
creator: reviewer_sl3u

<content>
summary:
This paper introduces DySpec, a novel approach to speculative decoding that employs a dynamic token tree structure to improve inference speed for large language models. The authors present both theoretical and empirical evidence showing that higher draft probabilities correlate with higher acceptance rates. Based on this insight, they develop a greedy strategy for dynamically expanding the token tree at runtime. The method achieves impressive results, demonstrating up to 9.1× throughput improvement and 9.4× reduction in latency on Llama2-70B under low temperature settings, outperforming existing methods like Specinfer and Sequoia.

strengths:
1. Strong Theoretical Foundation
- Provides rigorous theoretical analysis linking draft distribution to acceptance rate
- Includes formal proofs of optimality under stated assumptions
- Clearly bridges theoretical insights with practical implementation

2. Novel Technical Contributions
- Introduces an innovative dynamic token tree construction approach
- Develops efficient algorithms for both fixed-size and threshold-based tree construction
- Proposes block-sparsity friendly token ordering for optimization

3. Comprehensive Empirical Evaluation
- Tests across multiple model scales (7B to 70B parameters)
- Evaluates on diverse datasets (C4, OpenWebText, CNN DailyMail)
- Compares against strong baselines (Specinfer, Sequoia)
- Examines performance under different temperature settings

4. Implementation Efficiency
- Addresses practical concerns about overhead
- Provides C++ implementation to minimize token tree construction costs
- Includes detailed analysis of computation complexity

weaknesses:
1. Limited Discussion of Limitations
- Could elaborate more on scenarios where the method might not perform optimally
- More discussion of the trade-offs between fixed-size and threshold-based approaches would be valuable

2. Implementation Details
- Some implementation specifics about the C++ optimizations could be expanded
- Could provide more guidance on threshold selection for different scenarios

3. Experimental Validation
- Could include more ablation studies to isolate the impact of different components
- Additional experiments on more diverse model architectures would strengthen generalizability claims

questions:
None
</content>
</thread_object>

</thread>
<thread>
creation date:
1730684387291

thread_type:
official_review

<thread_object>
creator: reviewer_rcvf

<content>
summary:
This work proposes a method to dynamically expand the token tree based on the draft distribution. A key difference between this work and existing approaches is that it introduces a smooth dynamic draft token tree construction method, which expands the tree along both width and depth without hyperparameter setup, within the given token budget. However, related works are missing, and the experimental setup should be properly improved.

strengths:
1.	The correlation of Hypothesis 1 with the proposed method is well elaborated and explained.
2.	The structure of the paper is clear and easy to follow.

weaknesses:
1.	Lack of related work. Context-aware dynamic draft token tree is not a new idea. I would like to draw your attention to a very related work: “EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees” (EMNLP'24). This paper also proposes adopting dynamic draft token trees and claims a strong positive correlation between the draft model confidence score and the acceptance rate of the token. Another relevant paper is “Dynamic Depth Decoding: Faster Speculative Decoding for LLMs,” which further improves performance by dynamic depth. The methodology of EAGLE-2 is quite similar to this paper but differs in the tree construction method. Including this work in your paper and providing necessary discussion is essential.
2.	Experimental setup and results are weak. First, the dataset is too limited. I recommend adding more datasets like MT-Bench (Zheng et al., 2023), HumanEval (Chen et al., 2021), and GSM8K (Cobbe et al., 2021). Second, the 9.1× speedup is compared with the autoregressive method, and when compared with static tree methods like Sequoia, the results are much less overwhelming. Considering that the experimental setup lacks comparison with state-of-the-art dynamic draft token tree methods, the experiment does not fully convince me of the merits of this methodology.
3.	Writing needs improvement. Section 4.2 is hard to follow. The presentation of Figure 3 should be improved. Additionally, Figure 4 is not clear; please increase the font size.

questions:
More convincing experimental results that compare with more relevant related work.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732192647124

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response to Reviewer and AC

comment:
We thank all reviewers and AC for their efforts in reviewing and valuable feedback. We notice that many reviewers raise concerns about DySpec experiment coverage (rcVF, SL3U, Jh6N) and differences with some particular tree-based speculative decoding methods, e.g. EAGLE-2 (rcVF, Jh6N). Below we will address these common concerns and we hope that reviewers find our response well-clarified and helpful.

### **1. Comparison with EAGLE-2.**

As a concurrent work, the main difference between EAGLE-2 and DySpec (ours) is that
- **Feature-based vs. Model-based:** EAGLE-2 is a **self-speculative** method that makes draft predictions based on the target model's features, rather than a much smaller draft model. Due to the strong drafting capability, self-speculative methods (Medusa, EAGLE, and EAGLE-2) can usually guess with higher accuracy under the same budget. 
In this work, we mainly focus on **how to build near-optimal draft trees**
with a given target and draft, which is orthogonal to **using stronger draft models**. This is also why we do not compare self-speculative methods in our experiments section. Nevertheless, we would like to extend the DySpec tree construction method into self-speculative fashions in future works.
- **Tree construction:** EAGLE-2 builds their draft trees with an expand-rerank procedure: first **selects** top-k tokens at each node, and prunes the candidate tree with draft probability. While DySpec **samples** tokens with weights according to Eq (2), which is essentially the probability of **sampling with replacement**. In other words, DySpec **maximizes the sum of expected draft probabilities of all candidate nodes** with a given budget. We also implement a version of **tree construction with threshold and early exit** to accelerate sampling from the draft model by reducing sampling times, as shown by Algorithm 2.
- **Verification:** One following problem with EAGLE-2 is that it cannot accept tokens with standard verification, i.e. only reject the draft with probability $1 - \frac{target}{draft}$ when $draft > target$, since draft tokens are predicted by **selection rather than sampling**. Instead, EAGLE-2 is verified by sampling a new token from the target at each draft node, and seeing whether it matches the draft token. The problem here is that even in the case that **draft probability is identical to target probability**, the latter verification may yield a low acceptance rate. Imagine a distribution of $[0.25, 0.25, 0.25, 0.25]$ over a vocabulary of $4$. No matter how the draft model guesses, there is only a probability of $0.25$ that it aligns with the target. In contrast, the sampling method adopted by DySpec can yield a perfect acceptance rate of $1$ in this case.

### **2. Experiments with more settings.**

We totally understand the reviewer's concerns about experiment setups. We add experiment results on more corpus to thoroughly show the performance comparison between DySpec and baselines.

- **68m/7b**: **Speedup** to auto-regressive / **Accepted tokens** with a budget of 64

    | Dataset   | Temperature   | DySpec   | Sequoia | SpecInfer     |
    | :-:       | :-:           | :-:             | :-:           | :-:           |
    | GSM8K     | 0.0           | **3.93** / **6.86**   | 2.79 / 4.92   | 2.01 / 3.47   |
    | GSM8K     | 0.6           | **2.44** / **4.31**   | 2.20 / 3.55   | 1.65 / 3.03   |
    | MT-Bench  | 0.0           | **2.59** / **4.02**   | 2.35 / 3.55   | 1.68 / 2.70   |
    | MT-Bench  | 0.6           | **2.15** / **3.62**   | 2.11 / 3.18   | 1.50 / 2.71   |
- **68m/13b**: **Speedup** / **Accepted tokens**

    | Dataset   | Temperature   | DySpec                | Sequoia       | SpecInfer     |
    | :-:       | :-:           | :-:        | :-:      | :-:    |
    | GSM8K     | 0.0           | **3.17** / **5.29**   | 2.21 / 3.92   | 1.74 / 2.98   |
    | GSM8K     | 0.6           | **2.49** / **4.17**   | 2.10 / 3.39   | 1.51 / 2.72   |
    | MT-Bench  | 0.0           | **2.19** / **3.72**   | 2.15 / 3.46   | 1.50 / 2.62   |
    | MT-Bench  | 0.6           | **2.25** / **3.62**   | 1.93 / 3.11   | 1.40 / 2.60   |
- **7b/70b**: **Speedup** / **Accepted tokens**

    | Dataset   | Temperature   | DySpec    | Sequoia       | SpecInfer     |
    | :-:       | :-:           | :-:                   | :-:       | :-:     |
    | GSM8K     | 0.0           | **10.56** / **12.39** | 7.31 / 7.62   | 5.22 / 5.34   |
    | GSM8K     | 0.6           | **7.57** / **8.14**   | 6.62 / 6.89   | 5.75 / 5.89   |
    | MT-Bench  | 0.0           | **9.95** / **11.25**  | 6.96 / 7.46   | 4.75 / 4.85   |
    | MT-Bench  | 0.6           | **8.47** / **10.11**  | 6.96 / 7.46   | 5.52 / 5.67   |

In summary, DySpec consistently outperforms compared baselines, Sequoia and SpecInfer, under both greedy and non-greedy settings across different draft/target sizes. Compared with Sequoia, the relative speedup can be up to 45% and the average acceptance rate can be up to 63%.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732376289265

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Looking forward to the discussion

comment:
Dear Reviewers,

We sincerely appreciate your time and effort in reviewing our work. We fully understand your schedule may be quite busy right now. As the deadline for the Author-Reviewer discussion period is approaching, we would greatly value the opportunity to engage in further discussion with you. We look forward to your feedback on whether our responses effectively address your concerns and if there are any additional questions or points you would like to discuss.

Thank you again for your thoughtful consideration.

Best regards,
Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper provides a dynamic tree structure construction for multi-draft speculative decoding, and shows improvements over existing static tree construction techniques. However, during the rebuttal, the reviewers mentioned the existence of related work that also creates trees dynamically, specifically Eagle 2, which has significant overlap with the current proposal. The authors have responded with an explanation on the differences. However, the AC finds this to be insufficient and a more in-depth understanding of similarities and differences is warranted. In particular, the AC recommends that the authors replace the tree construction step in Eagle 2 with that proposed herein. This ablation is key to substantiating the claim of almost optimality of the tree construction in this work and can increase its impact further. We hope that the authors find the comments of the reviewers useful for a future iteration of their paper.