# Hydra: Sequentially-Dependent Draft Heads for Medusa Decoding

Zachary Ankner ∗,1,2 Rishab Parthasarathy∗,1 Aniruddha Nrusimha1
Christopher Rinard1 Jonathan Ragan-Kelley1 William Brandon1
1MIT 2MosaicML

###### 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 light-weight 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 speculation accuracy. Decoding with Hydra heads improves throughput
compared to Medusa decoding with standard draft heads. 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 $1.31\times$ and $2.71\times$ compared to Medusa
decoding and autoregressive decoding, respectively. Overall, Hydra heads are a
simple intervention on standard draft heads that significantly improve the
end-to-end speed of draft head based speculative decoding.

##  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 is able to achieve
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 bandwidth-
bound_ 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 weights, but only involves a
comparatively small number of FLOPs (processing just one token for each
sequence in the batch), LLM decoding tends to underutilize the GPU’s abundant
capacity for floating-point computation.

To mitigate the memory bandwidth bottleneck in sequential LLM decoding, recent
research has investigated accelerating LLM inference through _speculative
decoding_. Speculative decoding uses a smaller _draft model_ to propose a
multi-token candidate continuation of the current sequence on each generation
step. The original LLM then _verifies_ all tokens in the candidate
continuation in parallel, appending some subset of them to the sequence and
discarding the rest. Because each verification step requires only a 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.

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-trained language models, Stern
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
operates on the original LLM hidden states as _draft heads_. In the draft head
paradigm, each draft head is responsible for predicting the identity of a
token a fixed number of steps into the future.

All draft heads to date make predictions only as a function of the base
model’s hidden states from previously verified tokens, making them unaware of
earlier tokens in the current candidate continuation. Specifically, when
decoding the $i^{th}$ token in the continuation, a standard draft head is not
conditioned on what previous draft heads speculatively generated for the
$1^{st},2^{nd},...,(i-1)^{th}$ tokens in the 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 improve token prediction
accuracy and end-to-end decoding throughput. To construct sequentially
dependent draft heads, we set each head’s output 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 to 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.1\times$ better throughput than Medusa decoding.

In addition to proposing Hydra heads, we further explore the design space of
their training objective and architecture. We investigate adding noise to the
base model’s hidden states (Jain et al., 2024), using a teacher loss
objective, and adding an extra transformer decoder layer to the Hydra heads.
We find that the best combination of these techniques (teacher loss and extra
decoder layer) achieves up to $1.31\times$ and $2.7\times$ higher throughput
than standard Medusa decoding and regular autoregressive decoding,
respectively.

##### Contributions

In this work we present the following contributions:

  * •

We analyze the standard formulation of draft heads and observe that they are
sequentially independent during decoding. We propose Hydra heads as a
sequentially dependent alternative, that increase end-to-end decoding
throughput by up to $1.1\times$ as compared to Medusa decoding (section 5).

  * •

We perform a thorough empirical exploration of the design space of Hydra
heads. Using a teacher loss and adding an extra decoder layer to the Hydra
heads, we are able to further increase decoding throughput by up to
$1.31\times$ and $2.7\times$ over standard Medusa decoding and autoregressive
decoding respectively. (section 6, section 7).

##  2 Background

In this section we provide background details on both speculative decoding and
Medusa decoding.

##### Speculative decoding.

_Speculative decoding_ (Stern et al., 2018; Leviathan et al., 2023; Chen et
al., 2023) provides a 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_. On each step of
speculative decoding, we use the draft model to generate 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_{\leq t}$ be
the sequence that has been generated so far and fixing some speculation length
$K$, we query the draft model’s joint distribution
$p_{\text{draft}}(x_{t+1},\ldots,x_{t+K}\mid x_{\leq t})$ to generate a
candidate continuation $\hat{x}_{t+1},\ldots,\hat{x}_{t+K}$. We then invoke
the base model on the candidate continuation to compute the conditional
probabilities: $p_{\text{base}}(\hat{x}_{t+1}\mid x_{\leq
t}),\ldots,p_{\text{base}}(\hat{x}_{t+K}\mid x_{\leq
t},\hat{x}_{t+1},\ldots,\hat{x}_{t+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 chooses some prefix
$\hat{x}_{t+1},\ldots,\hat{x}_{t+n_{\text{accept}}}$ 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 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 a 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 _tree_ of candidate continuations, rather than a single candidate
continuation (Miao et al., 2023; Spector & Re, 2023; Cai et al., 2024). Nodes
of this candidate tree correspond to candidate tokens, and the children of a
node represent different possible tokens which might follow it in the
continuation. Each path along the tree thus represents a different candidate
continuation. To populate a node with $m$ children, we query the draft model
for the $m$ most likely tokens that might follow it, conditioned on the
generated sequence 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, the structure of the
tree is fixed at design time, so that the number of children $m$ at each
position in the tree does not depend on any runtime data.

After generating a candidate continuation tree from 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 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 decoding is an
independently trained language model, Stern et al. (2018) define the draft
model as a collection of lightweight heads that take as input the base model’s
hidden state, which we refer to as the _draft heads_. 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_{\text{draft},1},...,f_{\text{draft},K}$
responsible for predicting the tokens $1,\ldots,K$ steps into the future. The
predictions of these heads are statistically independent of each other;
letting $x_{\leq t}$ denote the sequence generated so far, and letting
$h_{t-1}$ denote the last-layer hidden state of the token most recently
processed by the base model, Stern et al. compute their draft predictions on
each generation step as

$p_{\text{draft}}(x_{t+i}\mid x_{\leq t+i-1})=f_{\text{draft},i}(h_{t-1})$ (1)

Because the hidden state $h_{t-1}$ only depends on the tokens $x_{\leq t-1}$,
the joint distribution of all draft heads decomposes as simply the product
distribution of individual heads conditioned on $x_{\leq t-1}$:

$p_{\text{draft}}(x_{t+1},\ldots,x_{t+K}\mid x_{\leq
t})=p_{\text{draft}}(x_{t+1}\mid x_{\leq t-1})\cdots
p_{\text{draft}}(x_{t+K}\mid x_{\leq t-1})$

fig. 1 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 a tree
of candidates where the draft model is a collection of draft heads.

While Medusa decoding is compatible with any specific model architecture used
for each draft head $f_{\text{draft},i}$, Cai et al. (2024) choose to use a
single layer MLP with a residual connection.

##  3 Hydra Heads

In this section, we provide the technical details on constructing Hydra heads
by adding sequential dependence to standard draft heads.

##### Sequentially dependent draft heads.

Figure 1:  A visualization of generating a candidate continuation using
existing draft heads and using our Hydra draft heads. Lines going into a head
represent inputs to the draft head. While the only input to existing draft
heads is the base model’s last-layer hidden state for the most recently
processed token, Hydra heads leverage earlier tokens in the candidate
continuation as additional inputs.

The key observation behind Hydra heads is that there is no sequential
dependence in standard draft heads, i.e., each draft head makes predictions
independently. A draft model defined by a collection of draft heads speculates
the identity of the $i^{th}$ future token (eq. 1) as
$f_{\text{draft},i}(h_{t-1})$. However, $h_{t-1}$ is only a function of the
already generated sequence $x_{\leq t-1}$. Thus, when using draft heads:

$p_{\text{draft}}(\hat{x}_{t+i}|x_{\leq
t},\hat{x}_{t+1},\dots,\hat{x}_{t+i-1})=p_{\text{draft}}(\hat{x}_{t+i}|x_{\leq
t-1})$

Intuitively, this means that there is no sequential dependence between draft
heads: when we use a draft head to speculate the $i^{th}$ token in a candidate
continuation, it is unaware of the $1^{st},2^{nd},...,(i-1)^{th}$ tokens in
the candidate continuation.

We propose Hydra heads, which are sequentially dependent draft heads. Hydra
heads are sequentially dependent as they are a function of both the base
model’s hidden state up to time $t$ as well as the input embeddings of the
tokens sampled by previous Hydra heads. Namely, the draft model is now a
collection of Hydra heads $\\{f_{\text{Hydra},1},...,f_{\text{Hydra},K}\\}$
and the $i^{th}$ future token’s distribution is parameterized by this
collection of Hydra heads as:

$p_{\text{draft}}(\hat{x}_{t+i}|x_{\leq
t},\hat{x}_{t+1},\dots,\hat{x}_{t+i-1})=f_{\text{Hydra},i}(h_{t-1},x_{t},\hat{x}_{t+1},...,\hat{x}_{t+i-1})$
(2)

where $h_{t-1}$ is again the base model’s hidden state of the final token in
the already decoded sequence. The sequential dependency of Hydra heads v.s.
standard draft heads is visualized in fig. 1.

We use the term _Hydra Decoding_ to refer to speculative decoding with tree
candidates and Hydra heads.

As with Medusa, the framework of Hydra decoding is compatible with any choice
of model architecture used to implement $f_{\text{Hydra},i}$. The most basic
architecture, which we use for our direct comparisons with Medusa, is simply a
single hidden layer MLP whose input is the hidden state $h_{t-1}$ concatenated
with the input embeddings of the proceeding tokens in the candidate
continuation $E_{x_{t}},E_{\hat{x}_{t+1}},...,E_{\hat{x}_{t+i-1}}$, where the
concatenation is performed along the feature dimension. We chose to use this
architecture as it is most comparable to the baseline Medusa architecture;
however, we explore alternative Hydra head architectures in section 6.2.

##  4 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.

##### Models.

We build on the Vicuna family of models (Chiang et al., 2023), which are
conversation-finetuned LLaMa models (Touvron et al., 2023), as the base models
for our speculative decoding experiments. We consider Vicuna 7B, 13B, and 33B.

##### Training.

While draft heads can be trained in conjuction 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), a collection of multi-turn
conversations. Training is performed on $8\times$ Nvidia A100-80gb and
conducted using the HuggingFace Trainer (HuggingFace, ). We use a cosine
learning rate schedule with warmup (Loshchilov & Hutter, 2017) and the AdamW
optimizer (Loshchilov & Hutter, 2019) with parameters
$\beta_{1}=0.9,\beta_{2}=0.999$.

##### Evaluation.

All evaluations are performed on MT-Bench (Zheng et al., 2023), a multi-turn
conversation benchmark. Unless otherwise specified, experiments are conducted
using the greedy verification criterion; since there is no stochasticity in
the greedy sampling procedure, we don’t report the quality of generations as
they are identical to the base model. Instead, we report the average
_throughput_ , 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 a single A100-80gb GPU.

##  5 Head to Head Comparison of Medusa and Hydra

In this section we perform a head to head comparison of the decoding speeds of
Medusa and Hydra.

##### Experiment setup.

In addition to the setup detailed in section 4, all heads are trained with an
initial learning rate of 1e-3, a training duration of 1 epoch, and a batch
size of 128 as done in Cai et al. (2024). We consider Medusa / Hydra schemes
consisting of 4 draft heads, each using a single hidden layer MLP with a skip
connection as its architecture.

###  5.1 Results

Figure 2:  Performance comparison on MT-Bench of Hydra, Medusa, and the
baseline of autogressive decoding. Hydra heads increase decoding throughput
and average acceptance length compared to all other methods.

The hypothesis that motivated us to propose Hydra heads is that introducing
sequential dependence among draft heads should improve their prediction
quality, leading to greater decoding throughput. We test this hypothesis by
directly comparing our Hydra decoding with Medusa decoding as well as the
baseline of standard autoregressive decoding. We plot the results of our head-
to-head comparison in fig. 2. We find that across all base model sizes
evaluated, Hydra decoding achieve the greatest average acceptance length,
which leads to a significant improvement in decoding throughput. Specifically,
Hydra heads achieve a $2.36\times,2.17\times,$ and $2.15\times$ improvement in
throughput as compared to autoregressive decoding for the 7B, 13B, and 33B
parameter base models respectively. This translates to a throughput
improvement over Medusa decoding of $1.11\times,1.10\times,$ and $1.11\times$
for the 7B, 13B, and 33B parameter base models respectively. These results
demonstrate that making draft heads sequentially dependent significantly
improves the prediction quality of said heads and thus their decoding speed.

##  6 Exploring the Design Space of Hydra Heads

In this section we explore modifications to the training procedure and
architecture of Hydra heads.

##### Basic improvements to Hydra head hyperparameters.

Before ablating more complex aspects of Hydra head training, we slightly
adjust the training hyperparameters used by Cai et al. (2024). We use these
adjusted hyperparameters for training all Hydra heads for the rest of our
experiments. Firstly, as we noticed that Hydra head quality does not saturate
after 1 epoch, we train Hydra heads for 10 epochs instead of 1 epoch.
Secondly, to mitigate over-fitting we set the weight decay to be $0.1$ and add
dropout at a rate of $0.2$ to each hidden layer of the Hydra head MLP.
Finally, we increase the size of each Hydra head’s MLP to be 4 hidden layers
instead of 1. We chose 4 hidden layers as we observed diminishing gains in
prediction quality beyond that MLP depth.

###  6.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 can improve the resulting model’s performance.
Specifically, they sample noise $\epsilon\in\mathbb{R}^{B\times L\times
d}\sim\text{Uniform}(-1,1)$, scale it by a factor
$\frac{\alpha_{\text{noise}}}{\sqrt{Ld}}$, and then add the scaled noise to
the input embeddings, where $B$ is the batch size, $L$ is the sequence length,
$d$ is the model dimension, and $\alpha_{\text{noise}}$ is a hyperparameter
that controls the strength of the noise. We consider whether applying such
noise to the hidden states of the base LLM can also improve the performance of
the Hydra heads. For our experiments, we set $\alpha_{\text{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 is 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. As also employed by Zhou et al. (2024), we
investigate using a _teacher loss_ where each Hydra head’s training loss is
the cross entropy between its predicted distribution and the base model’s next
token distribution.

####  6.1.1 Results

Figure 3: Performance comparison on MT-Bench of different Hydra head training
objectives. Training based on a teacher loss leads to the most performant
Hydra heads.

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 fig. 3. 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 a $1.04\times$ 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.

###  6.2 Exploring Alternative Hydra Head Architectures

##### Hydra-specific prefix attention.

For MLP-based Hydra heads, the only representation of the already-generated
sequence passed to the heads as 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 a single
layer MLP, they each now take as input the additional decoder layer’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 all 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 as the _PrefixMLP_ Hydra head architecture.

####  6.2.1 Results.

Figure 4:  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.

To test whether adding an additional Hydra-specific decoder layer improves
modelling performance, we compare decoding with our proposed PrefixMLP
architecture to decoding using the standard MLP-only Hydra head (fig. 4). We
find that the decoding with PrefixMLP heads improves the average acceptance
length by $1.12\times$ which leads to an improvement in average decoding
throughput of $1.08\times$. This result suggests that aggregating context from
the generated sequence in a Hydra head aware manner improves Hydra decoding
performance.

##  7 Hydra++: The Most Performant Hydra Model

In this section we propose _Hydra++_ , our recipe for the most performant
Hydra heads, based on our observations from section 6. Specifically, Hydra++
heads are trained using using the base model teacher loss, as well as using
the PrefixMLP head architecture.

###  7.1 Evaluating Hydra++

To understand the cumulative impact of our interventions to Hydra heads on
Hydra decoding, we evaluate Hydra++ against the the baselines of standard
autoregressive decoding, Medusa decoding, as well as the original formulation
of Hydra decoding.

##### Results.

Figure 5:  Performance comparison on MT-Bench of Hydra++ (the best-performing
Hydra head recipe), original Hydra decoding, original Medusa decoding, as well
as the baseline of autoregressive decoding. Hydra++ outperforms all other
decoding schemes.

We plot the decoding performance of Hydra++ compared to the other baselines in
fig. 5. Hydra++ produces a significant speedup, improving decoding throughput
by $2.7\times,2.5\times,$ and $2.53\times$ as compared to autoregressive
decoding for Vicuna 7B, 13B, and 33B respectively. Compared to Medusa
decoding, Hydra++ improves throughput by $1.27\times,1.27\times,$ and
$1.31\times$ for Vicuna 7B, 13B, and 33B respectively. Finally, compared to
decoding using the original Hydra head formulation, Hydra++ improves decoding
throughput by $1.14\times,1.15\times,$ and $1.18\times$ for Vicuna 7B, 13B,
and 33B respectively. These results suggest that sequentially dependent draft
heads in conjunction with our other training and architecture interventions
are substantial improvement to standard draft heads.

###  7.2 Typical Acceptance Sampling

All results presented so far have used the greedy acceptance criteria, where
tokens are only verified if they match the greedy next-token prediction of the
base LLM. In this section, we consider non-greedy sampling via the _typical
acceptance_ verification criterion proposed by Cai et al. (2024).

##### Typical acceptance criterion.

The purpose of the _typical acceptance_ verification criterion (Cai et al.,
2024) is to sample more diverse and creative sequences than greedy acceptance,
while preserving the efficiency benefits of speculative decoding. While it is
also possible to increase generation diversity in the context of speculative
decoding by performing rejection resampling with a high temperature parameter,
this tends to degrade acceptance rates and end-to-end decoding speed, as noted
by Gante (2023) and Spector & Re (2023).

The typical acceptance criterion specifies that a speculated token
$\hat{x}_{t+i}$ is accepted if:

$p_{\text{base}}(\hat{x}_{t+i}|x_{\leq
t},\hat{x}_{t+1},\dots,x_{t+i-1};\tau)>\min(\epsilon,\alpha\exp(-H(p_{\text{base}}(\cdot|x_{\leq
t},\hat{x}_{t+1},\dots,x_{t+i-1};\tau))))$

where $\epsilon$ is known as the _posterior threshold_ , $\alpha$ is known as
the _posterior alpha_ , $\tau$ is the sampling temperature, and $H(\cdot)$ is
the entropy. Both $\epsilon$ and $\alpha$ are hyperparameters that should be
tuned empirically. Intuitively, this acceptance criterion biases the
generation towards more likely tokens while accepting more diverse
continuations when there is greater uncertainty in the base model’s
distribution. For a more detailed analysis on typical acceptance, we refer the
reader to Cai et al. (2024).

##### Setup.

We evaluate how different settings of $\epsilon$ and $\alpha$ affect both
acceptance length and generation quality. Following Cai et al. (2024), we
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 al., 2023). We fix the sampling temperature $\tau=0.7$, vary
the posterior threshold $\epsilon~{}\in~{}\\{0.05,0.1,0.15,0.2,0.25\\}$, and
set the posterior alpha as $\alpha=\sqrt{\epsilon}$.

##### Results.

Figure 6: An analysis of the average acceptance length and MT-Bench score of
Medusa heads and Hydra++ heads for varying posterior thresholds when
performing typical sampling. We observe that over all posterior thresholds,
Hydra++ decoding has a significantly average higher acceptance length than
Medusa. Additionally, typical sampling of Hydra heads can achieve the same MT-
Bench score as random sampling from the base model with the temperature
recommended by MT-Bench ($\text{temperature}=0.7$ for the tasks we consider).

We plot how varying the posterior threshold affects both the average
speculation length as well as the quality of the resulting generations in fig.
6. For both Medusa and Hydra++, increasing the posterior threshold slightly
decreases the average speculation length, but for all posterior thresholds
examined, Hydra++ has a significantly higher average acceptance length than
Medusa. We do not find any monotonic relation between the posterior threshold
and MT-Bench score for Hydra++. However, for the best setting of the posterior
threshold ($\epsilon=0.15$), Hydra++ achieves a better score on MT-Bench than
Medusa, and achieves performance comparable to random sampling from the base
model. These results demonstrate that using typical acceptance as the
verification scheme in Hydra decoding can allow for comparable performance as
sampling the base model while maintaining a high average speculation length.

##  8 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 et al.
(2023) and Chen et al. (2023), and anticipated in a restricted form by Stern
et al. (2018). Recent work has explored alternatives to standard draft models
such as using retrieval mechanisms to propose continuations (He et al., 2023),
and reformulating language model sampling in terms of Jacobi iteration (Fu et
al., 2023). The alternative draft model architecture that we examine in this
work, draft heads, was first proposed by Stern et al. (2018). Another
direction of speculative decoding research has investigated verifying a tree
of candidate continuations rather than a 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 a hierarchy of draft models, with each aiding
speculative decoding for the next. Other contemporary 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 et al., 2024). We would also like to acknowledge the
concurrent work EAGLE (Li et al., 2024) which is the work most similar to
ours. We discuss EAGLE in appendix A.

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., 2024; Sheng et al., 2023). To decrease the memory
footprint of the KV-cache, Shazeer (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 a 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 at which inference is performed is another technique
for improving LLM inference throughput. Yu et al. (2022) propose iteration-
level scheduling for batched decoding, where executions are scheduled at the
granularity of each decoding step instead of at the granularity of an incoming
request. Additionally, Kwon et al. (2023) develop the vLLM serving system
which utilizes their proposed PagedAttention method for improving the memory
management of the KV-cache and that allows for greater inference batch sizes.
Contrary to these approaches, speculative decoding primarily targets the
small-batch inference regime, in which memory bandwidth constraints dominate
and there are abundant unused FLOPs available to spend on parallel
verification under the base model.

##  9 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-proposed draft heads are sequentially
independent, meaning that when speculating $i$ tokens into the future, the
draft head’s prediction is not conditioned on the
$1^{st},2^{nd},...,(i-1)^{th}$ speculated tokens. To fix this problem, we
propose Hydra heads: a drop-in, sequentially dependent 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.11\times$ improvement in end-to-end
throughput as compared to Medusa decoding. We also investigate different
training objectives and architectures for Hydra heads. We find that training
based on a teacher loss improves head quality and end-to-end decoding speed.
We experiment with a Hydra head architecture we call _PrefixMLP_ that better
encodes the base model’s hidden states by introducing an additional decoder
layer. Using PrefixMLP Hydra heads further increases decoding speed. Combining
all Hydra head improvements into a Hydra head recipe we call _Hydra++_ , we
increase decoding throughput by up to $1.31\times$ and $2.7\times$ as compared
to Medusa and autoregressive decoding respectively. Draft head based
speculative decoding is an efficient and simple alternative to the standard
speculative decoding paradigm, and our work takes an important step towards
maximizing the performance of decoding with draft heads through the
construction of sequentially dependent draft heads.

## Acknowledgments

This work was initially performed as a 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

  * Ainslie et al. (2023) 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.18653/v1/2023.emnlp-main.298.  URL https://aclanthology.org/2023.emnlp-main.298.
  * Alizadeh et al. (2024) Keivan Alizadeh, Iman Mirzadeh, Dmitry Belenko, Karen Khatamifard, Minsik Cho, Carlo C Del Mundo, Mohammad Rastegari, and Mehrdad Farajtabar.  Llm in a flash: Efficient large language model inference with limited memory, 2024.
  * Cai et al. (2024) Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao.  Medusa: Simple llm inference acceleration framework with multiple decoding heads, 2024.
  * Chen et al. (2023) Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper.  Accelerating large language model decoding with speculative sampling, 2023.
  * Chiang et al. (2023) Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, Ion Stoica, and Eric P. Xing.  Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023.  URL https://lmsys.org/blog/2023-03-30-vicuna/.
  * Dettmers et al. (2022) Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer.  Llm.int8(): 8-bit matrix multiplication for transformers at scale.  _arXiv preprint arXiv:2208.07339_ , 2022.
  * Frantar & Alistarh (2023) 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.
  * Frantar et al. (2023) Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh.  OPTQ: Accurate quantization for generative pre-trained transformers.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=tcbBPnfwxS.
  * Fu et al. (2023) Yichao Fu, Peter Bailis, Ion Stoica, and Hao Zhang.  Breaking the sequential dependency of llm inference using lookahead decoding, November 2023.  URL https://lmsys.org/blog/2023-11-21-lookahead-decoding/.
  * Gante (2023) Joao Gante.  Assisted generation: a new direction toward low-latency text generation, May 2023.  URL https://huggingface.co/blog/assisted-generation.
  * He et al. (2023) Zhenyu He, Zexuan Zhong, Tianle Cai, Jason D Lee, and Di He.  Rest: Retrieval-based speculative decoding, 2023.
  * (12) HuggingFace.  Hugging face trainer.  URL https://huggingface.co/docs/transformers/main_classes/trainer.
  * Jain et al. (2024) 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 Representations_ , 2024.  URL https://openreview.net/forum?id=0bMmZ3fkCk.
  * Kwon et al. (2023) Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica.  Efficient memory management for large language model serving with pagedattention, 2023.
  * Leviathan et al. (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.
  * Li et al. (2024) Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang.  Eagle: Speculative sampling requires rethinking feature uncertainty, 2024.
  * Liu et al. (2023a) Xiaoxuan Liu, Lanxiang Hu, Peter Bailis, Ion Stoica, Zhijie Deng, Alvin Cheung, and Hao Zhang.  Online speculative decoding, 2023a.
  * Liu et al. (2023b) Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Ré, and Beidi Chen.  Deja vu: contextual sparsity for efficient llms at inference time.  In _Proceedings of the 40th International Conference on Machine Learning_ , ICML’23. JMLR.org, 2023b.
  * Loshchilov & Hutter (2017) Ilya Loshchilov and Frank Hutter.  SGDR: Stochastic gradient descent with warm restarts.  In _International Conference on Learning Representations_ , 2017.  URL https://openreview.net/forum?id=Skq89Scxx.
  * Loshchilov & Hutter (2019) Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  In _7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_. OpenReview.net, 2019.  URL https://openreview.net/forum?id=Bkg6RiCqY7.
  * Miao et al. (2023) 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.
  * Sanh et al. (2020) Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf.  Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter, 2020.
  * ShareGPT (2023) ShareGPT.  Sharegpt, 2023.  URL https://huggingface.co/datasets/Aeala/ShareGPT_Vicuna_unfiltered.
  * Shazeer (2019) Noam Shazeer.  Fast transformer decoding: One write-head is all you need, 2019.
  * Sheng et al. (2023) Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Ré, Ion 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://proceedings.mlr.press/v202/sheng23a.html.
  * Spector & Re (2023) Benjamin Frederick Spector and Christopher Re.  Accelerating LLM inference with staged speculative decoding.  In _Workshop on Efficient Systems for Foundation Models @ ICML2023_ , 2023.  URL https://openreview.net/forum?id=RKHF3VYjLK.
  * Stern et al. (2018) Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit.  Blockwise parallel decoding for deep autoregressive 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 https://proceedings.neurips.cc/paper_files/paper/2018/file/c4127b9194fe8562c64dc0f5bf2c93bc-Paper.pdf.
  * Touvron et al. (2023) 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.
  * Xiao et al. (2023) Guangxuan Xiao, Ji Lin, Mickael Seznec, Hao Wu, Julien Demouth, and Song Han.  SmoothQuant: Accurate and efficient post-training quantization for large language models.  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. PMLR, 23–29 Jul 2023.  URL https://proceedings.mlr.press/v202/xiao23c.html.
  * Yu et al. (2022) Gyeong-In Yu, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, and Byung-Gon Chun.  Orca: A distributed serving system for Transformer-Based generative models.  In _16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22)_ , pp. 521–538, Carlsbad, CA, July 2022. USENIX Association.  ISBN 978-1-939133-28-1.  URL https://www.usenix.org/conference/osdi22/presentation/yu.
  * Zheng et al. (2023) 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 Ion 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 https://openreview.net/forum?id=uccHPGDlao.
  * Zhou et al. (2024) Yongchao Zhou, Kaifeng Lyu, Ankit Singh Rawat, Aditya Krishna Menon, Afshin Rostamizadeh, Sanjiv Kumar, Jean-François Kagy, and Rishabh Agarwal.  Distillspec: Improving speculative decoding via knowledge distillation.  In _The Twelfth International Conference on Learning Representations_ , 2024.  URL https://openreview.net/forum?id=rsY6J3ZaTF.

##  Appendix A Eagle Decoding

Concurrently to our work, Li et al. (2024) have proposed the EAGLE decoding
framework. 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 use a collection of draft heads and instead
uses a 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 a 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 in the continuation, but also predicts an
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, as well as the estimated hidden state.

Li et al. (2024) demonstrate that EAGLE decoding provides a speedup relative
to Medusa’s sequentially-independent draft heads. Given that EAGLE was
developed entirely independently 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.