Gumiho: A Hybrid Architecture to Prioritize Early Tokens in Speculative Decoding

Thank you for your suggestions and support.

Q1. There is a lack of comparison to the latest state-of-the-art papers, which propose better results than Eagle2, such as HASS [1] and DDD [2]

A1:(1) Hass improves Eagle2 by addressing the training-inference inconsistency in serial transformer heads within SPD. While Hass focuses on training methodology, our work enhances performance via architectural modifications. These two directions are orthogonal: our method could integrate Hass’s training improvements for further gains.

In the table below, we compare the performance of our Gumiho with Hass. Additionally, we evaluate a variant of our model (denoted as Gumiho w/ Hass ) where Hass’s training method is applied to our Transformer Head. All results are conducted on a Mi250 GPU using the LLaMA3-8B model with a temperature of 0.

(2) DDD improves Eagle2 by using a dynamic draft tree. Since DDD’s code is not open-sourced, we compare relative improvements over Eagle2. On the MT-Bench dataset, the performance gains of DDD and Gumiho are:

The aforementioned results and discussions will be included in the final version of our paper.

Q2. The method uses more parameters in the draft model compared to Eagle, Eagle2, and Medusa. Does the training process consume more GPU memory?

A2: Yes, you are correct. Our method does require more GPU memory during training compared to Eagle, Eagle2, and Medusa. Here’s why:

• Eagle and Eagle2 use one single-layer Transformer head and reuse it serially across multiple speculative steps.
• Medusa only maintains five parallel MLP heads.
• Our approach, however, employs a two-layer Transformer head alongside five parallel MLP heads, all trained simultaneously.

This architectural choice increases memory consumption. However, it enables a hybrid generation strategy: serial decoding for the first few tokens (to ensure accuracy) and parallel decoding for subsequent tokens (to maximize speed). While this introduces a memory-accuracy trade-off, it achieves a significant speedup improvement. We will explicitly discuss this trade-off in the final manuscript.

Q3. Is it able to combine the proposed method with the state-of-the-art Spd method to further improve the speed up ratio? It is better if the conclusion of the theorem can be applied to other methods.

A3: As illustrated in A1, our method could integrate Hass’s training strategy for further gains.

In addition to Hass, our core theoretical principle—prioritizing the initial tokens—can also be generalized to serial multi-head draft models (e.g., Hydra or Eagle). By reallocating Head parameters to emphasize early tokens, this principle could further enhance the speedup ratio of existing SOTA methods.

Moreover, existing methods typically use a single smaller model (e.g., LLaMA-3-3B) to accelerate larger counterparts like LLaMA-3-405B. However, our analysis reveals that adopting a differentiated draft model strategy—where a more capable model (e.g., LLaMA-3-8B) predict initial tokens, while a faster but lighter model (e.g., LLaMA-3-1B) handles subsequent tokens—could optimize the quality-speed trade-off beyond current homogeneous settings.

Exploring these extensions is a key direction for our future research.

Q4. What is the average draft time for each part of the method? In figure3 you show the overall draft time including the model forward time and full tree attention time. I would like to see each part of them.

A4: The following table presents the draft time of the Vicuna 7B on Mi250 GPU.

• Serial Head Forward represents the time for a single serial head forward pass once$\times$the number of forward passes. In Eagle2, a single serial head corresponds to a single-layer Transformer, while in Gumiho it represents a two-layer Transformer.
• Parallel Head Forward denotes the time for a parallel head forward pass once. Eagle2 does not have parallel heads, whereas in Gumiho, this refers to the parallel MLP heads.
• Tree Attention / Full Tree Attention represents the time required to construct the tree structure in Eagle2 and Gumiho.
• Additional Computations represents the time for additional computing. (e.g., using torch.top_k to retrieve current tokens).
• Total Time means the total time consumed for a draft generation.

Thank you for your suggestions and support.

Q1. The ablation study shows that using FTA has marginal improvement.

A1: The improvement contributed by our proposed FTA method is non-negligible in the context of the overall performance gain. In the FTA ablation study, the speedup ratio increases by 0.05, while the total improvement of our method over Eagle2 is 0.27. This means FTA accounts for approximately 18.5% of the total performance gain, thereby demonstrating its critical role in our approach.

Moreover, FTA introduces a lossless improvement compared to existing TA without incurring additional computational overhead. It only requires modifying the Tree Mask during verification by setting specific positions from 0 to 1 (i.e., unmasking certain tokens). This design ensures efficiency while enhancing performance.

Q2. The conclusion of the theorem is simple but it seems that the proof is rather complicated. Is there an easier way to give the proof?

A2: Note that although the conclusion of the theorem appears simple, the situation in reality is rather complicated since we have only imposed very few restrictions on the Theorem. Specifically, we only require that the sum of the increments in the former part and the decrements in the latter part must be equal. The probabilities $\lbrace p_i\rbrace_{i=1}^D$ in the original setting, the relative magnitude of increases and decreases $\zeta_i$, the length of the original sequence $D$, and the divided location $d$ of the former part and latter part are all unbounded and can change freely.

Therefore, although the theorem itself appears simple, the high degree of freedom results in highly complex scenarios requiring exhaustive case analysis. Under the current assumptions adopted in our paper, each step in the proof is indeed necessary.

Q3. The theorem assumes that the total amount of decrease to the second part is the same as that of increase to the first part of the sequence. However, this is hardly the case in reality. Is it able to give a more general assumption?

A3: Actually, the assumption that "the increment in the former part equals the decrement in the latter part" can be interpreted as an upper-bound condition. In fact, our conclusion holds not only when these quantities are equal, but also when the increment in the former part exceeds the decrement in the latter part, in which the Eq.6 can be modified from

$\sum_{i=1}^d \zeta_i=\sum_{j=d+1}^D \zeta_j$

to

$\sum_{i=1}^d \zeta_i\geq\sum_{j=d+1}^D \zeta_j$.

This is very easy to prove that Theorem 3.1 still holds in this situation if we remove the excess portion in the former part without decreasing the mean accepted tokens.

As for the scenario where the increment is less than the decrement, we believe that this is a more complex situation and we will investigate it in our future research.

The aforementioned analysis will be added in the final version of our paper.

Q4. I wonder if the idea of ‘the prior tokens in speculative decoding are more important’ has been investigated by other papers since the conclusion is obvious and easy to derive. If so, the novelty of this paper will be decreased, and the author should give a further comparison to other papers.

A4: To the best of our knowledge, no prior work has proposed the idea presented in our paper. We are the first to introduce the idea that "prior tokens are more important" and provide rigorous theoretical proof for this concept. Furthermore, we leverage this insight to design a novel model architecture, Gumiho.

Moreover, applying the principle of 'prior tokens are more important' to achieve speedup improvements is not straightforward. Simply enhancing prior token accuracy may increase computational overhead at the same time, which would prolong draft time and consequently undermine the overall speedup. In fact, we have no idea how this would affect the final speed-up ratio if we merely increase the computational budget on the prior tokens, since the mean accepted tokens and the draft time will increase together and they are two factors that play opposite roles.

In contrast, our Gumiho tries our best to fix the overall computational budget (by increasing the computational budgets of prior tokens and decreasing the latter) and shows that we can improve the mean accepted tokens. We achieve this by simultaneously boosting the accuracy of prior tokens while adopting simplified head structures (i.e., MLP) for subsequent tokens with parallel generation. This architectural innovation effectively balances accuracy and efficiency, ultimately achieving superior speedup.

Thank you for your suggestions and support.

Q1. It is not very clear why full tree attention works better than the tree attention methods in Eagle2, especially why this method cannot be applied in Eagle2. The authors should give an example to illustrate this better.

A1:

(1) Why FTA cannot be applied to Eagle2:

Let us refer to Fig. 2 for illustration. In the upper half of Fig. 2, the tokens [shines, glows, radiates] are outputs of Head 3, while [warmly, brightly, intensely] are outputs of Head 4. In Eagle2, the heads operate serially—both Head 3 and Head 4 are single-layer Transformers, and the output of Head 4 depends on the output of Head 3. In our method, however, these heads operate in parallel (i.e., Head 3 and Head 4 correspond to MLP1 and MLP2 in Fig. 1). Since the outputs of MLP2 are independent of MLP1, the outputs of MLP1 and MLP2 can be freely combined. In the lower half of Fig. 2, FTA leverages the independence between outputs of later heads to fill shorter paths in the tree with longer ones. In contrast, Eagle2’s Tree Attention (TA) requires strict dependency between tokens along each path (later tokens depend on earlier ones), making cross-path filling impossible.

(2) Why FTA outperforms TA:

The final speedup improvement is achieved by either increasing the mean accepted tokens or reducing draft time. Compared to TA, FTA increases the mean accepted tokens by extending candidate path lengths without incurring additional computational overhead, thereby maintaining the same draft time as TA. This allows FTA to achieve better performance than TA.

Q2. The experimental setup on A100 should be further claimed.

A2: The requirement of 4 GPUs applies exclusively to the scenario that the parameter size of target model is 70B. We clarify that our experiments on the NVIDIA A100 GPU did not include the 70B variant. All A100-based evaluations were conducted with 7B/8B/13B-scale Target Models.

Q3. Compare and discuss Cerberus[1], DDD[2], Hass[3].

A3:

(1) Cerberus [1]:

The observation in Cerberus that earlier heads exhibit higher accuracy aligns with our understanding. This is because, in parallel decoding, later heads predict tokens farther from the input (e.g., predicting the $i+2$-th token based on the $i$-th hidden state introduces greater error compared to predicting the $i+1$-th token). However, Cerberus still employs identical parameter scales and architectures across all heads, failing to prioritize earlier heads—a key distinction from our method. Their primary contribution lies in introducing sequential knowledge for parallel heads, whereas our approach explicitly optimizes model structure to prioritize early tokens.

(2) DDD [2]:

DDD improves Eagle2 by using a dynamic draft tree. Since DDD’s code is not open-sourced, we compare relative improvements over Eagle2. On the MT-Bench dataset, the performance gains of DDD and Gumiho are:

Our method consistently outperforms DDD across all tested models, demonstrating our effectiveness.

(3) Hass [3]:

Hass improves Eagle2 by addressing the training-inference inconsistency in serial transformer heads within SpD. While Hass focuses on training methodology, our work enhances performance via architectural modifications. These two directions are orthogonal: our method could integrate Hass’s training improvements for further gains.

In the table below, we compare the performance of our Gumiho with Hass. Additionally, we evaluate a variant of our model (denoted as Gumiho w/ Hass ) where Hass’s training method is applied to our Transformer Head. All results are conducted on an Mi250 GPU using the LLaMA3-8B model with a temperature of 0.

The aforementioned discussions will be included in the final version of our paper.

Q4. Why $p_i - \zeta < 0$ ($i=d+1$) is not assumed and discussed accordingly.

A4: Your concern is intuitively reasonable. From the perspective of probability definitions, it would indeed seem necessary to assume that probabilities should be greater than 0, even if this auxiliary variable is not a real probability but merely a constructed parameter for our proof.

However, throughout the proof process, we observed that such a constraint was unnecessary, and the proof could be completed without invoking this additional condition. Therefore, to maintain conciseness, we did not explicitly discuss this case separately.

Q5. In Appendix D, the authors show the comparison of Eagle2 and Gumiho on two datasets. What about the other four datasets?

A5: The other four datasets also exhibit similar results, and we will include the comparisons of these datasets in the final version of the manuscript.

Thank you for your suggestions and support.

Q1. The paper lacks a direct comparison with a strong baseline using flashdecoding.

A1: We have evaluated the baseline with FlashDecoding (FD) on the MT-Bench dataset using Mi250 GPU, as shown in the table below.

It is important to clarify that FD is an algorithm specifically optimized for attention computation. Its acceleration mechanism for LLMs is orthogonal to SpD – meaning these two approaches can be combined to achieve further acceleration beyond standalone SpD or FD implementations. The results of the combination will be shown in the final version of our paper due to the limited time of rebuttal.

Q2. The authors seem to introduce two different draft models (one large and one small), which could result in additional memory usage.

A2: Yes, you are correct. Our method does require more memory usage compared to Eagle2, and Medusa. Here’s why:

• Eagle2 use one single-layer Transformer head and reuse it serially across multiple speculative steps.
• Medusa only maintains five parallel MLP heads.
• Our approach, however, employs a two-layer Transformer head alongside five parallel MLP heads.

This architectural choice increases memory consumption. However, it enables a hybrid generation strategy: serial decoding for the first few tokens (to ensure accuracy) and parallel decoding for subsequent tokens (to maximize speed). While this introduces a memory-accuracy trade-off, it achieves a significant speedup improvement. We will explicitly discuss this trade-off in the final manuscript.

Q3. For the early draft model, using an approach similar to Eagle introduces an inconsistency between training and testing. I believe this issue becomes more pronounced as the network depth increases, as shown in Figure 4. Could the hass training method help mitigate this problem?

A3: Yes, our method could integrate Hass’s training strategy for further gains. Hass improves Eagle2 by addressing the training-inference inconsistency in serial transformer heads within SPD. While Hass focuses on training methodology, our work enhances performance via architectural modifications. These two directions are orthogonal, so our method could combine with Hass for further improvements.

In the table below, we compare the performance of our Gumiho with Hass. Additionally, we evaluate a variant of our model (denoted as Gumiho w/ Hass ) where Hass’s training method is applied to our Transformer Head. All results are conducted on an Mi250 GPU using the LLaMA3-8B model with a temperature of 0.

The aforementioned discussions will be included in the final version of our paper.

Q4. throughput improvement.

A4: We evaluated Gumiho under bs > 1 scenarios with Vicuna 7B, as shown in the table below. Experimental results demonstrate that the speedup effect degrades as batch size increases, aligning with observations in prior works like Eagle. Throughput speedup of baseline and Gumiho is computed at their respective maximum bs, as Eagle do.

The degradation in speedup with larger batch sizes can be primarily attributed to two key factors:

(1) Inter-Sequence Execution Time Imbalance. With larger batch sizes, the number of accepted tokens varies across sequences within a batch, leading to divergent completion times. Due to the "bucket effect" in batch processing — where the entire batch's completion time is determined by the slowest sequence — completed sequences must wait for unfinished ones. This waiting time exhibits growth as batch size increases. Notably, directly removing completed sequences and inserting new ones theoretically mitigates this issue but proves impractical in reality. Newly inserted sequences require full context re-computation (prefill phase), whereas existing sequences leverage pre-cached KV for efficient decoding. This computational asymmetry further increases overall latency.

(2) Computational Resource Contention. Large batch sizes push SpD into a computation-bound regime. During parallel verification, the target model must process a massive number of candidate tokens per step (quantity = draft tokens × batch size). Larger batch sizes prevent full parallelization of computing units, causing the verification process to degrade into partially sequential operations. This diminishes the theoretical acceleration benefits.

Given these constraints, most SpD solutions (e.g., Lookahead, Medusa, Hydra, Cerberus, Eagle2, Hass) primarily focus on optimizing single-sequence (bs=1) scenario. Our work follows them, prioritizing inference efficiency optimization for single-sequence processing.