Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding
Authors: Zachary Ankner
Rishab Parthasarathy
Aniruddha Nrusimha
Christopher Rinard
Jonathan Ragan-Kelley
William Brandon
Venue: COLM (2024)
Keywords: Speculative Decoding, Medusa, Draft Heads, Blockwise Parallel Decoding
TL;DR: We introduce sequential dependence to speculative deocding draft heads which signifigantly improves LLM decoding performance.

Abstract: To combat the memory bandwidth-bound nature of autoregressive LLM inference, previous research has proposed the speculative decoding framework. To perform speculative decoding, a small draft model proposes candidate continuations of the input sequence that are then verified in parallel by the base model. One way to specify the draft model, as used in the recent Medusa decoding framework, is as a collection of lightweight heads, called draft heads, that operate on the base model's hidden states. To date, all existing draft heads have been sequentially independent, meaning that they speculate tokens in the candidate continuation independently of any preceding tokens in the candidate continuation. In this work, we propose Hydra heads: a sequentially-dependent drop-in replacement for standard draft heads that significantly improves the accuracy of draft head speculation. We further explore the design space of Hydra head training objectives and architectures, and propose a carefully tuned Hydra head recipe, which we call Hydra++, that improves decoding throughput by up to 1.31x and 2.70x compared to Medusa decoding and autoregressive decoding respectively. Overall, Hydra heads are a simple and well-motivated intervention on standard draft heads that significantly improve the end-to-end speed of draft head-based speculative decoding. We make our code publicly available at https://github.com/zankner/Hydra.

Published as a conference paper at COLM 2024
Hydra: Sequentially-Deeedder Draft Heads for Medusa De- coding
Zachary Ankner 1 Rishab Parthasarathyl Aniruddha Nrusimha'
Christopher Rinard! Jonathan Ragan-Kelley! William Brandon MIT MosaicML.
Abstract
To combat the memory bandwidth- nature of autoregressive LLM inference, previous research has proposed the speculative decoding frame- work. To perform speculative decoding, a small draft model proposes candidate continuations of the input sequence that are then verified in parallel by the base model. One way to specify the draft model, as used in the recent Medusa decoding framework, is collection of lightweight heads, called draft heads, that operate on the base model's hidden states. To date, all existing draft heads have been sequentially independent, meaning that they speculate tokens in the candidate continuation independently of any preceding tokens in the candidate continuation. In this work, we propose Hydra hends: standard draft heads that significantly improves the accuracy of draft head speculation. We further explore the design space of Hydra head training objectives and architectures, and propose a carefully tuned Hydra head recipe, which we call Hydra++, that improves decoding throughput by up 1.31x and 2.70x compared to Medusa decoding and autoregressive de- coding respectively. Overall, Hydra heads are simple and well-motivated intervention on standard draft heads that significantly improve the end-to- end speed of draft head-based speculative decoding. We make our code publicly available at https:// http:///gthubbooo/m /aae///nnnnnenanee Hydra
weeeeniiiall-dppddeeee drop-in replacement for
1 Introduction
As transformer-based large language models (LLMs) have entered widespread deployment, research into improving the inference efficiency of these models has become increasingly important. While LLM pretraining achieves high hardware utilization by operating over the entire input sequence in parallel, the efficiency of LLM inference has traditionally been constrained by the need to generate tokens one by one in sequence. On current GPU hardware, the serial nature of LLM decoding makes it a memory bandwidtt--ounnd problem, with throughput limited by the movement of large weight matrices from GPU main memory to local registers. As each generation step requires accessing the entirety of the model's s weights, but only involves comparatively small number of FLOPs (processing just one token for each sequence in the batch), LLM decoding tends to under-utilize the GPU's abundant capacity for floating-pont -pinn computation.
To mitigate the memory bandwidth bottleneck in sequential LLM decoding, recent research has investigated accelerating LLM inference through speculative decoding. Speculative de- coding uses smaller draft model to propose multi-token candidate continuation of the current sequence on each generation step. The original LLM then verifies all tokens the candidate continuation in parallel, appending some subset of them to the sequence and discarding the rest. Because each verification step requires only single forward pass through the original LLM, but may result in more than one token being appended to the sequence, speculative decoding can accelerate decoding by reducing the amount of weight data movement required per generated token.
"qqual contribution. Correspondence antner@mmitedu

---

Published as a conference paper at COLM 2024
A critical component in any application of speculative decoding is the choice of draft model, which must be cheap enough such that the cost of querying it does not negate the efficiency gains from querying the base model in parallel, but accurate enough such that the acceptance rate in the verification step remains high. While the draft models used in speculative decoding have traditionally been stand-alone, independently-traiin
language models, [Stern ntter et al. (2018) and Cai et al. (2024) instead investigate structuring the draft model as a collection of lightweight heads operating on the base model's semantically rich hidden states. We refer to the lightweight heads that operate on the original hidden states of the LLM as draft heads. In the draft head paradigm, each draft head is responsible for predicting the identity of a token fixed number of steps into the future.
All draft heads to date make predictions only as function of the base model's hidden states from previously verified tokens, making them unaware of earlier tokens in the current candidate continuation. Because of the strong statistical dependencies between neighboring tokens in language, this sequential independence limits the prediction accuracy of existing draft head architectures. In this work, we propose Hydra heads: a drop-in sequentially dependent alternative to standard draft heads that improves token prediction accuracy and thus end-to-end decoding throughput. To construct sequentially dependent draft heads, set each head's output to be a function of the candidate continuation so far. This simple design change leads to significantly better speculation quality as compared to standard draft heads, increasing the average candidate continuation acceptance length by up 0.46 tokens. This improvement in speculation quality corresponds to a significant improvement in decoding speeds, with Hydra head-based decoding achieving up to 1.1x better throughput than Medusa decoding.
In addition to proposing Hydra heads, we further explore the design space of their training objective and architecture. We find that extending the depth of the draft head MLPs, using a teacher distillation objective, and adding an extra transformer decoder layer to better encode the already verified sequence achieves up to 1.31 x and 2.70x higher throughput than standard Medusa decoding and regular autoregressive decoding respectively.
Finally, we investigate Hydra and Hydra++ decoding in alternative inference settings. The first setting that we consider is batched inference. The next setting we examine is non- greedy decoding. We show that by using typical acceptance sampling (Cai et al. 20244,, con-distribution-rrssse preserving verification criterion, Hydra++ can achieve the same quality generations as non-greedy sampling of the base model while not compromising acceptance length.
Contributions
In this work, we present the following contributions:
We analyze the standard formulation of draft heads and observe that they are se- quentially independent during decoding. We propose Hydra heads as a sequentially dependent alternative and show that introducing sequential dependence increases end-to-end decoding throughput by up to 1.10x as compared to Medusa decoding (Section 6.1
We explore the design space of Hydra heads to produce draft head recipe Hy- dra++ that further increases decoding throughput by up to 1.31 and 2.70x over Medusa decoding and standard autoregressive decoding respectively (Section 3.1 Section 6.1
We analyze the performance of Hydra decoding in the batched inference setting, demonstrating that it achieves better throughput than Medusa at all batch sizes evaluated (Section 6.2)
We demonstrate that sampling using the typical acceptance criterion allows Hy- dra++ to achieve the same quality as sampling from the base model while preserving the throughput benefits of speculation (Section 6.3
2

---

Published as a conference paper at COLM 2024
2
Background
Speculative decoding. Speculative decoding (Stern al. 2018 Leviathan et al. 2023 Chen
et at. 2023) provides general framework for efficient LLM decoding. Speculative decoding generates text by combining an expensive, high-quality base model with a cheaper, lower- quality draft model. For each decoding step, the draft model generates one or more candidate continuations, each of which extends several tokens into the future. We then use a single forward pass through the base model to verify these candidate continuations in parallel based on some verification criterion. The verification process determines which candidate tokens will be appended to the sequence and which will be discarded.
In the simplest form of speculative decoding, the draft model only generates a single candidate continuation on each generation step. Letting x<i be the sequence that has been generated so far and fixing some speculation length K, we query the joint distribution of the draft model Paraft ( (X++1...., XtKK X<t) to generate candidate continuation I & We then invoke the base model on the candidate continuation to compute the conditional probabilities: pbase ($ xxst)..., Pbase($i+K X-st,
&(-K-1); )); querying the base model is done in a single forward pass. These base model probabilities become the input to the verification criterion, which selects some prefix f +ooces of the candidate continuation to accept, discarding the rest.
Common verification criteria for use with speculative decoding include rejection resampling (Leviathan et al. 2023 Chen et al. 2023), which guarantees that the output distribution matches the base model's 's distribution, and greedy acceptance (Stern et al. 2018), in which candidate tokens are accepted if they match the base LLM's most likely prediction. For all verification criteria in common use, using draft model whose predictions more accurately match those of the base model will result in increased average acceptance lengths, and consequently greater decoding throughput.
Tree decoding. Speculative decoding can be generalized to settings in which the draft model proposes a free of candidate continuations, rather than single candidate continu- ation (Miao et al. 2023 Spector & Re 2023 Cai et al. 2024). Nodes this candidate tree correspond to candidate tokens, and the children of node represent different possible tokens that might follow it in the continuation. Thus, each path along the tree represents different candidate continuation. To populate node with children, we query the draft model for the m most likely tokens that might follow it, conditioned on the sequence generated so far and the candidate continuation defined by the path to the node from the root of the tree. The children at each node are sorted in descending order of conditional probability. Typically, static trees are employed where the structure of the tree fixed at design time, meaning that the number of children m at each position in the tree does not depend on any runtime data.
After populating the candidate continuation tree using the draft model, we compute the conditional probabilities of all nodes in the tree using a single forward pass through the base model. We query the base model for these conditional probabilities in single forward pass by packing all of the tree's tokens into a single input sequence, and manipulating the attention mask to ensure that each token can only attend to its parents in the tree. The conditional probabilities obtained from querying the base model can then be used as input to the same verification criteria that are used in the single-candidate setting.
Lightweight heads as a draft model. While typically the draft model used in speculative language model, Stern et al. (2018) define the draft
decoding is an independently-trainet model as collection of lightweight heads, which we reter to as draft heads, that take as input the base model's hidden state. Taking K to be the maximum speculation length, the draft model used by Stern et al. is defined by a collection of small MLPs färaft,, ..., fdraf, responsible for predicting the tokens I,... K steps into the future. The predictions of these heads are statistically independent of each other; letting X<! denote the sequence generated so far, and letting h, denote the last-layer hidden state of the token most recently processed
3

---

Published as a conference paper at COLM 2024
Existing Draft Heads
Hydra Draft Heads (Ours)
"Pask"
"Pakk"
Hiddie
Head
Head
Base Model
Medel
Conteat
"Haka"
Figure A visualization of generating a candidate continuation using existing draft heads and using Hydra draft heads. Lines going into head represent inputs to the draft head. While the only input to existing draft heads is the base model's 's last-layer hidden state for the most recently processed token, Hydra heads leverage earlier tokens in the candidate continuation as additional inputs.
by the base model, Stern et al. compute their draft predictions on each generation step as Paran (xi+i Xg(+i-1) ( 1)
Figure provides a visualization of candidate continuation generation using draft heads.
Medusa decoding. Medusa decoding (Cai et al. 2024 is a particular configuration of the techniques listed above. Specifically, it is speculative decoding with tree of candidates, where the draft model is collection of draft heads.
While Medusa decoding is agnostic to the architecture used for each draft head fdrat, Cai et al. 2024)) choose to use single-layer MLP with residual connection.
Hydra Heads
The key observation behind Hydra heads is that there is sequential dependence in stan- dard draft heads, i.e., each draft head makes predictions independently. A draft model de- fined by collection of draft heads predicts the identity of the tt future token as faran 1). However, h4-1 is only function of the already generated sequence xgg-1- Thus, when using draft heads:
Paraft (51:)) &111 8 Paraft (f+1/
Intuitively, this means that there is no sequential dependence between draft heads: when we use a draft head to speculate the ith token in a candidate continuation, it is unaware of the 2nd ,,,,, (i 1hth tokens in the candidate continuation.
We propose Hydra heads, which are requentially-deeeddee draft heads. Hydra heads are sequentially dependent as they are a function of both the base model's hidden state up to time as well as the input embeddings the tokens sampled by previous Hydra heads. Namely, the draft model is now collection of Hydra heads {fHydra,1 ...drr,,....,/ HHdda,, and the ft future token's distribution is parameterized by this collection of Hydra heads as:
Pdraft X &.41 & 1) = fHydra, (h; Xt, 81+1, -++i-1)
where h4 is again the base model's hidden state of the final token in the already decoded sequence. The sequential dependence of Hydra heads v.s. standard draft heads is visualized
in Figure [ We use the term Hydra Decoding to refer to speculative decoding with tree candidates and Hydra heads.
4

---

Published as a conference paper at COLM 2024
As with Medusa, the framework of Hydra decoding is compatible with any choice of model architecture used to implement fHydraji. The most basic Hydra head architecture we examine is simply a single hidden layer MLP whose input is the hidden state 14- concate- nated with the input embeddings of the preceding tokens in the candidate continuation Exy, E E where the concatenation is performed along the feature dimension.
3.1 Hydra++
With the goal of further improving draft heads, we investigate draft head training objective and architecture improvements orthogonal to the introduction of sequential dependence. We detail this search for the optimal draft head recipe in Appendix A Ultimately, we find that three changes are beneficial:
1. Scaling: We extend the MLP of each head to layers. In our experiments we found that scaling 5 layers and beyond provides additional benefit.
2. Distillation: Following (Zhou et al. 20244, we train on self-distillation objective where the draft heads are trained to predict the base model's distribution for a given token instead of the true token.
3. Prefix Attention:
To improve our draft model's ability to condition on information from across the entire context, as opposed to just the most recently-verified token, we extend the base model with an additional self-attention decoder layer whose only role is to produce more informative hidden states for use input to the draft model. This added layer is only queried once per decoding step.
We refer to the version of Hydra that leverages all of the above changes as Hydra++.
4 Discovering performant decoding trees
Similarly to Medusa, we always perform tree-based speculative decoding with a static tree topology computed offline. Computing performant tree topology for given inference setting is nontrivial, because different settings call for different tree topologies; the choice of base model, draft model, batch size (discussed more in Section 6.2), and hardware can all affect the relative performance of different trees.
We derive our decoding trees in a data-driven manner using two-stage algorithm: first, we find a sequence of "proposal" trees Ti...., tree approximately maximizes expected acceptance length given its size; then, we choose the optimal tree size for given setup by empirically measuring the end-to-end throughput achieved using each Ti, and selecting the tree which maximizes throughput.
.,. Tn with sizes I,...,N N such that each proposal
To determine the sequence of proposal trees Ti...., "i....,TT,, TN, we follow simple iterative greedy procedure. We first initialize T, to the trivial one-node tree. Then, on each step i, we simulate speculative decoding using T;-1 1 existing node that would yield the greatest improvement in expected acceptance length added to the tree. We then add this node Ti-1 to form the tree Tü, and repeat.
corpus of sample text, and identify the child of
A fter computing Tı.... TN. we measure the throughput of speculative decoding using each Ti in our desired inference configuration (i.e., batch size, hardware, etc.), and select the tree which empirically maximizes decoding throughput.
In practice, we set the maximum tree size as N 100, and use a 100-question subset of the A Alpaca dataset (Taori et al. 2023) to gather our simulated acceptance length and throughput statistics. We provide more details on the trees discovered for each decoding strategy and batch size in Appendix B
5 Shared training and evaluation details
In this section, we detail the elements of our training and evaluation procedure that are common across all our experiments.

---

they are Published as a conference paper at COLM 2024
Hydra vs, Medusa Decoding
3.72x
3,62x
(Tok)
3.04
Avg. Througho (Tek:
2.50
2.5
2,39
2,00
Baseline
2,17x
1,47
Aug. Acceptanee
Hydra Hydra+++
1.00
Lo
1,0
1,00
1.00
25
Vicuna
Vicuna 130 Base Model
Vicuna 330
Vicuna
Vicuna 130 Base Model
Vituna 330
Figure 2:
Performance comparison MT-Bench of Hydra++, Hydra, Medusa, and the baseline of autogressive decoding. Hydra heads increase decoding throughput and average acceptance length compared to all other methods.
Models. As our base model we build on the Vicuna family of models CChingg et al. 2023 which are conversation-iintunne LLaMa models (Touvron al. 2023). We consider 7B, 13B, and 33B parameter Vicuna models.
Training. While draft heads can be trained in conjunction with the base model, in this work we only study base models with frozen weights. All models are trained on the ShareGPT dataset (ShareGPT, 2023), collection of multi-turn conversations. Training is performed on 8 NVIDIA A100-80GB GPUs and conducted using the HuggingFace Trainer (Flgg) We use cosine learning rate schedule with warmup (Loshchilov & Hutter [2utt)][[01]] and a peak learming rate 1e-3, and we use the AdamW optimizer (Loshchilov & Hutter 2019) with parameters 31 0.9, B2 0.999, All Hydra and Medusa heads are trained for one epoch, as we observed that performance for those models saturates at one epoch and fails to improve with further training. All Hydra++ heads are trained for ten
epochs.
Evaluation. All evaluations are performed on MT-Bench (Zheng et al., 2023), a multi-turn conversation benchmark. Unless otherwise specified, experiments are conducted using speculative decoding with the greedy verification criterion; since there is no stochasticity in the greedy sampling procedure, we do not report the quality of generations identical to the base model. Instead, we report the average throughpuf, which is the number of tokens generated per second, and the average acceptance length, which is the number of tokens generated per decoding step, to evaluate the speed and quality of Hydra decoding, We benchmark all 7B and 13B parameter experiments on a single A100-40GB GPU and all 33B parameter experiments on single A100-80GB GPU.
Results
In this section we investigate the performance characteristics of Hydra decoding. We examine the effect of draft model architecture on throughput and latency across range of batch sizes, and also investigate the effect of draft model architecture on generation quality when decoding with non-distribution-preervii verification criteria.
6.1 Batch-size-1 decoding throughput experiments
To assess the effect of our interventions draft model prediction accuracy (and thus decoding throughput), we compare the batch-size-l decoding throughput achievable using Medusa draft heads, our basic Hydra draft heads, and our enhanced Hydra++ draft heads.

---

Published as a conference paper at COLM 2024
Effect Batch Hydra Medusa
W1-
Awg. Lakency 15 UUee7)
Avg. Throug Tox
Baseline
X
Medusa
Hydra
Hydra+++
Batch Size
Batch Size
Figure 3: Performance comparison on MT-Bench of Hydra++, Hydra, Medusa, and the baseline of autogressive decoding for batched inference. Hydra heads increase decoding throughput for all batch sizes examined.
We also include the throughput of hon-speculative autoregressive decoding baseline. We summarize the results of these decoding throughput experiments in Figure
We find that across all base model sizes evaluated, Hydra achieves higher average accep- tance lengths than Medusa, and Hydra++ achieves higher acceptance lengths than Hydra, leading to significant improvements in decoding throughput in both cases. Specifically, Hydra heads achieve 2.36x, 2.17x, and 2.15x improvement in throughput as compared to autoregressive decoding for the 7B, 13B, and 33B parameter base models respectively. This translates to throughput improvement over Medusa decoding of 1.11x, x, 1.10x, and 1.11x for the 7B, 13B, and 33B parameter base models respectively. Furthermore, Hydra++ even more performant and achieves throughput improvements compared to autoregres- sive decoding of 2.70x, 2,500,, and 2.53x which translates to throughput improvements over Medusa of 1.27 1.27x, and 1.31 x for the 7B, 13B, and 33B parameter base models respectively. These results demonstrate that making draft heads sequentially dependent significantly improves their prediction accuracy, and thus their decoding speed. Moreover, these results show that any overheads introduced by passing from Hydra to the more expressive Hydra++ architecture are more than compensated for by the increase in accuracy that those changes enable. We further investigate draft head overheads in Appendix D
6.2 Speculative decoding for batched inference
Speculative decoding techniques are typically evaluated in the batch-size-1 setting. When there is only single sequence in the batch, decoding is extremely memory bandwidth- bound, and large numbers of FLOPs can be consumed in the verification step of speculative decoding "for free" without significantly increasing the latency per decoding step. However, at larger batch sizes it is easier for verification to become compute-bound, and the number of tokens verified per sequence per step must be more tightly controlled to avoid saturating the GPU's compute capacity and entering the regime where speculative decoding becomes unprofitable.
To assess whether not the performance gains from Hydra and Hydra++ observed at batch size continue to hold in the batched inference regime, we evaluated the throughput and latency of Medusa, Hydra, and Hydra at batch sizes {1,2,4, 8, using greedy verification and a 7B base model. We derived the decoding tree used for each draft model and batch size configuration using the algorithm described in Section
We plot the relationship between batch size, throughput, and latency using the 7B base model in Figure 3] While find that all speculative decoding techniques outperform standard autoregressive decoding for all batch sizes examined, the relative improvement of speculative decoding decreases the batch size increases. Specifically, for batch size
Results
7

---

Published as a conference paper at COLM 2024
Posterior Threshold's Effect Decoding
(0)
(b)
Avg. Acceptance Lenoth (Tok)
7.8
7.6
Bas Model
2
7.4
Hyora
Hydra++4
7.2
0.25
0.15
0.25
Posterior Threshold
Posterior Threshold
Figure 4: Average speculation length and MT-bench score as a function of the posterior threshold for typical acceptance-basee decoding. Hydra++ achieves comparable generation quality to the base model while preserving acceptance length.
Hydra++ has 2.70x improvement in throughput compared to standard decoding, but this gain decreases to 1.63 x at batch size 8. These results suggest that while the gain from Hydra decoding is less sigmificant at larger batch sizes, Hydra decoding is still an improvement over both Medusa and standard decoding at larger batch sizes.
6.3 Typical acceptance sampling
So far we have used the greedy acceptance criterion, where candidate tokens are only accepted if they match the greedy next-token prediction of the base LLM. We now evaluate the impact of Hydra decoding on throughput and quality using the non-greedy, non- distribution-presevvinn typical acceptance verification criterion introduced by Cai al. (2024).
Typical acceptance criterion. The purpose of the typical acceptance verification criterion to sample more diverse and creative sequences than greedy acceptance, while preserving the efficiency benefits of speculative decoding by avoiding the degradation in acceptance rate observed when employing rejection resampling (Gante) 2023 Spector Re 2023).
The typical acceptance criterion specifies that speculated token 44ti is accepted if:
Pbase (f)4 Xgt, f,41,..., T))-); > min(e, a
(P (- ( & ((+)11;)))
where is known as the posterior threshold, is known as the posterior alpha, is the sampling temperature, and H(.) ) is the entropy. Both and are hyperparameters to be tuned. For further analysis of typical acceptance, we refer the reader to Cai et al. (2024).
Setup. We evaluate how different settings of e and a affect both acceptance length and generation quality. Following Cai al. (2024), evaluate typical acceptance on the "Writing" and "Roleplay" categories of MT-Bench, and report the average LLM-as-a-judge score to quantify generation quality (Zheng et 2023). We fix the sampling temperature T = 0.7, vary the posterior threshold E {0.05, 0.1, 0.2,0,,55 and set the posterior alpha as & Ve.
Results. We plot how varying the posterior threshold affects both the average speculation length and the quality of the resulting generations in Figure For Medusa, Hydra, and Hydra++, increasing the posterior threshold slightly decreases the average speculation length, but for all posterior thresholds examined, Hydra and Hydra++ have significantly higher average acceptance length than Medusa. While neither Medusa nor Hydra is able to achieve the same quality random sampling from the base model for any of the posterior thresholds considered, for e = 0.15 Hydra++ achieves the same generation quality as sampling directly from the base model. These results demonstrate that the improved head quality achieved from Hydra++ is necessary to match the generation quality of the baseline for non-greedy inference while still maintaining high average acceptance length.
8

---

Published as a conference paper at COLM 2024
7
Related Work
Accelerating LLM inference is an area of active research. The technique our work is based on, speculative decoding, was first proposed by Leviathan al. (2023) and Chen et al.) (2023), and anticipated in restricted form by Stern et al. (2018). Recent work has explored alternatives standard draft models such as using retrieval mechanisms to propose con- tinuations (He et al. 2023), and reformulating language model sampling in terms of Jacobi iteration (Fu et al. 2023). Another direction of speculative decoding research has investi- gated veritying tree of candidate continuations rather than single continuation (Miao et al. 2023 Spector & Re 2023 Cai et al.) 2024). In addition to tree decoding, Spector & Re
(2023) also propose extending the basic speculative decoding framework by constructing hierarchy of draft models, with each aiding speculative decoding for the next. Other drary directions of research on speculative decoding include online training of the draft model based on user queries (Liu et al. 2023a) and knowledge distilattion based alignment of the draft model (Zhou al. 2024). We would also like to acknowledge the concurrent work EAGLE al.
contemporary
w which is the work most similar to ours. We discuss
We would also like to acknowledge ZZhnn et al. (2024); Wertheimer al. 2024) who concurrently investigated sequentially dependent draft heads.
EAGLE in Appendix C
Another direction for accelerating LLM inference is minimizing the memory impact of
LLMs. A common technique is to compress the LLM either by quantizing its weights or
pruning the features of the model (Dettmers et al., 2022 Xiao et al., 2023 Frantar et al.
2023: Frantar & Alistarh| 2023 Liu et al. 2023b: Alizadeh et al. Sheng et al. 2023)
1o decrease the memory footprint of the KVcache, Shazzer 2019 and Ainslie et al. (2023
introduce multi-query and grouped-query attention respectively. These works reduce the
size of the KV-cache by using fewer key and value heads as compared to the number of query heads in attention. Another method for decreasing the memory footprint of LLMs is
knowledge distillation, where smaller student network is trained to be as accurate as the original larger model (Sanh et al. 2020). These memory-reduction and inference acceleration techniques are orthogonal to, and potentially complementary with, speculative decoding,
Increasing the batch size which inference is performed is another technique for improving LLM inference throughput. Multiple works investigate better scheduling for batched
inference and improved management of shared resources during batched inference YYu
et al. 2022 Kwon et al., 2003)..
8 Conclusion
In this work, we systematically examine draft head-based speculative decoding and propose methods for improving the speculation quality of draft heads. We make the observation that previously-propose
draft heads are sequentially independent, leading to poor prediction quality. To fix this problem, we propose Hydra heads: a drop-in, replacement for standard draft heads. Hydra heads are made sequentially dependent by taking as input the base model's input embeddings of tokens in the candidate continuation. This simple change leads to significant improvements in decoding speed: Hydra decoding achieves up to a 1.11x x improvement in end-to-end throughput as compared to Medusa decoding. We also investigate different training objectives and architectures for Hydra heads, ultimately proposing a Hydra head recipe we call Hydra++ that increases decoding throughput by up to 1.31 x and 2.70x as compared to Medusa and autoregressive decoding respectively. Finally, we demonstrate that Hydra++ continues to confer benefits when performing batched inference, and that by using typical acceptance sampling Hydra++ can achieve the same quality as non-greedy sampling of the base model without compromising accepted continuation length. Draft head based speculative decoding is an efficient and simple alternative to the standard speculative decoding paradigm, and our work takes important step towards maximizing the performance of decoding with draft heads through the construction of sequentially dependent draft heads.

---

Published as a conference paper at COLM 2024
Acknowledgments
This work was initially performed as class project for MIT's NLP class, 6.8611, and as such we would like to thank the 6.8611 teaching staff for their collective help. In particular, we would like to thank Jacob Andreas, Yoon Kim, and Chris Tanner for teaching the course, and Marco Nocito and Michael Maune for their feedback on the paper throughout the course.
References
Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. GQA: Training generalized multi-query transformer models from multi-head checkpoints. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 4895-4901, Singapore, December 2023, Association for Computational Linguistics. doi: 10. 10.18653/v /v/ 2023.emnlp-manl 298. URL/httpss///claathhllos
org/2023. 2023. emnlp-man mapp..min.... 298
Keivan Alizadeh, Iman Mirzadeh, Dmitry Belenko, Karen Khatamifard, Minsik Cho, Carlo C Del Mundo, Mohammad Rastegari, and Mehrdad Farajtabar. Llm in flash: Efficient large language model inference with limited memory, 2024.
Tianle Cai, Yuhong Li, Zhengyang, Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao. Medusa: Simple Ilm inference acceleration framework with multiple decoding heads, 2024.
Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sa sampling, 2023.
Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, lon Stoica, and Eric P. Xing. Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023.
URL httK://ppp////sssoorr/////////00 /-tp:////0022////000002-------10
Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm.int8(): 8-bit matrix multiplication for transformers at scale. arv preprint arXiv:2208.0733) 2022.
Elias Frantar and Dan Alistarh. Sparsegpt: massive language models can be accurately pruned in one-shot. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023.
Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. OPTQ: Accurate quanti-
zation for generative pre-trained transformers. In The Eleventh International Conference on Learning Representations, 2023. URL/htpps///peereviee http://sss/ppppeeeeeeee//e////ea dd
Yichao Fu, Peter Bailis, Ion Stoica, and Hao Zhang. Breaking the sequential dependency of Ilm inference using lookahead decoding, November 2023. URL https://lmsys: org/ 11-0023-11-1--111
Joao Gante. Assisted generation: new direction toward low-latency text generation, May generratonnne
2023. URL
Zhenyu He, Zexuan Zhong, Tianle Cai, Jason D Lee, and Di He. Rest: Retrieval-based speculative decoding, 2023.
HuggingFace. Hugging face trainer. URL http:p///wugggigee....oooo//ssss mainn classes/trainer!
Neel Jain, Ping-yeh Chiang, Yuxin Wen, John Kirchenbauer, Hong-Min Chu, Gowthami Somepalli, Brian R. Bartoldson, Bhavya Kailkhura, Avi Schwarzschild, Aniruddha Saha, Micah Goldblum, Jonas Geiping, and Tom Goldstein. NEFTune: Noisy embeddings improve instruction finetuning. In The Twelfth International Conference on Learning Repre- sentations, 2024. URL / http:////ppppoeeeeeeeeeee/r/rono
10

---

Published as a conference paper at COLM 2024
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and lon Stoica. Efficient memory management for large language model serving with pagedattention, 2023.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In Proceedings of the 40th International Conference on Machine Learning, ICML'23, JMLR.org, 2023.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle: Speculative sampling requires rethinking feature uncertainty, 2024.
Xiaoxuan Liu, Lanxiang Hu, Peter Bailis, Ion Stoica, Zhijie Deng, Alvin Cheung, and Hao Zhang. Online speculative decoding, 2023a.
Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shri- vastava, Ce Zhang, Yuandong Tian, Christopher Ré, and Beidi Chen. Deja vu: contextual sparsity for efficient lims at inference time. In Proceedings of the 40th International Conference on Machine Learning, ICML'23, JMLR.org, 2023b.
Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations, 2017. URL https::/ net/ forum/1d=Skg895oxx
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In 7th Interna- tional Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL SEKLl//teer/eeeereeeeeeeeerrrrrr
Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Rae Ying Yee Wong, Alan Zhu, Lijie Yang, Xiaoxiang Shi, Chunan Shi, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerating generative large language model serving with speculative inference and token tree verification, 2023.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter, 2020.
ShareGPT. Sharegpt, 2023.
URL htttt:///iigiiigee..c///attsts/s
Vicuna_unfilterr
Noam Shazeer. Fast transformer decoding: One write-head is all you need, 2019.
Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, lon Stoica, and Ce Zhang. Flexgen: High-throughput generative inference of large language models with a single GPU. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA, volume 202 of Proceedings of Machine Learning Research, pp. 31094-31116. PMLR, 2023. URL. https://pproeedng.......esssssss htt.s//.......re////////e/a/aaa httll
Benjamin Frederick Spector and Christopher Re. Accelerating LLM inference with staged speculative decoding. In Workshop Efficient Systems for Foundation Models @ ICML,2023, 2023. URL httl:/ttps:/oopeeeiieeee http://wwwwwweeetttoorr//m/daaaa 3VY
Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep au- toregressive models. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL URL neur wwwppppecppppppppprreeeee.. htcc:/ppprrfillssppppr////// Ilama fieec472220444556000000652222222
Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto. Stanford alpaca: An instruction-following
http///tttth...../ http///tt....mttttaaaaaa/aaaaaaa alpaca 2023.
11

---

Published as a conference paper at COLM 2024
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023.
Davis Wertheimer, Joshua Rosenkranz, Thomas Parnell, Sahil Suneja, Pavithra Ranganathan, Raghu Ganti, and Mudhakar Srivatsa. Accelerating production Ilms with combined token/ boken/embeddiig speculators. Xiv preprint ar XrXoo:2404.11124 2024.
Heming Xia, Zhe Yang, Qingxiu Dong, Peiyi Wang, Yongqi Li, Tao Ge, Tianyu Liu, Wenjie Li, and Zhifang Sui. Unlocking efficiency in large language model inference: A compre- hensive survey of speculative decoding. arviv preprint ar 4r1.::2401.07851, 2024.
Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han. SmoothQuant: Accurate and efficient post-training quantization for large language mod- els. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 38087-38099, wtt..pprres//ee2///// http:///222xxxx2ccc.coo html
PMLR, 23-29 Jul 2023. URL httphttts//prrceeediiii
Gyeong-In Yu, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, and Byung-Got Gong Chun. Orca: A distributed serving system for Transformer-Bast Based generative models. In 16th USENIX Symposium Operating Systems Design and Implementation (OSDI 22), pp. 521- 538, Carlsbad, July 2022. USENIX Association. ISBN 978-1-9399333-2--1 URL https
wwwww.usenix.ook
Aonan Zhang, Chong Wang, Yi Wang, Xuanyu Zhang, and Yunfei Cheng. Recurrent drafter for fast speculative decoding in large language models. arXiv preprint ar Xrr::2403.09919, 2024.
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, Hao Zhang, Joseph E. Gonzalez, and lon Stoica. Judging LLM-as-a-judge with MT-bench and chatbot arena. In Thirty-seventh Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2023.
URL httpps: http::///pppeeeeeeeeeeeeeeerrree nett/forumiidduccHHOO
Yongchao Zhou, Kaifeng Lyu, Ankit Singh Rawat, Aditya Krishna Menon, Afshin Ros- tamizadeh, Sanjiv Kumar, Jean-François Kagy, and Rishabh Agarwal. Distillspee: Improv- ing speculative decoding via knowledge distillation. In The Twelfth International Conference on Learning Representations, 2024. URL htth:tts::///pptt http::////oeereeveeeee///////rtr
12

---

Published as a conference paper at COLM 2024
Exploring Training Objective
0.96x
3.08
0.00x
1.00%
Accopttnnc Langth CTOME
Avo. Throughout (Tok (5)
2.0
Hydr
Hydra Notse
Hydra Teach. Loss
Hydra
1/-
Base Modal
Base Model
Figure 5: Performance comparison on MT-Bench of different Hydra head training objectives. Training based on a teacher loss leads to the most performant Hydra heads.
A Exploring the Design Space of Hydra Heads
In this section explore modifications to the training procedure and architecture of Hydra heads.
A.1 Exploring the Training Procedure of Hydra Heads
Adding noise to the input sequence. Jain et al. (2024) demonstrate that adding noise to the input embeddings of an LLM during finetuning improve the resulting model's performance. Specifically, they sample noise factor #noise and then add the scaled noise to the input embeddings, where B is the batch
R 8xxLxd
(,1,)) scale by
VId
size, L is the sequence length, d is the model dimension, and &noise is a hyperparameter that controls the strength of the noise. We consider whether applying such noise the hidden states of the base LLM can also improve the performance of the Hydra heads. For our experiments, we set noise = 75.
Distilling the base LLM. In the standard Medusa decoding framework, the draft heads are trained to predict the text of the underlying fine-tuning dataset. We question whether this the optimal training objective as, during inference, the goal of the draft heads is only to predict the token which the base LLM would have autoregressively predicted. Following Zhou et al. (2024), we investigate using a teacher loss where each Hydra head's training loss the cross entropy between its predicted distribution and the base model's next token distribution.
A.1.1 Results
To test how each of our training interventions affects Hydra head quality, we evaluate both interventions separately as well as jointly. We compare decoding using Hydra heads trained with the proposed interventions and decoding using Hydra heads trained in the standard manner. We report the decoding throughput as well as the average acceptance length for each intervention in Figure 5) We find that the most performant intervention is to just train on the teacher loss without any additional embedding noise. Specifically, decoding with heads trained with just the teacher loss achieves 1.04x improvement in throughput for Vicuna 7B compared to decoding with vanilla Hydra heads. Interestingly, we find that any addition of noise to the input sequence degrades the acceptance length and thus the decoding speed. Based on these result, we conduct all future experiments using the teacher loss instead of the next token prediction loss.
13

---

Published as a conference paper at COLM 2024
Exploring the Head Architecture
Hydra RUP
Hydra Profx NLP
Vicure
22
Vican
Model
Figure 6: Performance comparison on MT-Bench of the standard MLP only Hydra head architecture and the PrefixMLP head architecture, which introduces an additional decoder layer. The PrefixMLP Hydra head architecture outperforms a standalone MLP Hydra head.
A.2 Exploring Alternative Hydra Head Architectures
Hydra-specific prefix attention. For MLP-based Hydra heads, the only representation of the gready-gennrated sequence passed to the heads input is the base model's hidden state corresponding to the most recently processed token. However, as the base model is trained independently of the Hydra heads, it is not obvious whether sufficient information regarding the context is encoded in the last token's hidden state. To better aggregate relevant information over the entire context for use by the Hydra heads, we propose to extend the base LLM with an additional decoder layer which is used solely to produce better input representations for the draft model. While each Hydra head is still single layer MLP, they each now take input the additional decoder layer's s representation of the final token in the already generated sequence. As the additional decoder layer is trained in conjunction with the Hydra heads, it can learn what information from the already generated sequence is useful for the Hydra heads. We note that Hydra heads share the same additional decoder layer hidden state, meaning the additional decoder layer is only queried once per Hydra decoding step. We refer to the resulting Hydra head architecture consisting of an additional decoder layer along with the standard MLP the PrefixMIP Hydra head architecture.
A.2.1 Results.
To test whether adding an additional Hydra-specific decoder layer improves modelling performance, we compare decoding with our proposed PrefixMLP architecture to decod- ing using the standard MLP-only Hydra head (Figure(6). We find that the decoding with PrefixMLP heads improves the average acceptance length by 1.12x which leads to an im- provement in average decoding throughput of 1.08x. This result suggests that aggregating context from the generated sequence in a Hydra head aware manner improves Hydra decoding performance.
B Discovering Performant Decoding Trees
We plot the results from searching for the optimal tree at each batch size investigated in Figure 7 Figure and Figure for Medusa, Hydra, and Hydra++ decoding respectively. For all decoding strategies evaluated, as the batch size increases the size of the tree which maximizes throughput decreases as well.
C Eagle Decoding
Concurrently to work, et Similar to existing speculative decoding techniques based on draft heads, the draft model in EAGLE leverages the base model's hidden states as input. However, EAGLE does not
et al. (2024) have proposed the EAGLE decoding framework.
14

---

Published as a conference paper at COLM 2024
Medusa 7b Decoding Tree Performance
Batch Size 1
Batch Size
Best Tree Batch Size
Best Thee Batch Size
Best Tree Batch Size
Best Tree Batch Size
CTok
259
Nodes
Figure 7: Results from determining the optimal decoding tree for Medusa with a 7B base model. Each point represents the throughput measured when using the tree that maximizes the expected accepted length for a given tree size. For each setting considered, the point marked by star denotes the tree size that achieves the greatest throughput.
Hydra 7b Decoding Tree Performance
Best Tree Batch Size
Best Tree Batch Size
Batchh Size 2
Batch Size
Tree Batch Size
Best Thee Batch
200-
Throughput
190
Figure 8: Results from determining the optimal decoding tree for Hydra with 7B base model. Each point represents the throughput measured when using the tree that maximizes the expected accepted length for given tree size. For each setting considered, the point marked by star denotes the tree size that achieves the greatest throughput.
collection of draft heads and instead uses singular draft head as the draft model. Similar to our work, EAGLE introduces sequential dependence to their draft head. Concretely, the EAGLE draft head is structured as transformer decoder layer which takes as input both the hidden states and input embeddings of the entire sequence. During each step of generating a candidate continuation, the EAGLE draft model not only predicts the next token the continuation, but also predicts estimate of the hidden state that the base model would have computed for that candidate token. EAGLE then extends the input to its draft head with both the input embedding of the predicted token, well as the estimated hidden state.
15

---

Published as a conference paper at COLM 2024
+ydra++ 7b Decoding Tree Performance
Balch Size 1
Batchh Size
Best Tree Batch Size
Best Thee Batch Size
Best Tree Batch Size
Best Tree Batch Size
400
Throughput 300
Nodes
Figure 9: Results from determining the optimal decoding tree for Hydra++ with a 7B base model. Each point represents the throughput measured when using the tree that maximizes the expected accepted length for a given tree size. For each setting considered, the point marked by star denotes the tree size that achieves the greatest throughput.
Hydra vs. Eagle Decoding
1.00x
1.01x
4-2
1.0cc
(lok)
Avg. Throughput (Tak (5
2.5
Eagle
Hydra++
40
2.3
Figure 10: Comparison of Hydra++ and EAGLE. We find that both draft heads achieve comparable throughput.
Li al. (2024) demonstrate that EAGLE decoding provides speedup relative to Medusa's sequentially-indeeeddee draft heads. Given that EAGLE was developed entirely indepen- dently of Hydra, we believe that Hydra and EAGLE, taken together, constitute valuable evidence that the benefits of sequential dependence in speculative decoding are robust and replicable.
We independently train and evaluate FAGLE draft heads using Vicuna 7B as the base model. We compare the throughput and acceptance length of EAGLE and Hydra++ in Figure [0) We find that while EAGLE achieves a higher average acceptance length, both EAGLE and Hydra++ achieve comparable decoding throughput. We attribute this to the added overhead of EAGLE draft heads, as they require querying full self-attention block for each position in the candidate continuation, whereas Hydra++ only queries additional self-attention block once per decoding step, with the rest of the computation in the Hydra++ draft model being done by shallow MI.Ps.
16

---

Published as a conference paper at COLM 2024
Model
Prefix Attention
Head
Head
Head
Head
Medusa
0.3
0.3
0.3
0.3
Hydra++
1.2
0.6
1.4
1.2
0.7
Breakdown of overhead during speculative decoding in milliseconds. We report the time spent performing prefix attention and decoding for each speculative head. As can be seen, Hydra decoding incurs greater overhead than Medusa decoding. Despite the overhead, Hydra decoding leads to end-to-end throughput improvements over Medusa.
Table 1:
Model MT Chat
Translation
Summary
QA
Math
RAG
Avg.
Medusa
2.00x
1.68 2.02x
1.59 x
1.67 x
1.98x
1.54x
1.75x
Hydra++
2.52x
1.89x
2.08x
2.59x
1.87x
2.17x
Table 2: Relative improvement in decoding throughput over standard autoregressive decoding for both Medusa and Hydra++ the SpecBench evaluation suite. We find that across all task categories, Hydra++ achieves significantly better throughput than Medusa decoding.
D Analysis of Hydra Head overheads
To better understand the efficiency gains from Hydra decoding, we analyze the overhead in- troduced by both: the additional decoder layer from prefix attention and 2) the sequential dependence of Hydra heads. Namely, to introduce sequential dependence, Hydra heads additionally take input the embeddings of the preceding speculated tokens. As such, the first layer of Hydra head has greater number of input features than a Medusa head, and the number of input features grows for later Hydra heads. We report the average time spent performing prefix attention and the time spent in each speculative decoding head for both Medusa and Hydra++ at the 8B model scale with batch size one in Table [ To contextualize the time spent performing speculative decoding, the average time to perform decoding step through the base model is 28 milliseconds. While Hydra decoding incurs additional overhead as compared to Medusa, Hydra still achieves an end-to-end throughput improvement as improves the average acceptance length.
E
Evaluation on SpecBench
In addition to MT-Bench, we evaluate the performance of Hydra++ on the SpecBench (Xia et al. 2024) speculative decoding evaluation suite. While MT-Bench is chat based bench- mark, SpecBench includes broader range of tasks to test the performance of speculative decoding methods in variety of settings. Specifically, SpecBench is composed of multi-turn chat (MT Chat), translation, summary, QA, math, and retrieval augmented generation (RAG) tasks. We report the relative improvement in throughput over autoregressive decoding for both Medusa and Hydra++ 8B at batch size one in Table
We find that Hydra++ significantly outperforms Medusa across all tasks in SpecBench. Averaged across all tasks, Hydra++ achieves an improvement in throughout over Medusa of 1.24x. The speedup we observe over Medusa on SpecBench very similar to the speedup measured on MT-Bench (1.27 x), suggesting that the gains from Hydra++ relative to Medusa generalize to a broad range of tasks. For both Medusa and Hydra++, the summary and RAG tasks the smallest improvement in throughput over the baseline. An important area for further work is to investigate whether the speedup on these tasks can be improved from increasing the amount of summarization and RAG data seen during training, or whether the reduced performance from speculative decoding is inherent to the task.
17


Reviews:
[{"review_id": "OZs9SUsqoB", "reviewer": "Reviewer_ngjH", "summary": "The paper introduces Hydra heads to introduce sequential dependency for the draft heads based on Medusa decoding. The paper conducts experiments using the Vicuna family of models to evaluate the throughput and token acceptance length compared to autoregressive decoding and Medusa and shows that the proposed method can improve on both metrics. The paper also investigates other aspects of the proposed Hydra methods by incorporating some recent techniques such as teacher loss and PrefixMLP to improve the framework further and thus presents Hydra++, which shows better performance with the combination of these techniques.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Could we include a cost analysis for the overhead of using Hydra draft heads?\n\nWe thank reviewer ngjH for the suggestion, and agree that including a discussion of the overhead of running the draft model would improve the paper. We have measured the overhead of each draft head in our 7B Medusa and Hydra++ models on MT-Bench, and obtained these results:\n\n**Absolute time (milliseconds) spent running draft heads per decoding step:**\n| Model | Extra Decoder Layer | Head 1 | Head 2 | Head 3 | Head 4 |\n|-|-|-|-|-|-|\n|Medusa|-|0.3|0.3|0.3|0.3|\n|Hydra++|1.2|0.6|1.4|1.2|0.7|\n\nFor reference, a single forward pass through the base model during the verification step takes 28 milliseconds.\n\nOne might expect that later Hydra++ heads would always have more overhead than earlier heads, because their input dimension is larger. However, later layers of our speculation tree can include fewer nodes than earlier layers, providing a counter-balancing effect on run time.\n\n# What is the intuition behind e.g. adding noise?\n\nInspired by the Neftune paper, our original motivation for adding noise to the embeddings was as a form of data augmentation / regularization, to allow us to train for more epochs without overfitting. Although the Neftune paper focused on the instruction-tuning setting, our hope was that the regularization benefits it found might apply to use cases beyond instruction-tuning. We would like to clarify that we only experimented with adding noise to embeddings at training time, not during inference. We will make all of these points more clear in our camera-ready version.\n\nIf there are any other aspects of our draft head recipe experiments which the reviewers would like us to clarify, we would be happy to address them.\n\n# What about sampling with higher temperatures?\n\nWe would be happy to include throughput results for temperature sampling at nonzero temperatures in the final paper.\n\nWe note that in section 6.3 we already investigate \"typical-acceptance\" sampling, which is non-greedy and stochastic. We find that Hydra++ achieves commensurate quality to temperature sampling from the base model, while maintaining high candidate acceptance lengths. We present results for typical sampling in Figure 4 of our submission.\n\n# Where are the confidence intervals in Figures 2 and 3?\n\nWe plot confidence intervals for all throughput results in the paper, although in some cases the error is so small that they can be difficult to see.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The presentation is clear and easy to understand. The motivation for incorporating sequential dependence to Medusa decoding is simple but seems effective from the experiments.\n2. The studied problem is important and impactful for current LLM practices for decoding efficiencies.\n3. The exploration of other techniques is quite extensive seeking alternative implementations of Hydra heads. The results and observations can provide useful insights for practitioners.", "reasons_to_reject": "1. It would be good to have a cost and efficiency analysis for the proposed method since taking the previous tokens to the draft heads seems to scale with the number of heads.\n2. The optimal draft head recipe section is not comprehensive enough to understand the intuitions behind some observations. For example, could the authors elaborate more on why adding noise to Hydra did not improve the performance of the models? The motivation to add noise here is unclear since the referenced paper is about using noisy embeddings for instruction tuning. It would be better to see if adding noise to the speculative decoding generally fails or if there are other observations for the relationship between noisy embeddings and decoding performance.\n3. Other than greedy decoding, Some results on other temperatures for evaluation would be interesting to see how randomness would affect the observations.", "questions_to_authors": "1. See the weakness above.\n2. For Figures 2 and 3, why is there a confidence interval plotted (if I understand correctly) for the Hydra only on Avg. Throughput results? Please correct me if I miss anything.\n3. Typo: the word \"using\" is repeated twice in the 2nd line of section 7.", "extra_scores": {}}, {"review_id": "nIEL3Uh49C", "reviewer": "Reviewer_YfQn", "summary": "This paper proposed a new speculative decoding scheme based on Medusa. This is different from Medusa which only considers the LLM's hidden state as the input and predicts the future tokens, this paper proposed to take the hidden stated and predicted candidates as input, which enhances the \"dependencies\" rather than fully factorized. Additionally, authors explored coupled improved model architecture and the training objectives,  leading to the \"Hydra++\". Empirical results showed a two-digit margin increase w.r.t baselines regarding averaged accepted drafts, and throughput. In general, I think this is an elegant idea and tailored improvements to the model and objectives showed consistent gains.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# What about other sampling strategies (e.g. top-p)?\n\nWe would be happy to include throughput results for stochastic sampling using other sampling strategies (e.g. top-p) in the final paper.\n\nWe note that in section 6.3 we already investigate \"typical-acceptance\" sampling, which is non-greedy and stochastic. We find that Hydra++ achieves commensurate quality to temperature sampling from the base model, while maintaining high candidate acceptance lengths. We present results for typical sampling in Figure 4 of our submission.\n\n# Could we add concurrent work to the related work section?\n\nThe three concurrent works we are aware of are [1], [2], and [3]. We already include a comparison to [1] in the related work section and Appendix C of our submission. The preprints [2] and [3] were released 2 weeks before and 4 weeks after the final COLM submission deadline, respectively. We would be happy to include references to these preprints in our related work section.\n\n[1] Li, Yuhui, et al. \"Eagle: Speculative sampling requires rethinking feature uncertainty.\" arXiv:2401.15077 (2024).\n\n[2] Zhang, Aonan, et al. \"Recurrent drafter for fast speculative decoding in large language models.\" arXiv:2403.09919 (2024).\n\n[3] Wertheimer, Davis, et al. \"Accelerating Production LLMs with Combined Token/Embedding Speculators.\" arXiv:2404.19124 (2024).\n\n# Could we evaluate on benchmarks beyond MT-bench?\n\nWe thank reviewer YfQn for the suggestion. In response to your review, we have evaluated Hydra++ on the multi-task SpecBench benchmark [4]. SpecBench includes multi-turn conversation, summarization, RAG, translation, question-answering, and mathematical reasoning tasks. We evaluated our 7B Medusa and Hydra++ models on SpecBench, as well as a Speculative Sampling [5] baseline. We find that the relative gains from Hydra++ over Medusa generalize to SpecBench:\n\n**Throughput relative to non-speculative decoding by SpecBench task:**\n|Method|Multi-Turn Conversation|Translation|Summarization|QA|Math Reasoning|RAG|Overall|\n|-|-|-|-|-|-|-|-|\n|Speculative Sampling|1.66×|1.13×|1.62×|1.49×|1.46×|1.54×|1.49×|\n|Medusa|2.00×|1.68×|1.59×|1.67×|1.98×|1.54×|1.75×|\n|Hydra++|2.52×|2.02×|1.89×|2.08×|2.59×|1.87×|2.17×|\n\n[4] Xia, Heming, et al. \"Unlocking efficiency in large language model inference: A comprehensive survey of speculative decoding.\" arXiv:2401.07851 (2024).\n\n[5] Chen, Charlie, et al. \"Accelerating large language model decoding with speculative sampling.\" arXiv:2302.01318 (2023).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The idea itself is well-motivated and somewhat elegant to me. The predicted tokens are at hand information and could be used. It feels like adding an \"autoregressiveness\" in the draft heads.  I would argue the idea is simple yet effective.\n\n- Authors also explored the design space of the model architectures, sampling methods, and training objectives. The design choices are well illustrated and motivated. And the final \"hydra++\" is pretty well-tuned. \n\n- The evaluation is comprehensive and empirical results are showing consistent gains even w.r.t. strong baselines. \n\nAdditionally, the paper is pretty easy to follow and mathematical explanations are clear. \n\nIn general, this paper is clear and solid with well-justified motivation and clean techniques. Empirical results also demonstrated gains.", "reasons_to_reject": "- Sample methods: I would encourage authors to consider more sampling methods, for example, top-p. \n- Related work might consider adding some concurrent work?", "questions_to_authors": "- Will the different sampling methods affect the gain over acceptance? \n\n- What are the other benchmarks that might be a good fit instead of MT-bench?", "extra_scores": {}}, {"review_id": "7thk12RlCu", "reviewer": "Reviewer_71U9", "summary": "This paper works on improving the drafting accuracy of speculative decoding. Authors mention a important problem that in the Medusa like speculative decoding methods, the draft heads make predictions without considering earlier tokens in the current candidate continuation. To solve this problem, they propose a drop-in sequentially dependent method to standard draft heads. This method increases the average candidate continuation acceptance length and further accelerates LLM inference.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "# Could we compare to Eagle and Speculative Sampling?\n\nAlthough Eagle [1] is concurrent work to ours, we include a comparison with Eagle in our related work section and in Appendix C of our submission.\n\nFollowing your suggestion to compare with Speculative Sampling [2], we have evaluated Speculative Sampling, Medusa, and our Hydra++ model under the SpecBench [3] evaluation harness, the results of which we present in the next subsection.\n\n# Could we evaluate Hydra on more benchmarks?\n\nWe thank reviewer 71U9 for the suggestion to evaluate on SpecBench. We have benchmarked our 7B Medusa baseline and Hydra++ model, as well as Speculative Sampling [2], on the SpecBench [3] benchmark. We find that the relative gains from Hydra++ over Medusa generalize to this larger benchmark, and that Hydra compares favorably to Speculative Sampling [2]. We would be happy to include these results in the final paper.\n\n**Throughput relative to non-speculative decoding by SpecBench task:**\n|Method|Multi-Turn Conversation|Translation|Summarization|QA|Math Reasoning|RAG|Overall|\n|-|-|-|-|-|-|-|-|\n|Speculative Sampling|1.66×|1.13×|1.62×|1.49×|1.46×|1.54×|1.49×|\n|Medusa|2.00×|1.68×|1.59×|1.67×|1.98×|1.54×|1.75×|\n|Hydra++|2.52×|2.02×|1.89×|2.08×|2.59×|1.87×|2.17×|\n\n[1] Li, Yuhui, et al. \"Eagle: Speculative sampling requires rethinking feature uncertainty.\" arXiv:2401.15077 (2024).\n\n[2] Chen, Charlie, et al. \"Accelerating large language model decoding with speculative sampling.\" arXiv:2302.01318 (2023).\n\n[3] Xia, Heming, et al. \"Unlocking efficiency in large language model inference: A comprehensive survey of speculative decoding.\" arXiv:2401.07851 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. The motivation is clear. The method is consistent with the motivation. The sequentially dependent method is novel and naturally considers earlier tokens in the current candidate continuation. \n\n2. The experiments demonstrate that the proposed method significantly outperforms the baselines (2.70× acceleration).\n\n3. Authors explore a lot on the architecture and training objective (section 3.1) . I especially appreciate the detailed analysis and discussions on this method.", "reasons_to_reject": "The paper mainly compares with Medusa but lacks comparison with other speculative decoding methods. [1][2]\n\n[1] Li, Yuhui, et al. \"Eagle: Speculative sampling requires rethinking feature uncertainty.\" arXiv preprint arXiv:2401.15077 (2024).\n\n[2] Chen, Charlie, et al. \"Accelerating large language model decoding with speculative sampling.\" arXiv preprint arXiv:2302.01318 (2023).", "questions_to_authors": "A suggestion:  The task influences the speculative decoding method a lot, but the paper only evaluates on several benchmarks. According to this leaderboard (https://github.com/hemingkx/Spec-Bench), Hydra achieves a good performance on a wide range of tasks. It would be beneficial if the authors could provide additional experimental results to demonstrate the effectiveness of the method across different tasks.", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a technique (Hydra) to train multiple prediction heads in an autoregressive language models to speed up speculative decoding. By having different heads depend on each other, the proposed technique improves on Medusa (where the heads are independent). Careful choice of the training objective and head architectures yield up to 1.3x compared to Medusa (and 2.7x compared to standard decoding).
+ The design is well-motivated and elegant.
+ Strong empirical results on speculative decoding speedup.
Overall the paper is a valuable contributed to the area of speculative decoding.
Decision: Accept