DySpec: Faster Speculative Decoding with Dynamic Token Tree Structure

We thank the reviewer for his efforts and insightful feedback. Below, we will try to address your questions carefully. We hope the reviewer will consider raising your score in light of our response.

Q1: Regarding related works (EAGLE-2 and ReDrafter).

Thank you for bringing these methods to our attention. We will add one additional related works section (which was not included previously due to page restriction) to discuss the difference between these methods in detail.

1. 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 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 by 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. We also carefully review the tree-based selection method in ReDrafter. We would like to point out that ReDrafter is a self-speculative method that uses independent sampling in tree construction (similar to SpecInfer). As in DySpec, we use a much smaller draft model rather than target model features to make the predictions. Again, we want to emphasize that our paper concentrates on building and implementing near-optimal draft trees with a given target and draft, and we are open to leaving the combination with self-speculation as future works.

In summary, we thank the reviewer for raising these complementary-related works. We will include the discussions in our final version.

Q2: The effectiveness of Hypothesis 1 under different domains.

We totally understand your concerns about the generalization of DySpec. We would like to point out that the effectiveness of DySpec builds on the similar nature of draft distribution and target distribution: it turns out to work fine both theoretically and empirically when the two distributions are "close", according to the discussion in our paper. This is also the case where speculative decoding itself works.

Besides, we also conduct more experiments under practical corpus, i.e. GSM8K (math) and MT-Bench (general chat) to further show the effectiveness of DySpec. The detailed results can be found in our response to Reviewer rcVF, Q2. DySpec consistently outperforms compared baselines, Sequoia and SpecInfer, under both greedy and non-greedy settings across different draft/target sizes. Particularly, under the memory-bound 7b/70b setup, DySpec can hit 50%-60% more draft tokens with greedy decoding and 20%-35% with non-greedy decoding. In terms of speedup, the performance gain is around 45% and 15%-20%, respectively. We believe that such strong performance across various domains and settings is able to demonstrate the effectiveness of DySpec's tree construction.

Q3: Support for various batch sizes.

Thank you for this question. We agree with you that tree-based speculative decoding methods give more gain for small batch sizes, where the generation of the target model is more memory-bound, thus verifying a larger tree incurs less overhead relatively. This conclusion also aligns with the result that the 7b/70b setting achieves the highest speedup. For larger batch sizes, one can afford less budget for speculation, thus the relative gain of optimal tree structures compared to baselines (e.g. chains) gets smaller. However, it is still possible to use DySpec in real time to push the performance limits at larger batch sizes.

Dear Reviewer,

Thank you for handling our manuscript and providing valuable feedback. We hope that our responses have sufficiently addressed the concerns you raised. We welcome more discussion if you have more questions and suggestions. As the discussion deadline is approaching, we would be very grateful if you could take a moment to review our reply.

Thank you very much for your thoughtful review. We are glad that the reviewer appreciates our motivation, careful implementation, and significant speedups. Below, we will try to address your questions carefully. We hope the reviewer will consider raising your score in light of our response.

Q1: Inaccuracy in expressions.

Thank you for pointing out. We will express more precisely the competitor when comparing the improvement in our abstract and introduction. To clarify, we achieve up to 50% speedup in the 7b/70b setting compared with Sequoia, a strong baseline.

Q2: More explanation in Figure 2.

Section 3 mainly elaborates on our hypothesis of the strong correlation between draft probability and expected acceptance rate. The results in Figure 2 act as empirical evidence for the hypothesis: tokens with higher draft probability usually have higher target probability as well as higher acceptance rate.

Specifically, each number in the grid map represents the density w.r.t. column, which is draft probability in both cases. For instance, the top right number $0.95$ in Figure 2a shows that 95% of tokens with draft probability among $(0.8, 1.0]$ also have an acceptance rate among $(0.8, 1.0]$. We believe that such observation provides a strong empirical guarantee for Hypothesis 1 and motivates our algorithm design.

Q3: Training draft model for SpecInfer.

In the original paper, SpecInfer trains JF68m to close the gap between the draft model and the target model (Llama). We follow their setting and use identical draft/target pairs, which are also adopted by Sequoia. Therefore no retraining is needed in this case.

Q4: The significance of CUDA graphs.

In DySpec implementation, we adopt CUDA graphs to accelerate draft model computation by eliminating CPU overheads, which is particularly crucial under the 68m/7b setting. To minimize the latency of the draft model, we capture a graph for each input length, which requires 128 captures to support generating 128 tokens in our experimental setup.

When the draft model's latency is already sufficiently smaller than the target model's (i.e. the 7b/70b setting), the overhead of the draft model becomes less significant. Consequently, even without CUDA graph acceleration, a considerable end-to-end speedup ratio can still be achieved.

Thank you for raising this concern. We will clarify more on the usage of the CUDA graph in our final revision.

Q5: 7b vs. 68m on 70b target model.

Following our discussion in Q4, since the 70b target model has a sufficiently larger latency compared with the 7b draft model (as shown by Figure 4c), further replacing it with a 68m size becomes less significant. In fact, the size of the draft model is a "latency-accuracy" trade-off. In this case, the 7b model yields a much higher acceptance rate at the cost of a slightly higher draft/target latency ratio, overall resulting in better performance.

Q6: About optimality.

Thank you for your question. We would like to explain this problem in detail.

First, we want to clarify that DySpec and Sequoia employ a sampling with replacement in tree construction, while EAGLE-2 uses selection and pruning. During verification, one following problem of EAGLE-2 is that it cannot accept tokens by 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 approach 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 yields a perfect acceptance rate of $1$ in this case.

Back to the cases of DySpec and Sequoia. Since the trees are built by sampling, it is impossible to find a fixed optimal tree in advance. We can only compute and compare the expected accepted tokens once the tree is sampled and determined. Nevertheless, we present a case study on initial guesses of DySpec and Sequoia in Figure. It shows a vocabulary space with depth and width both set to 3. Nodes chosen by DySpec are colored in blue and Sequoia are colored in orange. As shown, at the divergent step, DySpec will not sample from the node "com" with an acceptance rate of $0$ (since $target=0$) due to its low weight (draft probability). Instead, it samples the node "The" with an acceptance rate of 1 (since $draft < target$). In other words, DySpec can better escape from the local minimum and fully utilize the speculation budget.

Dear Reviewer,

Thank you for handling our manuscript and providing valuable feedback. We hope that our responses have sufficiently addressed the concerns you raised. We welcome more discussion if you have more questions and suggestions. As the discussion deadline is approaching, we would be very grateful if you could take a moment to review our reply.

Thank you very much for your thoughtful review. We are glad the reviewer found our work novel, practical, and theoretically and empirically sound.
Below, we will try to address your questions carefully. We hope the reviewer will consider raising your score in light of our response.

Q1: Discussion of limitations.

Thank you for this question. In short, DySpec has the same bottleneck as most speculative methods:

• The performance is bounded by how similar draft and target distributions are. In particular, an accurate draft can build a near-optimal draft tree at the tree construction stage and accept more tokens at the verification stage.
• The method has a larger gain in memory-bound scenarios. As shown, the 7b/70b setting has the highest speedup and drops as the generation turns to computation-bound. The reason is that verifying a larger tree incurs less overhead relatively in memory-bound scenarios, and the target model benefits more from near-optimal draft trees provided by DySpec. Nevertheless, we would like to point out that other prevalent speculative decoding methods, e.g. Sequoia and SpecInfer, also share the same trend.

Q2: Trade-offs between fix-size and threshold-based implementations.

Thank you for this question. Given the optimization objective of accepted tokens, it is clear that the fix-size algorithm, which expands the tree greedily, yields the best returns. However, in implementation, the fix-size algorithm may incur many calls of draft model generation, which is inefficient (since we do not know which leaf node to expand from each time). Therefore, we propose the threshold-based algorithm to pack the calls of the draft model into batches to reduce the number of passes. We also incorporate an early-exit mechanism to further improve efficiency: when the weight of every leaf node (i.e. the expected accepted probabilities) is below the threshold, we stop generating new tokens and verify the current tree with the target model. In short, the threshold-based algorithm is a practical approximation of the fixed-sized algorithm with higher efficiency and less execution time in the real world.

Q3: More details in C++ optimizations.

The overhead introduced by DySpec can be primarily attributed to draft tree construction. Even in the faster threshold-based algorithm, the tree expansion remains time-consuming due to its complex logic. To mitigate this, we maintain the tree structure with C++ implementation, which includes adding new nodes, generating input positions, and creating tree attention masks (here we also use PyTorch's C++ API to optimize the retrieval of parents' indices). As a result, we finally reduced the tree construction time by more than 50% compared to vanilla Python implementation.

Q4: Experimental validations.

We understand the reviewer's concerns about the effectiveness of our different components in implementation, and whether it can generalize to other model architectures.

• For the first question, we would like to point out that our dynamic tree algorithm is rather monolithic, and we walk through various optimization techniques to eliminate and overlap the additional cost induced by building the dynamic tree, which serves the same purpose. Therefore, we do not conduct much performance breakdown for these optimization tricks.
• For the latter question, we would like to kindly remind that we have included various range of model sizes (68m/7b, 68m/13b, 7b/70b) in our experiments, and this paper mainly focuses on accelerating the generation of transformer models. Extending to other model architectures (e.g. Mamba) is interesting but beyond the scope. In cases where the reviewer has concerns about generalization to other domains, we also add additional experiments on GSM8K and MT-Bench, two common evaluation benchmarks, which can be found in our response to Reviewer rcVF, Q2. We hope that the reviewer finds these results helpful.

Dear Reviewer,

Thank you for handling our manuscript and providing valuable feedback. We hope that our responses have sufficiently addressed the concerns you raised. We welcome more discussion if you have more questions and suggestions. As the discussion deadline is approaching, we would be very grateful if you could take a moment to review our reply.

Thank you very much for your thoughtful review. We are glad that the reviewer found our work well motivated and easy to follow.
Below, we will try to address your questions carefully. We hope the reviewer will consider raising your score in light of our response.

Q1: Regarding related works (EAGLE-2 and Dynamic Depth Decoding).

Thank you for bringing these methods to our attention! We will add one additional related works section (which was not included previously due to page restriction) to discuss the difference between these methods in detail.

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.

Q2: Extending experiment 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.

(Note: HumanEval is not included here because the draft model, JF68m, is only trained on Wikipedia and C4-en, with no coding corpus.)

• 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%.

As we mentioned in Q1, we do not compare with EAGLE-2 due to the unfair nature of self-speculative methods. The combination of DySpec's tree construction and self-speculative methods is beyond the scope of this paper, and we are looking forward to addressing it in future works.

Q3: Writing issues.

Thank you for pointing out our writing issues. We will polish our algorithm section ($\S$ 4.2) and refine the presentation of Figure 3 and Figure 4. We believe that we will present a better final version through revision.

Dear Reviewer,

Thank you for handling our manuscript and providing valuable feedback. We hope that our responses have sufficiently addressed the concerns you raised. We welcome more discussion if you have more questions and suggestions. As the discussion deadline is approaching, we would be very grateful if you could take a moment to review our reply.