Learning Harmonized Representations for Speculative Sampling

## Abstract

Abstract Speculative sampling is a promising approach to accelerate the decoding stage for Large Language Models (LLMs). Recent advancements that leverage target LLM’s contextual information, such as hidden states and KV cache, have shown significant practical improvements. However, these approaches suffer from inconsistent context between training and decoding. We also observe another discrepancy between the training and decoding objectives in existing speculative sampling methods. In this work, we propose a solution named HArmonized Speculative Sampling (HASS) that learns harmonized representations to address these issues. HASS accelerates the decoding stage without adding inference overhead through harmonized objective distillation and harmonized context alignment. Experiments on four LLaMA models demonstrate that HASS achieves 2.81x-4.05x wall-clock time speedup ratio averaging across three datasets, surpassing EAGLE-2 by 8%-20%.

## 1 Introduction

Generative Large Language Models (LLMs), such as GPT-4 (Achiam et al., ) and LLaMA (Touvron et al., ), have demonstrated remarkable capabilities across a wide range of tasks. Nevertheless, efficiently decoding from these models poses a significant challenge due to the inherent auto-regressive decoding mechanism, which restricts their applicability in time-sensitive scenarios. Speculative sampling (Chen et al., ; Leviathan et al., ) offers a solution by leveraging additional resources to increase concurrency. Specifically, it employs an efficient draft model to generate draft tokens auto-regressively, which are then concurrently verified by the target LLM. Based on the verification results, a subset of draft tokens that preserves the same distribution as the target LLM is accepted as the final output.

Leviathan et al. () show that the practical performance of speculative sampling is highly related to two factors: the decoding cost of the draft model and its alignment with the target LLM. To develop efficient draft models that are well-aligned with the target LLM, previous works propose to leverage the target LLM’s contextual information (Xiao et al., ; Li et al., ; ; Du et al., ). For instance, EAGLE (Li et al., ; ) employs previous hidden states of the target LLM as the draft model’s input features. However, these approaches introduce inconsistent context between training and decoding, as illustrated in Figure 2. During training, the draft model always has access to the target LLM’s hidden states in previous timesteps. However, during decoding, the draft model cannot access the target LLM’s hidden states for unverified timesteps, resulting in a context misalignment between training and decoding. This issue can be viewed as a form of exposure bias (Bengio et al., ; Wang & Sennrich, ) at the feature level in speculative sampling.

Another discrepancy is also observed between the objectives of the training and decoding stages. During the decoding stage, the objective of the draft model is to propose tokens that the target LLM is likely to assign high probabilities to (Li et al., ; Miao et al., ; Sun et al., ). In this scenario, the draft model should focus more on recalling the desired tokens, while the specific order of these tokens can be somewhat de-emphasized. Moreover, most LLM applications perform nucleus sampling (Holtzman et al., ) or top-k sampling (Fan et al., ). For these decoding objectives, tokens with high probabilities play a more significant role in determining the output. Therefore, to develop efficient draft models, their training objectives should consider these properties encountered in the decoding stage. To the best of our knowledge, previous works on training draft models for speculative sampling have largely overlooked these decoding considerations.

Figure: Figure 1: Speedup ratios of different methods on LLaMA2-Chat 7B/13B and LLaMA3-Instruct 8/70B with temperature T $\in\{0,1\}$, averaging over MT-bench, HumanEval, and GSM8K datasets.
Refer to caption: x1.png

Figure: Figure 2: We use EAGLE (Li et al., ) as an example to illustrate the context misalignment, where the speculation starts from timestep $t$. $f^{(l)}$ and $f^{(s)}$ represent hidden states from the target LLM and the draft model, respectively. When decoding draft token $x_{t+2}$, the input context is inconsistent between training and decoding.
Refer to caption: x2.png

In this paper, we introduce HArmonized Speculative Sampling (HASS), a novel approach designed to address the aforementioned issues by learning harmonized representations. Specifically, to make draft models aware of the decoding strategy, HASS extends the idea of ranking distillation (Tang & Wang, ) from the recommender system to speculative sampling, resulting in a distillation loss focused on the most probable tokens within the target distribution. To mitigate the previously discussed context misalignment between training and decoding, HASS employs a context-aligned training strategy. Together, these two strategies of HASS improve the acceleration performance without any inference overhead and maintain training efficiency.

We conduct experiments across dialogue, code generation, and mathematical reasoning using the MT-bench, HumanEval, and GSM8K datasets, respectively. Building with EAGLE-2 (Li et al., ), HASS achieves 8%-16% acceptance length improvement over it on LLaMA2-Chat 7/13B and LLaMA3-Instruct 8/70B, resulting in 2.81x-4.05x wall-clock time acceleration compared with the vanilla inference on NVIDIA H800 GPU.

## 2 Preliminary

Speculative sampling leverages the concept of speculative execution (Kung & Robinson, ; Hennessy & Patterson, ) to reduce wall-clock time from more concurrency.
Specifically, given the target LLM $\mathcal{M}^{(l)}$ that is the focus of acceleration, speculative sampling employs a draft model $\mathcal{M}^{(s)}$ to speculatively and efficiently generate draft tokens.
The conventional approach (Leviathan et al., ; Chen et al., ) decomposes the next step generation into three steps:

- •
$\mathcal{M}^{(s)}$ proposes an unverified draft sequence with length $L$ by auto-regressive decoding.
- •
$\mathcal{M}^{(l)}$ evaluates posterior probabilities of $L$ draft tokens in parallel.
- •
$\tau$ tokens that retain the target distribution are accepted by a modified rejection sampling schema based on the draft sequence and the distribution gap.

Leviathan et al. () demonstrate that the wall-clock time improvement ratio is directly proportional to $\tau$, while the arithmetic operation increment ratio is inversely proportional to $\tau$. Consequently, $\tau$, also known as the acceptance length, plays a crucial role in determining the performance of acceleration. This analysis also applies when using multiple draft sequences (Miao et al., ; Li et al., ; ; Sun et al., ).
Note that $\tau$ is closely related to the distribution gap between the target LLM and the draft model. With efficient decoding requirements, the draft model typically has limited capacity, resulting in a significant distribution gap compared to the target LLM. Fortunately, during inference, the acceptance rate is primarily influenced by the alignment of distributions on the desired tokens, i.e., the tokens to which the target LLM assigns high probabilities. However, previous speculative sampling works mainly focus on the entire vocabulary set w.r.t. knowledge distillation from the target LLM (Li et al., ; Zhou et al., ), thereby disconnecting the training process from the practical decoding requirements.

EAGLE (Li et al., ) is a lightweight draft model design, as shown in Figure 2. During decoding, it utilizes the LM Head of the target LLM to generate draft tokens. Specifically, we assume that the speculation starts from timestep $t$, meaning the first draft token is at timestep $t+1$. To generate the draft token $x_{t+1}$, the target LLM’s hidden state $f_{t-1}^{(l)}$ in the second-to-top layer is concatenated with the embedding $e_{t}$ to perform the input of the draft model. During training, EAGLE constructs a regression task between $f^{(l)}$s and the predicted hidden states $f^{(s)}$s of the draft model.
However, due to the auto-regressive decoding, the draft model only accesses the target LLM’s features at the beginning of the speculation. It uses the features produced by the draft models as input for subsequent steps. This context misalignment, stemming from feature inaccuracies, leads to error accumulation and hinders the performance of generating later draft tokens (Li et al., ; Du et al., ). EAGLE-2 (Li et al., ) employs the same model design but works on dynamic drafting structures instead of a static tree structure during the decoding stage, yet the aforementioned issue remains unresolved.

## 3 Methodology

As outlined before, previous speculative sampling methods suffer from disharmonies between training and decoding. This section introduces HArmonized Speculative Sampling (HASS) to tackle objective misalignment and context inconsistency through harmonized objective distillation and harmonized context alignment, respectively, as described below.

### 3.1 Harmonized Objective Distillation

HASS prioritizes the most decoding-desired tokens by leveraging the ranking distillation (Tang & Wang, ) idea from the recommender system. Specifically, ranking distillation aims to train a student model to assign higher ranks to the items that are top-ranked by the teacher model. In the context of speculative sampling, the draft model and the target LLM serve as the student and the teacher, respectively. Draft models with similar properties will perform at a higher acceptance rate in the decoding stage. Consider the set of K tokens with the highest probabilities from the target LLM’s probability distribution as $\hat{\Omega}\subset\Omega$, where $\Omega$ represents the entire vocabulary. HASS considers the following Top-K distillation loss:

$$ $L_{\text{Top-K}}=-\sum_{x\in\hat{\Omega}}q(x)\log p(x),$ (1) $$

where $q$ and $p$ are the next token probability distributions of the target LLM and the draft model, respectively.
Note that, when integrated with EAGLE, the training stage can obtain $\hat{\Omega}$ from hidden states of the target LLM. This implies that the proposed loss function benefits from the same efficient training cost as EAGLE. We evaluate the proposed Top-K distillation loss against six alternative losses, such as BiLD (Li et al., ) and Recall@k Surrogate loss (Patel et al., ), through ablation studies.

### 3.2 Harmonized Context Alignment

Figure: Figure 3: Training with harmonized context alignment, where $q$ and $k$ refer to the query and key states in the transformer layer, respectively. Superscript $(l)$ denotes tensors from the target LLM, and superscript $(s_{j})$ denotes tensors from the $j$-th draft model forward. Note that during training $(s_{j})$ refers to calling $j$ times draft model in a batch, while during inference $(s_{j})$ refers to $j$-th auto-regressive decoding.
Refer to caption: x3.png

HASS follows a context alignment schema that aligns training and decoding on their contexts. In particular, HASS divides the training procedure into $n$ steps, enabling the draft model to utilize contextual features consistent with those in the decoding stage. The procedure is outlined as follows:

- •
The first step mirrors the training stage of EAGLE. At timestep $t+1$, the draft model takes the target model’s feature $f^{(l)}_{t}$ as input and produces the draft feature $f^{(s_{1})}_{t+1}$. In this step, the attention mask remains as the original causal mask without any modification.
- •
In the second step, features from the first step are incorporated. For instance, in the self-attention mechanism at timestep $t+1$, $f^{(s_{1})}_{t}$ is used to calculate the current query. Keys and values are derived from $f^{(l)}_{:t}\oplus f^{(s_{1})}_{t}$, where $\oplus$ denotes concatenation and $f^{(l)}_{:t}$ includes features from timesteps earlier than $t$. The attention mask is adjusted to ensure that the previous feature seen by $f^{(s_{1})}_{i}$ is always $f^{(l)}_{i-1}$, as shown in the ‘HASS Training Step 2’ part of Figure 3.
- •
For step $j\geq 3$, the feature from the previous step $f^{(s_{j-1})}_{t}$ is utilized to calculate the query at timestep $t+1$, while keys and values are generated by $f^{(l)}_{:t-j+2}\oplus f^{(s_{1})}_{t-j+2}\oplus\ldots\oplus f^{(s_{j-1})}_{t}$.

It takes $n$ times the training overhead of EAGLE while maintaining the same decoding overhead. We empirically demonstrate that the acceleration effect converges with a small $n$ so that the training of HASS is cost-efficient.

**Table 1: Acceptance lengths $\tau$ of different methods on MT-bench, HumanEval, and GSM8K datasets with temperature $\text{T}\in\{0,1\}$. L2 represents LLaMA2-Chat, while L3 represents LLaMA3-Instruct.**
|  |  | Temperature = 0 | Temperature = 1 |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | Method | MT-bench | HumanEval | GSM8K | Mean | MT-bench | HumanEval | GSM8K | Mean |
| L2 7B | PLD | 1.43 | 1.59 | 1.37 | 1.46 | - | - | - | - |
| Lookahead | 1.66 | 1.77 | 1.65 | 1.69 | - | - | - | - |  |
| EAGLE | 3.68 | 3.90 | 3.77 | 3.78 | 3.45 | 3.67 | 3.62 | 3.58 |  |
| EAGLE-2 | 4.44 | 4.78 | 4.60 | 4.61 | 4.23 | 4.47 | 4.50 | 4.40 |  |
| HASS | 4.99 | 5.29 | 5.17 | 5.15 | 4.84 | 4.91 | 5.01 | 4.92 |  |
| L2 13B | PLD | 1.46 | 1.70 | 1.44 | 1.53 | - | - | - | - |
| Lookahead | 1.64 | 1.85 | 1.69 | 1.73 | - | - | - | - |  |
| EAGLE | 3.86 | 4.50 | 4.17 | 4.18 | 3.62 | 4.27 | 3.98 | 3.96 |  |
| EAGLE-2 | 4.74 | 5.57 | 5.17 | 5.16 | 4.60 | 5.41 | 5.03 | 5.01 |  |
| HASS | 5.13 | 6.05 | 5.55 | 5.58 | 4.98 | 5.86 | 5.41 | 5.42 |  |
| L3 8B | EAGLE | 2.91 | 3.66 | 3.57 | 3.38 | 2.67 | 3.35 | 3.30 | 3.11 |
| EAGLE-2 | 4.21 | 4.93 | 4.42 | 4.52 | 3.90 | 4.73 | 4.30 | 4.31 |  |
| HASS | 4.68 | 5.54 | 5.02 | 5.08 | 4.26 | 5.30 | 4.85 | 4.80 |  |
| L3 70B | EAGLE | 3.24 | 4.07 | 3.79 | 3.70 | 3.06 | 3.85 | 3.66 | 3.52 |
| EAGLE-2 | 4.10 | 5.02 | 4.37 | 4.50 | 4.00 | 4.93 | 4.35 | 4.43 |  |
| HASS | 4.62 | 5.78 | 5.24 | 5.21 | 4.59 | 5.68 | 5.20 | 5.16 |  |

**Table 2: Speedup ratios of different methods on MT-bench, HumanEval, and GSM8K datasets with temperature T $\in\{0,1\}$. L2 represents LLaMA2-Chat, while L3 represents LLaMA3-Instruct.**
|  |  | Temperature = 0 | Temperature = 1 |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | Method | MT-bench | HumanEval | GSM8K | Mean | MT-bench | HumanEval | GSM8K | Mean |
| L2 7B | EAGLE | 1.90x | 2.10x | 2.04x | 2.01x | 1.50x | 1.91x | 1.87x | 1.76x |
| EAGLE-2 | 2.66x | 3.06x | 2.72x | 2.81x | 2.39x | 2.87x | 2.54x | 2.60x |  |
| HASS | 2.99x | 3.41x | 3.32x | 3.24x | 2.70x | 3.13x | 2.87x | 2.90x |  |
| L2 13B | EAGLE | 1.80x | 2.46x | 2.41x | 2.22x | 1.84x | 2.10x | 2.21x | 2.05x |
| EAGLE-2 | 3.02x | 3.64x | 3.23x | 3.30x | 3.04x | 3.45x | 3.13x | 3.21x |  |
| HASS | 3.23x | 4.24x | 3.48x | 3.65x | 3.28x | 3.78x | 3.37x | 3.48x |  |
| L3 8B | EAGLE | 1.29x | 2.00x | 1.85x | 1.71x | 1.25x | 1.41x | 1.67x | 1.44x |
| EAGLE-2 | 2.64x | 3.31x | 2.54x | 2.83x | 2.39x | 2.54x | 2.48x | 2.47x |  |
| HASS | 2.78x | 3.43x | 3.06x | 3.09x | 2.49x | 3.05x | 2.89x | 2.81x |  |
| L3 70B | EAGLE | 2.14x | 2.74x | 2.42x | 2.43x | 1.80x | 2.34x | 2.12x | 2.09x |
| EAGLE-2 | 2.94x | 3.98x | 3.19x | 3.37x | 3.02x | 3.61x | 3.21x | 3.28x |  |
| HASS | 3.40x | 4.68x | 4.08x | 4.05x | 3.43x | 4.25x | 3.87x | 3.85x |  |

## 4 Experiment

### 4.1 Experimental Setup

Target LLMs. LLaMA2-Chat 7B/13B and LLaMA3-Instruct 8/70B.

Tasks. We conduct evaluations on three generation tasks. For multi-turn conversation, code generation, and mathematical reasoning tasks, we choose the MT-bench (Zheng et al., ), HumanEval (Chen et al., ), and GSM8K (Cobbe et al., ) datasets, respectively.
The batch size is set as 1 under all the experiments following Leviathan et al. () and Zhou et al. ().

Metrics. HASS neither fine-tunes the target LLMs’ weights during training nor relaxes the acceptance conditions during decoding, making it a lossless acceleration method. Thus, the generation quality is promised with no need for evaluation. We use the following two metrics to measure the acceleration performance:

- •
Speedup Ratio: The actual test speedup ratio relative to vanilla auto-regressive decoding.
- •
Acceptance Length $\tau$: The average number of tokens generated per drafting-verification cycle, indicating the number of tokens accepted by the target LLM from the draft model.

Note that the speedup ratio is sensitive to the hardware due to variations in computing power, and the acceptance length may be slightly affected by hardware due to numerical errors. Therefore, all inference processes are conducted on NVIDIA H800 GPU.

Comparisons. The vanilla auto-regressive decoding is taken as the baseline, which serves as the benchmark for speedup ratios (1.00x). We compare HASS with recent lossless speculative sampling methods, including PLD (Saxena, ), Lookahead (Fu et al., ), EAGLE (Li et al., ), and EAGLE-2 (Li et al., ).

Implementation. Our code is built based on EAGLE-2’s open-source repository(^1^11https://github.com/SafeAILab/EAGLE). Experiments on EAGLE and EAGLE-2 reuse draft model weights trained by Li et al. (). For harmonized objective distillation, K is set as 10, and the loss of harmonized objective distillation is added to EAGLE’s original loss with a coefficient of $w=1.0$. For harmonized context alignment, the draft model is aligned for 3 steps during training. For dynamic tree structure, we set the total number of draft tokens to 60 for all experiments with a draft tree depth of 6. We keep other settings, such as the fixed training dataset and the optimizer, consistent with EAGLE-2.

### 4.2 Effectiveness & Ablation Study

In this section, we first evaluate the effectiveness of HASS by comparing it with existing speculative sampling methods on acceptance length and speedup ratio. Then, we conduct ablation studies on harmonized objective distillation and harmonized context alignment.

#### 4.2.1 Effectiveness

We present different methods’ acceptance lengths and speedup ratios across three datasets in Tables 1 and 2, respectively. HASS performs the largest acceptance length and highest speedup ratio across all datasets and LLMs we tested.
Most methods achieve their best performance on the HumanEval dataset, as the fixed templates in the code generation task are easier to draft and accelerate.
Though PLD and Lookahead are free of training, they consistently show poorer performance than EAGLE, EAGLE-2, and HASS.

#### 4.2.2 Ablation Study on Harmonized Objective Distillation

Figure: Figure 4: Acceptance lengths $\tau$ of HASS with varied Ks and weights of the Top-K loss. The results are conducted on LLaMA2-Chat 7B and averaged over MT-bench, HumanEval, and GSM8K datasets with temperature T $\in\{0,1\}$.
Refer to caption: x4.png

We first study the effects of different K and the weight $w$ of the Top-K loss by varying these hyper-parameters and summarize the results in Figure 4. Training with the Top-K loss ($w>0$) always improves performance compared to training without the Top-K loss ($w=0$). HASS achieves the largest acceptance length when $w=0.5$.
A small value of K may result in performance degeneration, as the draft model only focuses on the token with the highest probability and consequently neglects other potential tokens. With a larger K, the Top-K loss generally brings better results, while the acceptance length is the largest when K $=5$.

**Table 3: Acceptance lengths $\tau$ of HASS with different kinds of loss functions for harmonized objective distillation. The results are conducted on LLaMA2-Chat 7B over the MT-bench dataset with temperature T $\in\{0,1\}$.**
| Loss Function | Temperature = 0 | Temperature = 1 | Mean |
| --- | --- | --- | --- |
| Top-K Loss | 4.99 | 4.84 | 4.92 |
| Top-P Loss | 5.03 | 4.76 | 4.90 |
| Normed Top-K Loss (Linear) | 4.97 | 4.83 | 4.90 |
| Normed Top-K Loss (Softmax) | 4.98 | 4.72 | 4.85 |
| Bi-directional Top-K Loss | 4.99 | 4.72 | 4.86 |
| Recall@k Surrogate Loss | 4.97 | 4.76 | 4.87 |
| BiLD Loss | 5.04 | 4.75 | 4.90 |

Since the harmonized objective distillation can be implemented with any loss function that focuses on the most probable tokens w.r.t. the target distribution, we further consider the following loss functions and compare them with the Top-K Loss:

- •
Top-P Loss, where the $\hat{\Omega}$ is formed by the most probable tokens whose cumulative probability is just larger than P.
- •
Normed Top-K Loss, where the target and draft distributions are both normalized over $\hat{\Omega}$. The normalization operation can be either linear or softmax.
- •
Bi-directional Top-K Loss, where the distillation is conducted over the most probable tokens w.r.t. the target distribution as well as the draft distribution.
- •
Recall@k Surrogate Loss (Patel et al., ), where a smooth approximation of the recall metric is obtained and is differentiable for direct optimization.
- •
BiLD Loss (Li et al., ), where the internal logits ranking information is captured by constructing logits differences with long-tail noise filtered out.

After searching the optimal hyper-parameters for each of the compared loss functions, we summarize their best results in Table 3.
BiLD loss outperforms other loss functions under temperature T $=0$, while Top-K loss outperforms others under temperature T $=1$. Generally, Top-K loss shows the best performance.
A better loss function may exist than Top-K loss to exploit the target LLM further. We leave this topic in future works.

We also conduct an experiment with LLaMA2-Chat 7B, where the fixed training dataset is replaced by the dataset generated by the target LLM (see Appendix 8). We observe that when using non-greedy decoding, the acceptance length increases from 4.92 to 5.19 averaging over three datasets. Therefore, information obtained from harmonized objective distillation is not equivalent to direct distillation from target-model-generated data.

#### 4.2.3 Ablation Study on Harmonized Context Alignment

**Table 4: Acceptance lengths $\tau$ of HASS with varied aligning steps in the harmonized context alignment. The results are conducted on LLaMA2-Chat 7B with temperature T $\in\{0,1\}$.**
|  | Aligning Step | MT-bench | HumanEval | GSM8K | Mean |
| --- | --- | --- | --- | --- | --- |
| T=0 | EAGLE-2 + Top-K | 4.59 | 4.97 | 4.77 | 4.78 |
| HASS Align-2 | 4.95 | 5.25 | 5.12 | 5.11 |  |
| HASS Align-3 | 4.99 | 5.29 | 5.17 | 5.15 |  |
| HASS Align-4 | 4.99 | 5.30 | 5.18 | 5.16 |  |
| HASS Align-5 | 4.98 | 5.26 | 5.09 | 5.11 |  |
| T=1 | EAGLE-2 + Top-K | 4.46 | 4.61 | 4.64 | 4.57 |
| HASS Align-2 | 4.71 | 4.89 | 4.98 | 4.86 |  |
| HASS Align-3 | 4.84 | 4.91 | 5.01 | 4.92 |  |
| HASS Align-4 | 4.77 | 4.93 | 5.03 | 4.91 |  |
| HASS Align-5 | 4.71 | 4.92 | 4.95 | 4.86 |  |

We propose the harmonized context alignment, which eliminates the feature inconsistency of draft models between the training and decoding stages.
To study the effect of increasing the aligning steps in the harmonized context alignment, we conduct experiments by varying the step number and summarize the results in Table 4.

As the first training step of HASS is the same as EAGLE-2, we continually train EAGLE-2’s draft model weights with the Top-K loss and consider it the baseline.
Without harmonized context alignment (EAGLE-2 + Top-K), the draft model performs the worst across all datasets.
Training with 3/4 steps of harmonized context alignment generally obtains the most considerable acceptance length.
When training with five steps of context alignment, the acceptance length decreases. We believe this is caused by the draft model’s limited capacity, as it predicts less accurately on former steps’ tokens when paying too much attention to the latter ones. Figure 5 shows the acceptance rate $\alpha$ across speculation steps on the MT-bench dataset following Li et al. (). In later speculation steps, HASS performs better acceptance rates than EAGLE-2, demonstrating the effectiveness of harmonized context alignment.

Figure: Figure 5: Acceptance rates $\alpha$ of HASS and EAGLE-2 across different speculation steps on the MT-bench dataset with temperature T $\in\{0,1\}$.
Refer to caption: x5.png

As shown in Figure 5, the acceptance rates of HASS decrease compared to EAGLE-2 on LLaMA2-Chat 13B and LLaMA3-Instruct 70B at the first step (0-$\alpha$).
The draft models degrade on the first speculation step with much attention paid to the latter speculation steps, while the first step’s acceptance rates are crucial to larger acceptance lengths.
We conduct experiments to emphasize the significance of former speculation steps by reweighting the training loss from each step with a factor $\beta$. In specific, the step $j$’s training loss will be multiplied by $\beta^{j-1}$. Table 5 and Figure 6 show the acceptance lengths and acceptance rates of HASS with different reweight factors on LLaMA3-Instruct 70B over the MT-bench dataset, respectively.
With the factor $\beta$ decreasing from 1.0 to 0.5, HASS achieves better acceptance lengths with different temperatures.
Correspondingly, we perceive that the acceptance rate at the first speculation step is consistently higher with a smaller $\beta$, while the acceptance rates at the latter speculation steps generally decline.
When the factor $\beta$ decreases to 0.3, too much emphasis is assigned to the first speculation step, leading to deterioration in acceptance length.
Since the further exploration of an appropriate trade-off between different speculation steps is out of this paper’s scope, we leave it for future work.

| Reweight Factor $\beta$ | T = 0 | T = 1 | Mean |
| --- | --- | --- | --- |
| 1.0 (Default) | 4.62 | 4.59 | 4.61 |
| 0.7 | 4.65 | 4.61 | 4.63 |
| 0.5 | 4.67 | 4.62 | 4.65 |
| 0.3 | 4.65 | 4.61 | 4.63 |

Table: Table 5: Acceptance lengths $\tau$ of HASS with different reweight factors $\beta$ for harmonized context alignment. The results are conducted on LLaMA3-Instruct 70B over the MT-bench dataset with temperature T $\in\{0,1\}$.

Figure: Figure 6: Acceptance rates $\alpha$ of HASS with different reweight factors $\beta$ for harmonized context alignment. The results are conducted on LLaMA3-Instruct 70B over the MT-bench dataset with temperature T $\in\{0,1\}$.

## 5 Related Work

There have been a number of works on improving the acceptance rate of speculative sampling while maintaining the target distribution. Most of them fall into two categories. (1) The former category is aligned training that tries to obtain draft models aligned with the target LLM before the decoding stage. Zhou et al. () propose a knowledge distillation approach and study several strategies to improve the alignment. Li et al. () demonstrate that hidden states of the target LLM as input of the draft model provide extra feature uncertainty information. Xiao et al. () also utilize hidden states of the target LLM and introduce an RNN-based draft model design that achieves a comparable acceptance rate. GLIDE (Du et al., ) instead reuses the KV cache of the target LLM. It also notices the context misalignment when using information from the target LLM, but the proposed blockwise attention mask method can not solve the misalignment completely. (2) The latter category is efficient decoding, which designs sophisticated decoding strategies to utilize concurrency efficiently. Miao et al. () propose to utilize multiple draft models and design a tree-based attention mechanism to verify multiple draft sequences efficiently. Li et al. () introduce a dynamic structure to save computation by pruning inefficient paths in the draft tree. Sun et al. () study improving the verification stage through optimal transportation. However, these works tend to only consider training or decoding, ignoring the linkage of these two stages. This work instead aims to link training and decoding, leading to harmonized speculative sampling.

## 6 Conclusion

This paper introduces HASS, a harmonized speculative sampling solution that addresses disharmonies between training and decoding on their objectives and contexts. Compared to its closest baseline, EAGLE-2, HASS improves the acceptance rate without any inference overhead. Experiments conducted on LLaMA2-Chat 7/13B and LLaMA3-Instruct 8/70B demonstrate the effectiveness and efficiency of HASS. Averaging on MT-bench, HumanEval, and GSM8K, HASS is 2.81x-4.05x faster than vanilla auto-regressive decoding, 8%-20% faster than EAGLE-2.

## Appendix for ‘Learning Harmonized Representations for
Speculative Sampling’

## Appendix A Implementation of Harmonized Context Alignment

We present the pseudo code of harmonized context alignment, which is implemented without the customized attention mask, for better understanding.
The actual implementation in our experiments is achieved by the customized attention mask as shown in Figure 3.

## Appendix B Harmonized Context Alignment on Tokens

In this section, we attempt to verify whether applying token alignment as well as feature alignment brings better performance.
In specific, we use the tokens generated by the draft model for training in harmonized context alignment instead of using the tokens from training data.
We apply feature and token alignment to EAGLE-2’s draft model weights and summarize the results in Table 6 and Figure 7.

**Table 6: Acceptance lengths $\tau$ of applying feature and token alignment to EAGLE-2’s draft model weights, where ‘Token ($x$)’ denotes tokens from training data are replaced by draft-model-generated tokens with a probability of $x$. The results are conducted on LLaMA2-Chat 7B over the MT-bench dataset with temperature T $\in\{0,1\}$.**
|  | Temperature = 0 | Temperature = 1 | Mean |
| --- | --- | --- | --- |
| EAGLE-2 | 4.44 | 4.23 | 4.34 |
| Feature Only | 4.83 | 4.60 | 4.72 |
| Feature + Token (0.1) | 4.81 | 4.57 | 4.69 |
| Feature + Token (0.2) | 4.78 | 4.51 | 4.65 |
| Feature + Token (1.0) | 4.28 | 4.11 | 4.20 |

Figure: Figure 7: Acceptance rates $\alpha$ of applying feature and token alignment to EAGLE-2’s draft model weights, where ‘Token ($x$)’ denotes tokens from training data are replaced by draft-model-generated tokens with a probability of $x$. The results are conducted on LLaMA2-Chat 7B over the MT-bench dataset with temperature T $\in\{0,1\}$.
Refer to caption: x7.png

Feature only alignment brings the best performance, while token alignment leads to degeneration. With the probability of applying token alignment increasing from 0.1 to 1.0, the acceptance length decreases consistently.
As shown in Figure 7, more token alignment generally causes lower acceptance rates.
As a result, training with tokens generated by the draft model in harmonized context alignment hurts the acceleration performance.

## Appendix C Hyper-Parameters of Top-K Loss

We conduct an ablation study on Top-K loss’s hyper-parameters, i.e., K and $w$, in Section 4.2.2 and show the averaged acceptance lengths over three datasets in Figure 4. Here, we present the speedup ratios and acceptance lengths of HASS with varied Ks and $w$s in Table 7.

**Table 7: Speedup ratios and acceptance lengths $\tau$ of HASS with varied Ks and $w$s of the Top-K loss on LLaMA2-Chat 7B over MT-bench, HumanEval, and GSM8K datasets with temperature T $\in\{0,1\}$.**
|  |  |  | MT-bench | HumanEval | GSM8K | Mean |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | K | $w$ | Speedup | $\tau$ | Speedup | $\tau$ | Speedup | $\tau$ | Speedup | $\tau$ |
| T=0 | 1 | 1.0 | 2.89x | 4.94 | 3.24x | 5.19 | 3.11x | 5.10 | 3.08x | 5.08 |
| 5 | 1.0 | 2.90x | 5.00 | 3.44x | 5.29 | 3.33x | 5.18 | 3.22x | 5.16 |  |
| 10 | 1.0 | 2.99x | 4.99 | 3.41x | 5.29 | 3.32x | 5.17 | 3.24x | 5.15 |  |
| 50 | 1.0 | 2.85x | 5.01 | 3.46x | 5.29 | 3.41x | 5.17 | 3.24x | 5.16 |  |
| 100 | 1.0 | 2.93x | 5.00 | 3.45x | 5.29 | 3.45x | 5.18 | 3.28x | 5.16 |  |
| 10 | 0.0 | 2.77x | 4.93 | 3.38x | 5.22 | 3.18x | 5.11 | 3.11x | 5.09 |  |
| 10 | 0.1 | 2.98x | 4.96 | 3.40x | 5.26 | 3.51x | 5.16 | 3.30x | 5.13 |  |
| 10 | 0.2 | 2.87x | 4.98 | 3.41x | 5.29 | 3.35x | 5.16 | 3.21x | 5.14 |  |
| 10 | 0.5 | 3.00x | 5.02 | 3.32x | 5.31 | 3.50x | 5.18 | 3.27x | 5.17 |  |
| 10 | 2.0 | 2.94x | 4.98 | 3.37x | 5.29 | 3.34x | 5.17 | 3.22x | 5.15 |  |
| T=1 | 1 | 1.0 | 2.58x | 4.70 | 2.79x | 4.80 | 2.83x | 4.95 | 2.73x | 4.82 |
| 5 | 1.0 | 2.64x | 4.81 | 3.13x | 4.94 | 2.93x | 5.02 | 2.90x | 4.92 |  |
| 10 | 1.0 | 2.70x | 4.84 | 3.13x | 4.91 | 2.87x | 5.01 | 2.90x | 4.92 |  |
| 50 | 1.0 | 2.62x | 4.77 | 3.01x | 4.88 | 2.99x | 5.03 | 2.87x | 4.89 |  |
| 100 | 1.0 | 2.66x | 4.74 | 3.14x | 4.97 | 2.90x | 5.03 | 2.90x | 4.91 |  |
| 10 | 0.0 | 2.61x | 4.71 | 2.76x | 4.84 | 2.79x | 4.96 | 2.72x | 4.84 |  |
| 10 | 0.1 | 2.69x | 4.75 | 3.05x | 4.94 | 2.87x | 5.00 | 2.87x | 4.90 |  |
| 10 | 0.2 | 2.66x | 4.75 | 3.16x | 4.95 | 2.88x | 5.01 | 2.90x | 4.90 |  |
| 10 | 0.5 | 2.68x | 4.80 | 3.15x | 4.93 | 2.96x | 5.03 | 2.93x | 4.92 |  |
| 10 | 2.0 | 2.68x | 4.75 | 3.11x | 4.89 | 2.89x | 5.03 | 2.89x | 4.89 |  |

## Appendix D Self-Distillation

Following existing speculative sampling methods (Zhou et al., ; Cai et al., ), we further use target-model-generated outputs to distill the draft model from the target model’s real output distribution, dubbed as self-distillation.
In specific, we feed the prompts from the ShareGPT dataset into the target models recursively and collect the responses as multi-turn conversations for self-distillation.
To study the effect of self-distillation, we conduct experiments by training the draft model with fixed data or model-generated data and summarize the results in Table 8.

**Table 8: Speedup ratios and acceptance lengths of HASS with fixed or target-model-generated training data. F and MG stand for ‘Fixed’ and ‘Model-Generated’, respectively. L2 represents LLaMA2-Chat.**
|  |  |  | MT-bench | HumanEval | GSM8K | Mean |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Model | Data | Speedup | $\tau$ | Speedup | $\tau$ | Speedup | $\tau$ | Speedup | $\tau$ |
| T=0 | L2 7B | F | 2.99x | 4.99 | 3.41x | 5.29 | 3.32x | 5.17 | 3.24x | 5.15 |
| MG | 3.13x | 5.25 | 3.85x | 5.70 | 3.40x | 5.57 | 3.46x | 5.51 |  |  |
| L2 13B | F | 3.23x | 5.13 | 4.24x | 6.05 | 3.48x | 5.55 | 3.65x | 5.58 |  |
| MG | 3.34x | 5.27 | 4.42x | 6.00 | 3.63x | 5.61 | 3.80x | 5.63 |  |  |
| T=1 | L2 7B | F | 2.70x | 4.84 | 3.13x | 4.91 | 2.87x | 5.01 | 2.90x | 4.92 |
| MG | 2.75x | 4.97 | 3.39x | 5.24 | 3.13x | 5.35 | 3.09x | 5.19 |  |  |
| L2 13B | F | 3.28x | 4.98 | 3.78x | 5.86 | 3.37x | 5.41 | 3.48x | 5.42 |  |
| MG | 3.33x | 5.02 | 3.76x | 5.80 | 3.60x | 5.42 | 3.56x | 5.41 |  |  |

On LLaMA2-Chat 7B, self-distillation consistently brings improvements.
On LLaMA2-Chat 13B, self-distillation only achieves marginally better or comparable results, as the corresponding draft model’s performance is saturated.
Especially, the acceptance lengths of the self-distilled draft model are lower than that of the vanilla one on the HumanEval dataset. It may be due to the code generation dataset HumanEval being less similar to the training dataset ShareGPT compared with MT-bench and GSM8K.

## Appendix E Drafting Hyper-Parameters

Li et al. () find that the draft token’s confidence score is strongly positively correlated with the acceptance rate, and accordingly propose the context-aligned dynamic draft tree, which can be dynamically adjusted with two hyper-parameters: ‘depth’ and ‘number of tokens’.
‘Depth’ decides the draft tree’s depth during the expansion phase, while ‘number of tokens’ decides how many draft tokens will be kept during the reranking phase.
Increasing both these hyper-parameters surely leads to a larger acceptance length. Nevertheless, sending more draft tokens into the target model for verification causes a higher overhead in real applications.
Therefore, we vary these hyper-parameters and report the speedup ratios in Table 9 to find a better trade-off.

**Table 9: Speedup ratios of EAGLE-2 and HASS-MG with varied depths and numbers of tokens on the MT-bench dataset with temperature T $\in\{0,1\}$, where HASS-MG denotes HASS trained with self-distillation. L2 represents LLaMA2-Chat.**
| Depth | 5 | 6 | 7 | 8 | 9 |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| # Tokens | 40 | 60 | 80 | 100 | 40 | 60 | 80 | 100 | 40 | 60 | 80 | 100 | 40 | 60 | 80 | 100 | 40 | 60 | 80 | 100 |  |  |
| T=0 | L2 7B | EAGLE-2 | 2.48x | 2.78x | 2.61x | 2.69x | 2.69x | 2.66x | 2.70x | 2.79x | 2.71x | 2.86x | 2.91x | 2.95x | 2.60x | 2.60x | 2.88x | 2.89x | 2.28x | 2.54x | 2.55x | 2.65x |
| HASS-MG | 3.09x | 3.02x | 3.04x | 3.22x | 3.08x | 3.13x | 3.19x | 3.22x | 3.14x | 3.11x | 3.29x | 3.27x | 3.07x | 3.16x | 3.31x | 3.32x | 2.78x | 2.75x | 2.95x | 3.03x |  |  |
| L2 13B | EAGLE-2 | 2.63x | 2.96x | 3.04x | 3.06x | 3.01x | 3.02x | 3.18x | 3.22x | 2.78x | 3.12x | 3.14x | 3.24x | 2.98x | 3.12x | 3.19x | 3.26x | 2.49x | 2.64x | 2.69x | 2.72x |  |
| HASS-MG | 3.25x | 3.31x | 3.24x | 3.25x | 3.33x | 3.34x | 3.49x | 3.40x | 3.19x | 3.42x | 3.36x | 3.40x | 3.15x | 3.40x | 3.40x | 3.37x | 2.70x | 2.74x | 3.09x | 3.02x |  |  |
| T=1 | L2 7B | EAGLE-2 | 2.31x | 2.37x | 2.55x | 2.36x | 2.42x | 2.39x | 2.33x | 2.40x | 2.49x | 2.66x | 2.65x | 2.44x | 2.36x | 2.48x | 2.38x | 2.64x | 2.29x | 2.22x | 2.27x | 2.42x |
| HASS-MG | 2.79x | 2.89x | 2.86x | 2.88x | 2.72x | 2.75x | 2.92x | 2.82x | 2.83x | 2.76x | 2.81x | 2.75x | 2.49x | 2.68x | 2.77x | 2.77x | 2.30x | 2.35x | 2.58x | 2.50x |  |  |
| L2 13B | EAGLE-2 | 2.92x | 3.11x | 2.88x | 2.79x | 3.06x | 3.04x | 3.16x | 2.93x | 3.05x | 3.14x | 3.14x | 3.11x | 3.00x | 3.13x | 3.15x | 2.98x | 2.61x | 2.72x | 2.65x | 2.54x |  |
| HASS-MG | 3.24x | 3.30x | 3.27x | 3.19x | 3.33x | 3.33x | 3.40x | 3.28x | 3.19x | 3.26x | 3.24x | 3.26x | 3.15x | 3.26x | 3.19x | 3.17x | 2.62x | 2.77x | 2.74x | 2.84x |  |  |

When $\text{`depth'}=5$, the acceptance length is relatively small. When $\text{`depth'}=9$, the verification overhead is extremely high. Thus, neither of these settings achieves a promising speedup ratio.
For both HASS-MG and EAGLE-2, the best performances are achieved when $\text{`depth'}\in\{6,7,8\}$ and $\text{`\# tokens'}\in\{60,80,100\}$.
HASS-MG consistently obtains a superior performance compared with EAGLE-2 through hyper-parameter tuning across different LLMs and temperatures.

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 3
presentation: 3
questions: Seek weakness.

Will update the score after looking at the results of Weakness.
soundness: 3
strengths: The experiments are well-detailed, with clear metrics for comparison against baselines (e.g., EAGLE and multiple architectures).

The method makes sense in the perspective of typical acceptance process in SD.
summary: Hass method includes leveraging a hybrid autoregressive and non-autoregressive (HASS) approach to optimize decoding time. The authors present experimental results demonstrating improved performance compared to baseline SD methods.
weaknesses: 1. While the paper introduces a novel approach, it does not sufficiently explore the construction and utilization of the self-distillation dataset. Isn't the quality and configuration of this dataset more crucial than the framework itself? A deeper discussion on how the dataset is designed and its influence on model performance would strengthen the claims [Referring to A].

2. The paper lacks experiments analyzing how the framework performs across different token counts and task types [Referring to A]. Adding these analyses could provide a more comprehensive understanding of the proposed method's scalability and robustness. For instance, how does performance vary when the token count significantly increases or decreases? Are there certain task types where the proposed method excels or struggles compared to baselines?

[A] Yi et al., 2024. Towards Fast Multilingual LLM Inference: Speculative Decoding and Specialized Drafters, EMNLP 2024-main.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: - Could you please specify what data was used for training the draft model in the experiments? Does the proposed method (and EAGLE) suffer from generalization issues when speculative decoding is tested on data different from what was used to train the draft model?
- Can training on Model Generated data mitigate the context misalignment? In other words, will HASS outperform EAGLE-2 trained on Model-Generated data? It'd be interesting to see a comparison to EAGLE-2 in Table 8.

My rating for this paper is (7), which is not available as an option. If my questions and clarity concerns are addressed, I will be happy to consider raising the rating to (8).
soundness: 4
strengths: The paper carefully considers the two important aspects of speculative decoding efficiency: the ability of the draft model to predict top-K tokens and the number of tokens the draft model can successfully predict. Authors identify respective shortcomings of the recent method EAGLE and propose practical solutions.

- Empirical results are convincing and demonstrate meaningful improvements in terms of acceptance rate and speedups
- Authors present extensive ablation studies providing additional insights
summary: The paper proposes two improvements to the speculative decoding method EAGLE. First, it regularizes training of the draft model to focus on top-K tokens. Second, it trains the draft model to use its own predictions when drafting more than one token. Both improvements are inspired by how speculative decoding is used during inference.
weaknesses: A more rigorous mathematical presentation of the loss function in Section 3.2 (harmonized context alignment) would help to improve clarity. It is not easy to parse what is the input to the draft model at every step and what is it trying to predict. Writing out the objective function for training the draft model would help. I am also confused by Figure 3 - should Training Step 2 have $f^{(l)}\_{t-2}$ and $f^{(l)}\_{t-1}$ in the bottom row and only use superscript $(s_1)$ in the bottom right entry $f^{(s_1)}_{t}$ (analogously for steps 3 and 4)?

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. What does the vertical axis of the line charts in Figures 5 and 6 represent? The ranges shown do not appear to correspond to the acceptance rate.
If the authors are willing to address our concerns, we are open to raise the scores during the rebuttal stage.
soundness: 3
strengths: 1. The motivation to improve the current speculative decoding method is clear and well-founded.
2. The paper is well-written and easy to follow.
3. The proposed HASS significantly outperforms other baseline methods and achieves impressive speed-up performance. Without introducing extra inference costs, HASS can be nearly considered a 'free lunch' while incurring only acceptable additional training costs (3 or 4 extra alignment steps).
4. The authors conduct interesting ablation experiments to provide a better understanding of this method, such as the variations in reweight factors demonstrated in the experiments shown in Table 5.
summary: This paper identifies two drawbacks—objective misalignment and context inconsistency—within the current speculative approach (EAGLE) and proposes a progressive training framework, HASS, for better adapting draft models to actual decoding scenarios. The framework incorporates a distillation loss and a progressive training paradigm to reduce error accumulation caused by misaligned input features, empowering draft models to generate longer and more reliable token sequences. Experiments across various datasets and model sizes demonstrate that HASS significantly outperforms other speculative competitors and achieves impressive inference acceleration.
weaknesses: 1. The training efficiency of this method appears to have limitations. During the additional training steps (step 2, 3, 4...), each token’s corresponding KV matrix differs, making it infeasible to compute all attention scores through a single matrix multiplication. Could the authors offer a solution for this issue or provide detailed reports on the specific training time costs?
2. Section 3.2 (HARMONIZED CONTEXT ALIGNMENT) is the core of the proposed method, but the description is overly brief. To enhance readers' understanding, the authors could consider adding more technical details and method insights into the design choices. For instance, highlighting that different steps in model training focus on generating tokens for specific positions within the draft sequence would provide valuable context.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: - The authors state "It takes $n$ times the training overhead of EAGLE", suggesting HASS training steps have the same overhead as EAGLE training steps. However, HASS training steps (except step 0) actually have larger KV caches, attention scores, etc. For the $i$-th step, how much does the HASS training step increase compared to the EAGLE training step in terms of overhead (computational FLOPs, memory access bytes, GPU memory usage, actual execution time, and actual peak GPU memory)?

- The presentation could be clearer. In Figure 3, the light green text is difficult to read. In Table 1, the "-" marks for PLD and Lookahead Temperature=1 should be explained, considering readers may not be familiar with these methods.
soundness: 3
strengths: - The research addresses LLM latency reduction, which is highly relevant for LLM applications.

- The experiments are comprehensive, covering various tasks and models, achieving 3-4x speedup compared to vanilla decoding.

- The ablation studies are thorough. They clearly demonstrate the impact of HASS's two improvements and how different parameter settings affect the final results.
summary: This paper proposes HASS, a speculative sampling algorithm. HASS makes two main improvements:

- Training objective. The loss focuses on the k most likely tokens in the target distribution, rather than fitting the complete target distribution.

- Training strategy. Methods like EAGLE have inconsistent distributions between training and inference phases - training uses hidden states from target mode, while inference involves predicted values. HASS simulates this process during training by using model predictions as training data.

The paper evaluated HASS on dialogue, code, and math tasks. On LLaMA models, HASS improved the acceptance length by 8%-16% compared to EAGLE-2, achieving 2.81-4.05x speedup.
weaknesses: - HASS builds upon EAGLE but requires higher training overhead compared to EAGLE.
- The paper lacks comparisons with more related methods, such as vanilla speculative sampling and Medusa.

Thank you for the author's response. The additional results are consistent with my expectations, so I have decided to maintain my previous score.

Meta-review: The paper proposes an approach to enhance the speedup offered by speculative decoding. The idea is to reduce the discrepancy between drafter training and inference, particularly for the recent EAGLE series of methods. This involves employing a ranking-based distillation loss, and a training loss that incorporates the drafter's autoregressively generated features as additional context (to complement the verifier's features, as in regular EAGLE). These are shown to yield consistent gains over the EAGLE series of methods.

Reviewers were unanimously supportive of the paper. Speculative decoding is widely used to improve LLMs' inference speed, and further enhancing this technique is of clear interest. The core idea of reducing train-inference discrepancy is sensible, and the paper's empirical results are impressive.

Reviewers did point out one drawback of method being the increased training time, although the authors empirically demonstrated that it is not exorbitant (and has no impact on the decoding process). There were also suggestions on improving the clarity of presentation, which in the AC's reading could be further enhanced (see below). Overall, however, the paper's contribution is expected to be useful to the field.

_Suggestions for presentation_:
- Figure 1, font sizes can be increased
- Figure 3, several things are not clear. What do the dashed arrows emanating from $\mathcal{M}_s$ represent? Presumably the rectangular array on the right is the attention matrix; but why is there a single label for $q^{(s_i)}$? What is the significance of the colours changing across the timesteps?

Decision: Accept (Poster)