Learning Harmonized Representations for Speculative Sampling
Authors: Lefan Zhang
Xiaodan Wang
Yanhua Huang
Ruiwen Xu
Venue: ICLR 2025 Poster (2025)
Keywords: speculative sampling
large language model
Primary area: generative models

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%. The code is available at https://github.com/HArmonizedSS/HASS.

Published as a conference paper at ICLR 2025
LEARNING HARMONIZED REPRESENTATIONS FOR SPECULATIVE SAMPLING
Lefan Zhang, Xiaodan Wang, Yanhua Huang, Ruiwen Xu
Xiaohongshu Inc.
Shanghai, China
rui ruixenxx)@exiachongshh.
(lefan, xi xiaoaa2,, yanhuahuang,
ABSTRACT
Speculative sampling is promising approach accelerate the decoding stage for Large Language Models (LLMs). Recent advancements that leverage tar- get 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 specula- tive sampling methods. In this work, we propose a solution named HArmo- nized Speculative Sampling (HASS) that learns harmonized representations to address these issues. HASS accelerates the decoding stage without adding infer- ence overhead through harmonized objective distillation and harmonized context alignment. Experiments on four LLaMA models demonstrate that HASS achieves 2.81 2.81x-4.05x wall-clock time speedup ratio averaging across three datasets, sur- passing EAGLE-2 by 8%-20%. The code available at https: https://github. Com/HArmonizen zedss HASSS.
INTRODUCTION
Generative Large Language Models (LLMs), such as GPT-4 (Achiam et, 2023) and LLaMA (Tou- vron et al., 2023), have demonstrated remarkable capabilities across wide range of tasks. Never- theless, 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., 2023; Leviathan et al., 2023) offers a solution by leveraging addi- tional resources increase concurrency. Specifically, it employs efficient draft model generate draft tokens auto-regressively, which are then concurrently verified by the target LLM. Based on the verification results, subset of draft tokens that preserves the same distribution as the target LLM is accepted as the final output.
Leviathan al. (2023) show that the practical performance of speculative sampling is highly re- lated 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 pro- pose to leverage the target LLM's 's contextual information (Xiao et al., 2024; Li et al., 2024b;c; Du et al,, 2024). For instance, EAGLE (Li al., 2024b;c) 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 's hidden states in previous timesteps. However, during decoding, the draft model cannot access the target L.M''s hidden states for unverified timesteps, resulting in a con- text misalignment between training and decoding. This issue can be viewed bias (Bengio et al., 2015; Wang & Sennrich, 2020) at the feature level in speculative sampling.
form of exposure
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 assign high probabilities to (Li et al., 2024c; Miao et al., 2024; Sun et al., 2024). In this scenario, the draft model should focus more on recalling the desired tokens, while the specific
Corresponning author.

---

Published as a conference paper at ICLR 2025
40.
30.
30.
2,00
20
2--3c
215
1.0
05
00
11 LeaASiin
Figure 1: Speedup ratios of different methods on LLaMA2-Chat 7/13B and LLaMA3-Instrue 8/70B with temperature T E {0, 1), averaging over MT-bench, HumanEval, and GSM8K datasets.
order these tokens can be somewhat de-emphasized. Moreover, most LLM applications perform nucleus sampling (Holtzman et al., 2020) top-k sampling (Fan et al., 2018). For these decoding objectives, tokens with high probabilities play 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.
In this paper, we introduce HArmonized Spec- ulative Sampling (HASS), a novel approach designed to address the aforementioned is- sues by learning harmonized representations. Specifically, to make draft models aware of the decoding strategy, HASS extends the idea of ranking distillation (Tang & Wang, 2018) from the recommender system to speculative sam- pling, resulting in distillation loss focused on the most probable tokens within the tar- get distribution. To mitigate the previously discussed context misalignment between train- ing 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.
Unverilted tokens
Prediction Head & Sampling
Consistent forward
Training
Decoding forward
One Auto regreseegr Layer
,10
We conduct experiments across dialogue, code generation, and mathematical reasoning tasks using the MT-bench, HumanFivll, and GSMSK datasets, respectively. Building with EAGLE-2 (Li et al., 2024c), HASS achieves 8%-16% acceptance length improvement over it LLaMA2-Chal 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.
Target LLM
Figure We use EAGLE (Li al., 2024b) as an example to illustrate the context misalign- ment, where the speculation starts from timestep t. f()) and f(*) represent hidden states from the target LLM and the draft model. When decod- ing draft token Z442, the input context is incon- sistent between training and decoding.
PRELIMINARY
Speculative sampling leverages the concept of speculative execution (Kung & Robinson, 1981; Hennessy & Patterson, 2011) to reduce wall-clock time from more concurrency. Specifically, given the target L.M ( that is the focus of acceleration, speculative sampling employs draft model
2

---

Published as a conference paper at ICLR 2025
M to speculatively and efficiently generate draft tokens. The conventional approach (Leviathan et al., 2023; Chen et al., 2023) decomposes the next step generation into three steps:
-M proposes an unverified draft sequence with length L by auto-regressive decoding.
MM (1) evaluates posterior probabilities of I. draft tokens in parallel.
T 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. (2023) demonstrate that the wall-clock time improvement ratio is directly propor- tional to 7, while the arithmetic operation increment ratio is inversely proportional т. Conse- quently, T, 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., 2024; Li et al., 2024c;b; Sun et al., 2024). Note that 7 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 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 L.MM assigns high probabilities. However, previous speculative sampling works mainly focus on the entire vocabulary set w.r.t. knowledge dis- tillation from the target LLM (Li et al., 2024b; Zhou ct al., 2024), thereby disconnecting the training process from the practical decoding requirements.
EAGLE (Li et al., 2024b) is lightweight draft model design, shown in Figure 2. During decod- ing, it utilizes the LM Head of the target LLM to generate draft tokens. Specifically, we assume that the speculation starts from timestep 6, meaning the first draft token is at timestep + 1. To generate
the draft token T++1, the target LLM's hidden state f(() in the second-to-top layer is concatenated with the embedding e, to perform the input of the draft model. During training, EAGLE constructs regression task between f()ss and the predicted hidden states f(t)ss of the draft model. However, due to the auto regressseee decoding, the draft model only accesses the target LLM's features at the be- ginning of the speculation. It uses the features produced by itself 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., 2024b; Du al., 2024). EAGLE-2 (Li et al., 2024c) 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.
METHODOLOGY
As outlined before, previous speculative sampling methods suffer from disharmonies between train- ing and decoding. This section introduces HArmonized Speculative Sampling (HASS) to tackle objective misalignment and context inconsistency through harmonized objective distillation and har- monized context alignment, respectively, as described below.
3.1 HARMONIZED OBJECTIVE DISTILLATION
HASS prioritizes the most decoding recoding-desired tokens by leveraging the ranking distillation (Tang & Wang, 2018) idea from the recommender system. Specifically, ranking distillation aims to train 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 Ω C 2 where Ω represents the entire vocabulary. HASS considers the following Top-K distillation loss:
Lopp-
(x) Σ9(x) log p(x),
where g and 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 from hidden states of the target L.M. This implies that the proposed loss function benefits from the same efficient.
3

---

Published as a conference paper at ICLR 2025
training cost EAGLE. We evaluate the proposed Top-K distillation loss against six alternative losses, such as BiLD (Li et al., 2024a) and Recall @k Surrogate loss (Patel et al., 2022), through ablation studies.
3.2
HARMONIZED CONTEXT ALIGNMENT
6--2
F--1
Sex
HASS Tminnin Step
RAGLE Training
9--2
G--
45
fe-s
HASS Tminning Step
4-
(9
M0
6-4
HASS Training Step 3
fu2
Fe-1
Sat
12
(+)
4
x!ta
HASS Training Step
HASS EAGLE Decoding
Figure 3: Training with harmonized context alignment, where q and l refer to the query and key states in the transformer layer, respectively. Superseript (1) denotes tensors from the target LLM, and superscript (sy) denotes tensors from the j-th draft model forward. Note that during training (sj) refers to calling times draft model in batch, while during inference (sz) refers j-th auto- regressive decoding.
HASS follows a context alignment schema that aligns training and decoding on their contexts. The training procedure of HASS is divided into steps, enabling the draft model to utilize contex- tual features consistent with those in the decoding stage and addressing the context inconsistency by adapting the inaccurate features generated in previous HASS training steps. Specifically, it is achieved by first taking the inaccurate feature from the last draft model as query, and then consider- ing the inaccuracy accumulation in key-value part of the transformer block.
Formally. in the HASS training step j. given input token sequence II, T2, ... ZT, we optimize the draft model ((s) with the following objective function:
T-1
min Σ Min
(,,,,..,,, xx)) + Aux-loss), where [CrossEntropy (()) (11,1/21,... xt), p(o) (1 (()
(24411 xx) = Head(fle)) des))
From draft
target
From previous draft models
= Head(M((
(1) @ ... @ Je-j+1 J+1 1 @ ... 1 f (85-1)
query
key & value
4

---

In the second step, features from the first step are incorporated. For instance, in the self- attention mechanism at timestep t + p((1) used to calculate the current query. Keys and values are derived from f(( f(((), where f!() includes features from timesteps carlier Published as a conference paper at ICLR 2025
Temperature =
Temperature =
Model
Method
M'T-hench
HumanFval
GSM8K
Mean
MT-bench
HumanEval
GSMBK
Mean
PLD
1.43
1.59
1.37
1.46
Lookahead
1.66
1.77
1.65
1.69
SpS (V-68M)
2.02
2.03
2.04
2.03
1,72
1,30
1.65
1.62
SpS (L-68M)
1.83
1.81
1.83
1.82
1.47
1.36
1.46
1.43
L2
Medusa
2.34
2,48
2.37
2.40
2.35
2.56
2.40
2.44
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
PLD
1.46
1.70
1.44
1.53
Lookahead
1.64
1.85
1.69
1.73
SpS (V-68M)
2.13
2.61
2.21
2.32
1.73
2.25
1.81
1.93
SpS (L-68M)
1.83
1.70
1.73
1.50
1.34
1.45
1.43
L2 13B
Medusa
2.51
2.70
2.59
2.53
2.89
2.72
2.71
EAGLE
3.86
4.17
4.18
3.62
3,98
3,96
EAGLE-2
4.74
5.57
5.17
5.16
4.60
5.41
5.03
5.01
5.13
5.55
5.58
4.98
5.41
5.42
EAGLE
2.91
3.57
3.38
2.67
3.30
3.11
133
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
4.85
4.80
EAGLE
3.24
4.07
3.79
3.70
3.06
3.85
3.66
3.52
L3 70B
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 Acceptance lengths T of different methods on MT-bench, HumanEval, and GSM8K datasets with temperature T {0, 13. L2 represents LLaMA2-Chat, while L3 represents LLaMA3-Instruet. SpS stands for Vanilla Speculative Sampling, while V-68M and L-68M represent Vicuna-68M and LLaMA-68M, which the draft models of SpS.
p()) is the auto-regressive probability distribution provided by the target LLM, the Aux-loss consists of the proposed Top-K loss and the feature regression loss (following EAGLE), Head and @ stand for the language modeling head and the concatenation operation respectively. When training tokens in the entire sequence in parallel, the above formulation adapts the inaccurate features from previous j -1 steps for all positions except the first j positions. Note that compared to EAGLE, HASS takes additional training overhead due to the extra features, while maintaining the same decoding overhead. To accelerate the training procedure, we propose a modification to the attention mask mechanism, as outlined below:
training steps for adapting inaccurate
The first step mirrors the training stage of EAGLE. At timestep ( + 1, the draft model takes the target model's 's feature f(() as input and produces the draft feature f(at). In this step, the attention mask remains the same as the original causal mask without any modification.
than The attention mask is adjusted to ensure that the previous feature seen by g(x1) is always f(2) shown in the 'HASS Training Step 2' part of Figure 3.
For step j 3, the feature from the previous step f(e)-1) is utilized to calculate the query timestep 1++ while keys and values are generated by f 2 @ ((11) @... [ (())-1)
We empirically demonstrate that the acceleration effect converges with small n that the train- ing of HASS is cost-efficient. The actual training overhead of HASS in terms of training speed, computational cost, and GPU memory is investigated in Appendix A.8.
5

---

Published as a conference paper at ICLR 2025
Temperature =
Temperature =
Model
Method
M'T-bench
HumanFval
GSM8K
Mean
MT-bench
HumanEval
GSMBK
Mean
SpS (V-68M)
1,35x
1.38x
1.37x
1.37x
1.17x
1.02x
1.12x
1.10x
SpS (1-68M)
1,23x
1.24x
1.25x
1.24x
1.00x
0.94x
0.99x
0.98x
Medusa
1.91x
1.96x
2.20x
2.02x
2,00x
2.25x
2.15x
2.13x
L2
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
SpS (V-68M)
1.63x
1.98x
1.68x
1.76x
1.33x
1.72x
1.39x
1.48%
SpS (L-68M)
1.41x
1.29x
1.30x
1.33x
1.12x
1.04x
1.11x
1.09x
Medusa
2.26x
2.25x
2.71X
2.41x
2.31x
2.47x
2.36X
2.38x
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
EAGLE
1.29x
2.00x
1.85x
1.71x
1.25x
1.41x
1.67x
1.44x
1.3
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
EAGLE
2.14x
2.74x
2.42x
2.43x
1.80x
2.34x
2.12x
2.09x
L3
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
Speedup ratios of different methods on M'T-bench, HumanEval, and GSM8K datasets with temperature T 1). L2 represents LLaMA2-Chat while L3 represents LLaMA3-Instruc.. SpS stands for Vanilla Speculative Sampling, while V-68M and L-68M represent Vicuna-68M and LLaMA-68M, which are the draft models of SpS.
Table 2:
4 EXPERIMENT
4.1 EXPERIMENTAL SETUP
Target LLMs. LLaMA2-CHe 7/13B and LLaMA3-Instrut 8/70B.
Tasks. We conduct evaluations on three generation tasks. For multi-turn conversation, code genera- tion, and mathematical reasoning tasks, choose the MT-bench (Zheng et al., 2024), HumanEval (Chen et al., 2021),, and GSM8K (Cobbe et al., 2021) datasets, respectively. The batch size set as 1 under all the experiments following Leviathan al. (2023) and Zhou et al. (2024).
Metrics. HASS neither fine-tunes the target LLMs' weights during training nor relaxes the accep- tance conditions during decoding, making 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 speedup ratio relative vanilla auto-regressive decoding.
Acceptance Length r: The average number tokens generated per drafting-veriffcatiinn 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 numerical errors. Therefore, all inference processes are conducted on NVIDIA H800 GPU.
Comparisons. The vanilla auto-regressive decoding taken as the baseline, which serves as the benchmark for speedup ratios (1 (1.00x). We compare HASS with recent lossless speculative sampling methods, including PLD (Saxena, 2023), Lookahead (Fu et al., 2023), Vanilla Speculative Sampling (Chen et al., 2023), Medusa (Cai et al., 2024), EAGLE (Li et al., 2024b), and EAGLE-2 (Li et al., 2024c). PLD and Lookahead are free traning, which respectively use string matched from the prompt and cached n- n-grams draft tokens instead of generating draft tokens from draft model's predicted probability distribution. Therefore, the results of PLD and Lookahead under temperature I are not reported in Table 1.

---

Published as a conference paper at ICLR 2025
Implementation. Our code is built based on EAGLE-2's open-source repository'. Experiments on EAGLE and EAGLE-2 reuse draft model weights trained by Li et al. (2024b). For harmonized objective distillation, K is set 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, 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, i.e., the ShareGPT dataset with 68,000 dialogues, 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 acceptance length and speedup ratio. Then, conduct ablation studies on harmonized objective distillation and harmonized context alignment. Inspired by Yi et al. (2024), further conduct experiments by training on different proportions of the ShareGPT dataset to investigate HASS's scalability in the face of data sparsity (see Appendix A.6), and by evaluating on the translation tasks investigate HASS's robustness across different task types (see Appendix
A.7). As shown from the results, HASS is more scalable than EAGLE-2 with fewer training data and achieves promising improvements over EAGLE-2 on translation tasks consistent with results on MT-bench, HumanEval, and GSM8K.
4.2.1 EFFECTIVENESS
We present different methods' acceptance lengths and speedup ratios across three datasets in Tables and 2, respectively. HASS performs the largest acceptance length and highest speedup ratio across all datasets and L.MSs we tested. Most methods achieve their best performance on the HvmanEval dataset, the fixed templates in the code generation task are easier draft and accelerate. Though PLD and Lookahead are free of training, they consistently show poorer performance than Medusa, EAGLE, EAGLE-2, and HASS.
4.2.2 ABLATION STUDY ON HARMONIZED OBJECTIVE DISTILLATION
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 re- sults in Figure 4. Training with the Top-K loss (w >> always improves performance com- pared to training without the Top-K K loss (w 0). HASS achieves the largest acceptance length when = 0.5. A small value of K may result in performance degeneration, as the draft model only focuses on the token with the highest prob- ability and consequently neglects other potential tokens. With a larger K, the Top-K loss gener- ally brings better results, while the acceptance length is the largest when K 5.
6.3
5.3
Acceytance Length
52
6.1
511
4.9
4.9
0
0.1
0.2
0.5 Weight
1.0 2.0
10 100
Figure 4: Acceptance lengths T of HASS with varied Ks and weights of the Top-K loss. The results are conducted on LLaMA2-Cha 7B and averaged over MT-bench, HumanEval, and GSM8K datasets with temperature T E (0, 1}.
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 Ω is formed by the most probable tokens whose cumulative proba- bility is just larger than P.
Normed Top-K Loss, where the target and draft distributions are both normalized over Ω. The normalization operation can be either linear or sofumax.
htnps//gghhhb.cmmmsseeelll....ap hhtp:://hhggigggaacccooassssssss
Vicuna unfiltered
7

---

Published as a conference paper at ICLR 2025
Loss Function
Temperature
Temperature
Mean
Top-K 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- Loss
4.99
4.72
4.86
Recall @+k k Surrogate Loss BiLD Loss
4.97
4.76
4.87
5.04
4.75
4.90
Table 3: Acceptance lengths of HASS with different kinds of loss functions for harmonized ob- jective distillation. The results are conducted on LLaMA2-Chat 7B over the MTTbench dataset with temperature (0, 1).
Bi-directional Top-K Loss, where the distillation is conducted over the most probable kens w.r.t. the target distribution as well as the draft distribution.
Recall Surrogate Loss (Patel et al., 2022), where a smooth approximation of the recall metric is obtained and differentiable for direct optimization.
BiLD Loss (Li al., 2024a), 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 summa- rize their best results in Table 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 K loss 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 replaced by the dataset generated by the target LLM (see Appendix A.4). 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 -enerated data.
4.2.3 ABLATION STUDY ON HARMONIZED CONTEXT ALIGNMENT
Aligning Step
MT-bench
HumanEval
GSM8K
Mcan
EAGLE-2 Top-K
4.59
4.97
4.77
4.78
HASS Align-2
4.95
5.25
5.12
5.11
T=0
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
EAGLE-2 Top-K
4.46
4.61
4.64
4.57
HASS Align-2
4.71
4.89
4.98
4.86
T=l
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
Table 4: Acceptance lengths of HASS with varied aligning steps in the harmonized context align- ment. The results are conducted on ILaMA2-Chat 7B with temperature T (0, 1).
We propose the harmonized context alignment, which climinates 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, conduct experiments by varying the step number and summarize the results Table 4.
As the first training step of HASS is the same EAGLE-2, continually train EAGLE-2's draft model weights with the Top-K loss and consider the baseline. Without harmonized context align- ment (EAGLE-2 Top-K), the draft model performs the worst across all datasets. Training with 3/4
8

---

Published as a conference paper at ICLR 2025
steps of harmonized context alignment generally obtains the most considerable acceptance length. When training with steps of context alignment, the acceptance length decreases. We believe this caused by the draft model's 's limited capacity, it predicts less accurately former steps' tokens when paying too much attention to the latter ones. Figure 5 shows the acceptance rate a across speculation steps the MT-bench dataset following Li et al. (2024c). In later speculation steps, HASS performs better acceptance rates than EAGLE-2, demonstrating the effectiveness of harmo- nized context alignment.
75
85
65
60
60-
60
35
55
56
65
2-g
3-g
4.
0-a
1-a
6-6
2-0
3-0
4.
0-d
4-a
5-
LLaMAR-Chal T-0
LLaMA2-Chal T-0
Lnatua- 88
T-0
LiaMMa T-0
75-
75-
70
65
60
60
65
55
66
3-g
4-g
0-g
1-g
2-g
S-0
4-1
5-q
LLAMAD-Chal
LLaMAD-Chal T-1
LLaMAS-Instrud 88
LLaMA3-hnsiret
Figure 5: Acceptance rates a (%) of HASS and EAGLE-2 across different speculation steps on the MT-bench dataset with temperature T {0, 1].
As shown in Figure 5, the acceptance rates HASS decrease compared to EAGLE-2 on LLaMA2- Chat 13B and -.atttttttt 70B at the first step (0-a). The draft models degrade the first speculation step with much attention paid to the latter speculation steps, while the first step's ac- ceptance rates are crucial larger acceptance lengths. We conduct experiments emphasize the significance of former speculation steps by reweighting the training loss from each step with factor 3. specific, the step j's training loss will be multiplied by 3i-1 Table and Figure 6 show the ac- ceptance lengths and acceptance rates of HASS with different reweight factors on LLLMA3--Instuu 70B over the MT-bench dataset, respectively. With the factor B decreasing from 1.0 to 0.5, HASS achieves better acceptance lengths with different temperatures. Correspondingly, we perceive that acceptance rate at the first speculation step consistently higher with a smaller 3, while the acceptance rates at the latter speculation steps generally decline. When the factor / decreases 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 out of this paper's 's scope, we leave it for future work.
72
72
Reweight Factor /
T=0
T= 1
Mean
1-31
70
70
1.0 (Default)
4.62
4.59
4,61
0.7
4.65
4.61
4.63
68
68
0.5
4.67
4.62
4.65
0.3
4.65
4.61
4.63
66
1-0 LLaMAB-instucl Te0
LLaMkalinstuu T=1
Table 5: Acceptance lengths of HASS with different reweight factors 3 for harmonized con- text alignment. The results are conducted on LLaMA3-Instrict 70B over the MT-bench dataset with temperature T E (0, 1}.
Figure 6: Acceptance rates a (%) of HASS with different reweight factors 3 for harmonized context alignment. The results are conducted L.LMMA-Instrrtt 70B over the M'T-bench dataset with temperature T E 13.
5 RELATED WORK
There have been a number of works on improving the acceptance rate speculative sampling while maintaining the target distribution. Most of them fall into two categories. (1) The former category

---

Published as a conference paper at ICLR 2025
is aligned training that tries to obtain draft models aligned with the target LLM before the decoding stage. Zhou et al. (2024) propose a knowledge distillation approach and study several strategies to improve the aligmment. Li et al. (2024b) demonstrate that hidden states of the target LLM as input of the draft model provide extra feature uncertainty information. Xiao et al. (2024) also utilize hidden states of the target LLM and introduce RNN-based draft model design that achieves comparable acceptance rate. GLIDE (Du et al., 2024) 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. (2024) propose to utilize multiple draft models and design a tree- based attention mechanism to verify multiple draft sequences efficiently. Li al. (2024c) introduce dynamic structure to save computation by pruning inefficient paths in the draft tree. Sun et al. (2024) study improving the verification stage through optimal transportation. However, these works tend to only consider training decoding, ignoring the linkage of these two stages. This work instead aims to link training and decoding, leading to harmonized speculative sampling.
CONCLUSION
This paper introduces HASS, a harmonized speculative sampling solution that addresses dishar- monies 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. Experi- ments conducted on LLaMA2-Chat 7/13B and LLaMA3-Instruct 8/70B demonstrate the effective- ness and efficiency of HASS. Averaging on MT-bench, HumanEval, and GSM8K, HASS is 2.81x- 4.05x faster than vanilla rugo-segressive decoding, 20% faster than EAGLE-2.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv, 2023.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. In Advances in Neural Information Processing Sys- lems, volume 28, 2015.
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. In Proceedings of the 41st International Conference Machine Learning, 2024.
Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, fean-Baptiss Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv, 2023.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda. Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv, 2021.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. Training verifiers to solve math word problems. arXiv, 2021.
Cunxiao Du, Jing Jiang, Xu Yuanchen, Jiawei Wu, Sicheng Yu, Yongqi Li, Shenggui Li, Kai Xu, Liqiang Nie, Zhaopeng Tu, and Yang You. Glide with a cape: A low-hassle method to accelerate speculative decoding. In Proceedings of the 41st International Conference on Machine Learning, 2024.
Angela Fan, Mike Lewis, and Yann Dauphin. Hierarchical neural story generation. In Proceedings of the 56th Annual Meeting the Association for Computational Linguistics, pp. 889-898, 2018.
Yichao Fu, Peter Bailis, Ion Stoica, and Hao Zhang. Breaking the sequential dependency of I1m inference using lookahead decoding, November 2023. URL https: / //ms/s/mss//////////////////ss.. org/blog/ 22023-1-22 waokkaaaaddeccaa ng/.
10

---

Published as a conference paper at ICLR 2025
John L Hennessy and David A Patterson. Computer architecture: quantitative approach. Elsevier, 2011.
Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. In Proceedings of the 8th International Conference Learning Representations, 2020.
Assiang-Tsungy Kung and John T Robinson. On optimistic methods for concurrency control. ACM Transactions on Database Systems, 6(2):213-226, 1981.
Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. In Proceedings of the 40th International Conference on Machine Learning, pp. 19274- 19286. PMLR, 2023.
Minchong Li, Feng Zhou, and Xiaohui Song. Bild: Bi-directional logits difference loss for large language model distillation. arXiv, 2024a.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. Eagle: Speculative sampling requires rethinking feature uncertainty. In Proceedings of the 41st International Conference on Machine Learning, 2024b.
Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang, Eagle-2: Faster inference of language models with dynamic draft trees. arXiv, 2024c.
Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Zhengxin Zhang, Rae Ying Yee Wong, Alan Zhu, Lijie Yang, Xiaoxiang Shi, al. Specinfer: Accelerating large lan- guage model serving with tree-based speculative inference and verification. Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, pp. 932-949, 2024.
Yash Patel, Giorgos Tolias, and Jiří Matas. Recall @ k surrogate loss with large batches and sim- ilarity mixup. In Proceedings of the 2022 IEEE Conference on Computer Vision and Pattern Recognition, pp. 7502-7511, 2022.
Apoorv Saxena. Prompt lookup decoding, November 2023. URL https::t///gggg thub. com/ wttprr/uaannpppopt-eeeeeeee www.oooouuddddigg////oe
Ziteng Sun, Ananda Theertha Suresh, Jae Hun Ro, Ahmad Beirami, Himanshu Jain, and Felix Yu. Spectr: Fast speculative decoding via optimal transport. In Advances in Neural Information Processing Systems, volume 36, 2024.
Jiaxi Tang and Ke Wang. Ranking distillation: Learning compact ranking models with high per- formance for recommender system. In Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pp. 2289-2298, 2018.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv, 2023.
Chaojun Wang and Rico Sennrich. On exposure bias, hallucination and domain shift in neural ma- chine translation. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp.
2020.
Bin Xiao, Lujun Gui, Lei Su, and Weipeng Chen. Clover-2: Accurate inference for regressive lightweight speculative decoding. arXiv, 2024.
Euiin Yi, Taehyeon Kim, Hongseok Jeung, Du-Seong Chang, and Se- Young Yun. Towards fast multilingual L.M inference: Speculative decoding and specialized drafters. In Proceedings of the 2024 Conference on Empirical Methods Natural Language Processing, pp. 10789-10802, November 2024. URL https: /aclanttoooooyy...ooooo...oooolt org/ 2024 www............mmm.mmmmmm....... emnlp-manaee waill-mmai... 602.
Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, et al. Judging llm-as-a-judge with mt- bench and chatbot arena. In Advances in Neural Information Processing Systems, volume 36, 2024.

---

Published as a conference paper at ICLR 2025
Yongchao Zhou, Kaifeng Lyu, Ankit Singh Rawat, Aditya Krishna Menon, Afshin Rostamizadeh, Sanjiv Kumar, François Kagy, and Rishabh Agarwal. Distillspec: Improving speculative decoding via knowledge distillation. In Proceedings of the 12th International Conference on Learning Representations, 2024.
12

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 IMPLEMENTATION OF HARMONIZED CONTEXT ALIGNMENT
We present the pseudo code of harmonized context alignment, which implemented without the customized attention mask, for better understanding. The actual implementation in our experiments achieved by the customized attention mask as shown in Figure 3.
def Lewie
dract_nedel,
draft aodel
in_head,
aodel
aptiniser,
optianiee
Lorward_aun,
aligning stepo 2ed context alfonnn
stddee__taate__caree.. input__da,,
target LIM' reature
imput Loksna
hiidenn abe dree de. = []
10
Fax idk in (Esseessinn))
optinizer. poro_prad0)))
predict drodt_ stabde targes, irput__ds, dtf seddeen_ataaesssdlae midden_atates_drate wtstt__aat..ppoeoo thi sttoonnnttte_arreeeeeee
13
torch.cat thidder aattaaes iil, predictis, : :-11l, dio-l)
14
benget__eed, bred bead inchead (hidde aaateaasssssseeee caages.))gge)), Iinheed (predict)
losa feature (hidden saaten predict) logit_ (Larget_head, pead)
17
less.backwardd)
18
det attention(
attontica_mack, nisden__tatte__dratt List,
causai actention
fiet of deatt model"s testures
):
hidder_statts, target.
aidden_states___ target.enaae(...
query q_prot (nidden, states litt-1l (hidonn statonnssttt___arreeeeee key_t, value_t
states dratt
t_projthidses states, states_targgtt
stin_veight torch.matmul rquery, sey_t.tronspose22 12, 31) nath, squt lquary, stapel11) attention_nask
indices for hiddea_states sey_e, ind_g, ind_k s shsrst sereeen_ttte__steeeeeeeeee
enumerate (hiddeiddddddaaaeeeeeeaaaae statessdsttteeeeeeeeeeeeeee Mot
12
indicss dilly indicss|
13
attn_aeight_d torch.natmul feuezy, atta_reight atc_neight_oommel
313 nath.agrt (qaary.)
14
ind_k] attn_ssichs_di., cottnnx(aett_neeigtt. din--1)
ind_k]
16
atcn_output torch widht_wwettnnoommee value_t]
17
for i, biddoa_states_dret droft cammerate stataa Hott
18
value_d,
ind_k v_pcoj fiiddn atate__aaaaa indicsalid) indiosatsseg_een
attngautout i, ind_g] atth_veicht(s, ind_a, ind [..., Xona] (valoe_d[s, i, ind_k] valee_tli, i ind_k()
20
autn_output (_troo[(tin (ett (lansoose 2) 2e.sasha attr_octpet
seg_len, -1]]
13

---

Published as a conference paper at ICLR 2025
A.2 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 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.
Temperature 0
Temperature
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
Table 6: Acceptance lengths 7 of applying feature and token alignment to EAGLE-2's draft model weights, where "Token (x)) denotes tokens from training data being replaced by draft-model- generated tokens with probability of z. The results are conducted on LLaMA2-Chat 7B over the MT-bench dataset with temperature T E {0,1].
Feeaus Only
Festurs (22)
72
72-
85
68
84 L1 LLaA5k2-Chel 7B Tw0
LLaNdkg-Cro 78
Figure 7: Acceptance rates a (%) of applying feature and token alignment to EAGLE-2's draft model weights, where "Token (x)) denotes tokens from training data being replaced by draft-model- generated tokens with probability x. The results are conducted on LLaMA2-Chat 7B over MT-bench dataset with temperature T 10.
Feature only alignment brings the best performance, while adding token alignment leads to degen- eration. 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 result, training with tokens generated by the draft model in harmonized context alignment hurts the acceleration performance.
14

---

Published as a conference paper at ICLR 2025
A.3 HYPER-PARAMENN 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, present the speedup ratios and acceptance lengths of HASS with varied Ks and ws in Table 7.
MT-bench
HvaanEval
GSMSK
Mean
K
w
Speedup
Speedup
7
Speodup
Speedup
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
T=0
10
2.77x
4.93
3.38x
5.22
3.18x
5.11
3.11x
5.09
10
2,98x
4.96
3,40x
5.26
3,51x
5.16
3,30x
5.13
2.87x
4.98
3.41x
5.29
3.35k
5.16
3.21x
5.14
10
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
2.58x
4.70
2.79x
4.80
2.83x
4.95
2.73x
4.82
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
T=1
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
Table 7: Speedup ratios and acceptance lengths 7 of HASS with varied Ks and ws of the Top-K loss on LLAMA2-Chat 7B over MT-bench, HumanEval, and GSM8K datasets with temperature T {0, 1}.
15

---

Published as a conference paper at ICLR 2025
A.4 SELF- SILL-DISTILAAION
In the main text, we use the fixed ShareGPT dataset to train draft models for a fair comparison with EAGLE and EAGLE-2. Following existing speculative sampling methods (Zhou et al., 2024; Cai et al., 2024), we further use target-moodelgeneratee outputs to distill the draft model from the target model's ' real output distribution, dubbed as distiiiiilati.n. In specific, we feed the prompts from the ShareGPT dataset into the target models recursively with temperature set to O and collect the responses as multi-turn conversations for self-distillation following Li al. (2024b). To study the effect of reff-distillation, we conduct experiments on HASS and EAGLE-2 by training the draft model with fixed data or model-generated data and summarize the results in Table 8.
MT-hench
HumanFval
GSM8K
Mean
Model
Method
Data
Speedup
Spocdup
Specdup
Specdup
F
2.66x
4.44
3.06x
4.78
2.72x
4.60
2.81x
4.61
EAGLE-2
MG
2.86x
4.70
3.30x
5.12
3.03x
3.06x(+0.25)
4.94(+0.33)
L2 7B
F
2.99x
4.99
3,41x
5.29
3.32x
3.24x
5.15
HASS
MG
3.13x
5.25
3.85x
5.70
3.40x
5.57
3.46x(+00.22)
5.51(+0.36)
T=0
F
3.02x
4.74
3.64x
5.57
3.23x
5.17
3.30x
5.16
EAGLE-2
MG
3.04x
4.80
3.47x
5.46
3.19x
5.16
3.23x(-0.07)
5.14(-0.02)
L2 13B
F
3.23x
5.13
4.24x
6.05
3,48x
5,55
3.65x
5.58
HASS
MG
3.34x
5.27
4.42x
6.00
3.63x
5.61
3.80x(+0.15)
5.63(+0.05)
F
2.39x
4.23
2.87x
4.47
2.54x
4,50
2.60x
4,40
EAGLE-2
MG
2.49x
4.38
2.94x
4.73
2.69x
4.80
2.71x(+0.11)
4.64(+0.24)
L2
F
2.70x
4.84
3.13x
4.91
2.87x
5,01
2.90x
4.92
HASS
MG
2.75x
4.97
3.39x
5.24
3.13x
5.35
3.09x(+0.19)
5.19(+0227)
T=1
F 3.04x
4.60
3.45x
5.41
3.13x
5,03
3.21x
5.01
EAGLE-2
MG
3.08x
4.63
3.23x
5.25
3.04x
4.95
3.12x(-0.09)
4.94(-0.07)
L2 13B
F 3.28x
4.98
3.78x
5.86
3.37x
5,41
3.48x
5,42
HASS
MG
3.33x
5.02
3.76x
5.80
3.60x
5.42
3.56x(+0.08)
5.41(-0.01)
Table 8: Speedup ratios and acceptance lengths 7 of HASS and EAGLE-2 with fixed target- model-generated training data. F and MG stand for 'Fiedd' and 'Model-Generated"" respectively. L2 represents LLaMA2-Char
On LLaMA2-Chat 7B, self-distillation consistently brings improvements for HASS and EAGLE- 2. On LLaMA2-Chal 13B, self-distillation only achieves marginally better or comparable results, which is consistent with the observation from Li at. (2024b) ( ('data from the target LLM marginally improves performance' in its section 4.3.3). Especially, the acceptance lengths of the self-distilled HASS are lower than that of the vanilla HASS the HumanEval dataset, while both the speedup tios and acceptance lengths of the self-distileed EAGLE-2 are lower than that of the vanilla EAGLE-2 on HumanEval and GSMSK datasets. It may be due the code generation dataset HumanEval and the mathematical reasoning dataset GSMBK being less similar to the training dataset ShareGPT compared with MT-bench.
HASS outperforms EAGLE-2 on either fixed training data or model- training data. It noted that HASS trained on the fixed dataset even achieves better performance than EAGLE-2 data consistently. With self-distillation, HASS consistently achieves
trained on the more improvement or less degeneration in terms the acceptance length compared with EAGLE-2.
16

---

Published as a conference paper at ICLR 2025
A.5 DRAFTING HYPER-PARAMETE PARAMETERS
Li et al. (2024c) 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 '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-parameten surely leads to larger acceptance length. Nevertheless, sending more draft tokens into the target model for verification causes higher overhead in real applications. Therefore, we vary these hyper-parameters and report the speedup ratios Table to find a better trade-off.
Depch
100
T-0
1.2 7B
HASSS MG
EAGLE-2
EAGLE-2
2.45%
3.09t
3.02x
2.78x
3,0kx
2.61x
2.69%
3.22s
3 3.08%
2.69%
3.13x
2.66%
3. 19%
2.70%
3 122
2.79%
2.71x
3.14x
31
2.86x
3.29s
2.91%
2.95%
3.27%
2 6.60%
3.07
2.60%
3.16%
2 2.88%
3.31x
3,3%
2.8%
2.78x
2.28%
2.75x
2.54s
2.95.
2.55A
3.03x
2.65x
3.31x
3.34x
3.19%
3.42x
3.13x
3.37
2.70x
3.02x
1.2 7B
EAGLE-2
2,51x
2,57%
2,58%
2.36x
2 242%
2.39X
2.33%
2.40X
2,4%
2,60X
2.6%%
2,48x
3.36
2.48%
2 338
2646
2.25%
2.228
2.27% 2.42%
3.15%
HASS-MG
3.24x
3.30x
3.27%
3.19%
3.33k
3.33x
3.19%
3,208
3.26x
3.13x
3.26X
3.17%
2.62x
2.77%
2.84x
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 {0, where HASS-MG denotes HASS trained with self-distillation. L2 represents LLaMA2-Chat.
When 'depth' = 5, the acceptance length is relatively small. When 'depth' 9, the verification overhead is extremely high. Thus, neither of these settings achieves promising speedup ratio. For both HASS-MG and EAGLE-2, the best performances are achieved when 'depth' E {6, 7, and '# tokens' {60, 80, 100). HASS-MG consistently obtains superior performance compared with EAGLE-2 through hyper-parameter tuning across different LLMs and temperatures.
17

---

Published as a conference paper at ICLR 2025
A.6 NUMBER OF TRAINING TOKENS
Inspired by Yi et al. (2024), we randomly sample different proportions of the training dataset, i.e.. the ShareGPT dataset with 68, 68,000 dialogues. investigate the influences of training token numbers. In specific, we train the draft models of HASS and EAGLE-2 with 1/8, 1/4, 1/2 and the entire ShareGPT dataset and summarize the results in Figure 8 and Table 10.
Accoplance
1/9
118
114 ILaMAD-instuu
1/2
LLa Chat
LLaMM2 13B
LLaMA2 Instruct
Socod.p Raiio
1/8
1/4
1/1
1/8
1/4
12
1/1
115
164
1/2
1/1
1/8
1/4
1/2
CLaMkk--Chtt77
LLaMA2 Chat 13B
Lz
LLaMk2-instrud
Figure 8: Acceptance lengths 7 and speedup ratios of HASS and EAGLE-2 averaging across MT- bench, HumanEval, and GSM8K with different proportions of training dataset, i.e., the ShareGPT dataset with 68,000 dialogues.
As shown from Figure 8, HASS consistently outperforms EAGLE-2 under different proportions of training dataset with temperature T {0, 1j. HASS with merely 1/4 training dataset achieves better or comparable performance compared to EAGLE-2 with the entire training dataset, which demonstrates HASS's superior data exploitation and scalability obtained through further aligning objectives and contexts between training and decoding. The speedup ratio and acceptance length of HASS and RAGLE-2 approximately logarithmically proportional to the scale of training data, which is consistent with the finding in Yi et al. (2024). As shown from Table 10, the decrease in training data contributes more severe degradation on EAGLE-2 than that on HASS, reflecting HASS's S robustness to data sparsity.
18

---

Published as a conference paper at ICLR 2025
MT-hench
HumanEval
GSM8K
Mcan
Model
Method
Proportion
Speedup
Speedup
Speedup
Speedup
1/8
1.74x
3.06
3.39
3.24
3.23
EAGLE-2
1/4
1/2
2.08x
2.36x
4.11
3.64
2.49x
2.76x
4.46
3.93
2.64x
2.21x
4.28
3.81
2.26x
2.59x
3.79
4.28
1/1
2.66x
4.44
3.06x
4.78
2.72x
4.60
2.81x
4.61
L2
1/8
2.32x
3.92
2.69x
4.30
2.56x
4.12
2.52x
4.11
HASS
1/2
1/4
2.85x
2.64x
4.79
4.42
3.36x
3.05x
4.76
5.05
3.18x
2.70x
4.94
4.59
3.13x
2.80x
4,93
4.59
1/1
2.99x
4.99
3.41x
5.29
3.32x
5.17
3.24x
5.15
1/8
1.88x
3.25
2.39x
3.79
2.27x
3.57
2.18x
3.54
EAGLE-2
1/1
2.38x
3.82
2.85x
4.53
2.79%
4.25
2.67x
4.20
1/2
2.74x
4.32
3.46x
5.18
3.08x
4.81
3.09x
4.77
L2 13B
1/1
3.02x
4.74
3.64x
5.57
3.23x
5.17
3.30x
1/8
2.59x
4.01
3.37x
4.79
3.01x
4.43
2.99x
4.41
1/1
2,89x
4.55
3.59x
5.51
3.18x
5,00
3.22x
5.02
HASS
T=0
1/2
3.16x
3.59x
1/1
3.23x
5.13
4.24x
6.05
3.48x
5.55
3.65x
5,58
1/8
1.54x
2.77
2.00x
3.12
1.61x
2.83
1.72x
2,91
1/4
1.89x
2.30x
3.59
2.08x
3.35
2.09x
EAGLE-2
1/2
2.24x
3.68
2.64x
4.16
2.42x
3.91
2.43x
3.92
1/1
2.64x
4.21
3.31x
4.93
2.54x
4.42
2.83x
4.52
13 8B
1/8
2.14x
3.61
2.84x
4.22
2.30x
3.78
2.43x
3.87
HASS
1/1
2.46x
4.04
3.24x
4.83
2.48x
4.27
2.73x
4.38
1/2
2.72x
4.43
3.38x
5.28
2.99x
4.71
3.03x
4.81
1/1
2.78x
4.68
3.43x
5.54
3.06x
5.02
3.09x
5.08
1/8
2.09x
2.87
2.72x
3.45
2.29x
2.92
2.37x
3.08
EAGLE-2
1/2
1/4
2.76x
2.47x
3.33
3.74
3.71x
3.25x
4.57
4.01
3.11x
2.71x
3,96
3,46
3.19x
2.81x
4.09
3.60
1/1
2.94x
4,10
3.98x
5.02
3.19k
4.37
3.37x
4,50
1,3 70B
1/8
2.73x
3.68
3.79x
4.61
3.10x
3.95
3.21x
4,08
HASS
1/1
3.05x
4.12
4.23x
5.18
3.56x
4.54
3.61x
4.61
1/2
3,27x
4.40
4.52x
5.57
3.87x
4.92
3.89x
4.96
1/1
3,40x
4.62
4.68x
5.78
4.08x
5.24
4.05x
5.21
1/8
1.60x
2.99
1.90x
3.21
1.86x
3.17
1.79x
3.12
EAGLE-2
1/2
1/4
2.26x
1 .89
3.93
3.48
2.28%
2.57x
3.71
4.14
2.27k
2.46x
3.71
4.17
2.43x
2.15%
3.63
1/1
2.39x
4.23
2.87%
4.47
2.54x
4.50
2.60x
4,40
1/8
2.19x
3.82
2.50x
4.06
2.41x
4.03
2.37x
3.97
HASS
1/4
2,50x
4.27
2.84x
4.46
2.61x
4.52
2.65x
4.42
1/2
2.63x
4.56
3.10x
4.75
2.81x
4.81
2.85x
4.71
1/1
2.70x
4.84
3.13%
4.91
2.87x
5.01
2.90x
4.92
1/8
1.1x
3.16
2.18X
3.71
2.23x
3.46
2.11x
3.44
EAGLE-2
1/4
1/2
2.31%
2.78x
4.20
3.68
2.72x
3.19%
4.40
5.00
2.77x
3.08x
4.67
4.11
2.60x
3.02x
4,06
4.62
1/1
3.45x
1.2 13B
1/8
2.49%
3.94
2.98x
4.70
2.99x
4.33
2.82x
4.32
HASS
1/4
2.87x
4.43
3.40x
5.35
3.11x
4.87
3.13x
4.88
T=1
1/2
3.22x
4.75
3.70x
5.69
3.26x
5.18
3.39x
$.21
1/1
3.28x
4.9%
3.78x
5.86
3.37k
5.41
3.48x
5.42
1/8
1.51x
2.64
1.63x
3.01
1.64x
2.81
1.59x
2.82
1/4
1.77x
2.99
1.87%
3.51
1.92x
3.29
1.85x
3.26
EAGLE-2
1/2
1.90x
3.40
2.25x
4,03
2.38x
3.82
2.18x
3.75
1/1
2.39x
3.90
2.54x
4.73
2.48x
4.30
2.47x
4.31
L3 8B
1/8
1.96x
3.42
2.31x
2.25x
2.17%
3.75
1/4
2.22x
3.77
2.51x
4.63
2.45x
4.18
2.39x
4.19
HASS
1/2
2.43x
4.10
2.96x
5.07
2.82x
4.56
2.74x
4.58
1/1
2.49%
4.26
3.05x
5.30
2.89x
4.85
2.81x
4,80
1/8
2.16x
2.85
2.52x
3.35
2. 2.19k
2.91
2.29x
3.04
EAGLE-2
1/4
2.26x
3.29
3.01x
3.94
2.57x
3,44
2.61x
3.56
1/2
2.70x
3.67
3.37%
4.47
3.00x
3.94
3.02x
4,03
1/1
3.02x
4.00
3.61x
4.93
3.21x
4.35
3.28x
4.43
L3 70B
1/8
2.80x
3.70
3.46x
4.52
2.98x
3.93
3.08x
4.05
HASS
1/4
1/2
3.10x
3.28x
4.10
4.36
3.90x
4.15%
5.11
5.46
3.41x
3.72x
4.91
4.51
3.47x
3.72x
4.91
4.57
1/1
3.43x
4.59
4.25x
Table 10: Speedup ratios and acceptance lengths of HASS and EAGLE-2 with different propor- tions of training dataset, i.c., the ShareGPT dataset with 68,000 dialogues. L2 represents LLaMA2- Chat, while L3 represents LLaMA3-Instruet..
19

---

Published as a conference paper at ICLR 2025
A.7 EVALUATION ON TRANSLATION TASKS
To investigate the robustness of HASS across different task types, we further evaluate HASS and EAGLE-2 on five translation tasks by following Yi al. (2024). It is noted that both HASS and EAGLE-2 are trained on the fixed ShareGPT dataset without adaptation for translation tasks. We
conduct experiments on LLaMA2-Chat 7/13B and LLaMA3- 8/70B and summarize the re-
sults in Table 11.
De-+En
Fr-
Ja-+Fn
Ru-FEn
Zh-+En
Mean
Model
Method
Speedup
Speedup
Speedup
Speedup
Speedup
Speedup
EAGLE-2
2.58x
4.06
2.47x
3.99
2.46x
3.79
2.23x
3.48
2.39x
3.68
2.43x
3.80
L2 7B
HASS
3.15x
4.55
2.99x
2.97x
4.26
2.64x
2.83x
4.10
2.92x
4.27
EAGLE-2
2.95x
4.51
3.00x
2.67x
3.80
2.61x
2.60x
3.92
2.77x
4.06
L2 13B
HASS
3.63x
5.01
3.64x
4.94
3.05x
4.07
3.06x
3.02x
4.22
3.28x
4.45
T=0
EAGLE-2
2.59x
3.89
2.34x
1.91x
2.97
1.90x
2.03x
3.17
2.15x
3.45
L3 8B
HASS
2.98x
4.30
2.79x
2.21x
3.19
2.28x
2.32x
3.38
2.52x
3.72
EAGLE-2
3.10x
4.17
3.13x
4.07
2.35x
3.16
2.79x
3.76
2.52x
3.39
2.78x
3.71
L3 70B
HASS
3.75x
4.71
3.61x
4.47
2.67x
3.41
3.39x
4.25
2.84x
3.72
3.25x
4.12
EAGLE-2
2.26x
3.86
2.41x
3.91
2.09x
3.58
1.98x
3.34
2.25x
3.61
2.20x
3.66
L2
HASS
2.80x
4.44
2.99x
4.59
2.66x
4.11
2.53x
3.74
2.65x
4.05
2.73x
4.19
EAGLE-2
2.97x
4.29
2.77x
4.31
2.45x
3.73
2.33x
3.51
2.47x
3.72
2.60x
3.91
13B
HASS
3.45x
4.88
3.22x
4.84
3.02x
4.13
2.79x
3.97
2.83x
4.01
3.06x
4.37
EAGLE-2
2.23x
3.67
2.21x
3.69
1.85x
2.79
1.94x
3.15
1.89x
3.03
2.02x
3.27
L3
HASS
2.80x
4.13
2.73x
4.08
2.21x
3,18
2.35x
3.54
2.11x
3.34
2.44x
3.65
EAGLE-2
2.97x
4.02
2.95x
3.89
2.37x
3.16
2.72x
3.65
2.44x
3.34
2.69x
3.61
70B
HASS
3.71x
4.67
3.49x
4.38
2.83x
3.47
3.25x
4.11
2.75x
3.70
3.21x
4.07
Table 11: Speedup ratios and acceptance lengths 7 of HASS and EAGLE-2 on five translation tasks, where draft models are trained with the fixed ShareGPT dataset. De, Fr, Ja, Ru, Zh, and En stand for German, French, Japanese. Russian, Chinese, and English, respectively. L2 represents LLaMA2- Chat, while L3 represents LLaMA3-Instruct
As shown from Table 11, HASS consistently outperforms EAGLE-2 under all settings. HASS achieves 2.44x-3.28x wall-clock time speedup ratio averaging across five translation tasks, surpass- ing EAGLE-2 by 17%-24%. In terms of acceptance length, HASS achieves 8%-14% improvement over EAGLE-2. consistent with results on dialogue (MT-bench), code generation (HumanEval), and mathematical reasoning (GSM8K) tasks, HASS shows promising improvements over EAGLE-2 on translation tasks, reflecting its robustness across different task types.
20

---

Published as a conference paper at ICLR 2025
A.8 TRAINING OVERHEAD
As shown from Table 4, training with 3/4 steps of harmonized context alignment generally obtains the most considerable acceptance length, and so the aligning step of HASS fixed to 3 (Standard) in this paper unless stated otherwise. To investigate the actual training overhead of HASS, draft models for LLaMA2-Chat 7/13B and LLaMA3- 8/70B on single NVIDIA H800 GPU with batch size set and varied aligning steps, and summarize the results of training speed, computational cost, and GPU memory in Figures 9, 10, and 11, respectively. It is worth mentioning that the training overhead of HASS with 1 aligning step is the same as that of EAGLE-2.
train
LiaM12
LLaMA2
LLaMA3
LLaMA3 instuct
Troaking Triiking Sowerbaane
3.00
4.51
2:51
3:9
2:47
200
0.52
Aig-1 Mag2
Alg-3 Aiges4
Akgs-
44ig-2
AlanS Adg1-4
Alge-S
Aign-1 Algh22
Migh-3 Aign4
Agg
Ala-1
A0g122
Alghs Akg14 Aigh-5
Figure 9: Training speed (batch/s) of HASS with varied aligning steps, where the speed of Align-1 the same as that of EAGLE/EAGLE-2.
The training speed is evaluated by how many batches can be processed in one second, i.e., batch/s, and the ratio between Align-1 and Align-j represents how much training time needed for executing the same amount of training data compared to EAGLE-2. As shown from Figure 9, the training speed decreases with more aligning steps. However, the actual training time of standard HASS (Align-3) only 66.34% more than EAGLE-2 averaging over four target models, and the highest extra time cost compared to EAGLE-2 is just 91.47% (on LLaMA3- totally affordable, while HASS achieves superior performance and requires unchanged inference overhead.
8B). The training overhead of HASS
Cha 7B
LLa Cha 13B
LLaMV3-Indrudt 8B
LLaNA3--natuut 70B
Alig-5
Alre-5
A0gn4
A0ge-
Alge-
(EAGLE)
(SAGLE))
(EAGLE
40
103
45
Allgnd
Align5
Alge-3
Alge--
Allgnd
Aigh-
Alge-t
Aigh-5
Alge-3
Alie-5
Alge-3
Alg2
Alge-
Algh-
Allgnn
Alge-1
Alge-
1P.03%
TPEOPM
THLo
Figure 10: Training FLOPs of HASS with varied aligning steps, where the computational cost of Align-I is the same that of EAGLE/EAGLE-2. The upper figures show the FLOPs of the forward pass, while the lower figures show the FLOPs in total (forward and backward passes).
The computational cost is evaluated by TFLOPs and can be divided into forward pass and backward pass, we depict the cost of forward and backward passes in upper and lower figures Figure 10, respectively. The cost of forward pass is consisted of three parts:
Constant part is invariant to the number of aligning steps. Mapping target LLM's hidden state into q(x) (refer to section 3. with the LM head for distilling the draft model is included in constant part.
21

---

Published as a conference paper at ICLR 2025
Attention part is linearly proportional to the hidden state number fed into the draft model, which is accumulated across HASS training steps, i.e., Σi= for Align-j. Fusing token embeddings with hidden states, projecting hidden states into keys and values, and conduct- ing attention operations between query and several kv-pairs sourced from different hidden states are included in attention part.
Others linearly proportional to the number of aligning steps, i.e., j for Align-j. Compu- tational costs except for constant and attention parts included in others.
The cost of backward pass can be considered as (A + Q) x 2, where A and O represent attention part and others respectively, as computation constant part requires no gradient. Generally, the standard HASS (Align-3) requires approximately 3x computational cost of EAGLE-2.
LLaMA2-Chat 7B8
LLatAk2-Chst 1.385
LL LLaMA3--neett
LLaMAS-immtuct
60-
20
18-
32
(CAOLE)
(SAOLE)
(EAOLE)
(SAGUL))
Figure 11: Training GPU memory of HASS with varied aligning steps, where the GPU memory of Align-1 is the same as that of EAGLE/EAGLEE Avg and Peak stand for the average and peak GPU memory across the training process, respectively.
The GPU memory is evaluated by GB and report the average and peak GPU memory across the training process. Both the average and peak GPU memory increase with more aligning steps. The GPU memory requirement be covered by single NVIDIA H800 GPU even at Align-5 and batch size set to 2.
22


Reviews:
[{"review_id": "WgpCDrbPfv", "reviewer": "Reviewer_RwqU", "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.", "questions": "Seek weakness.\n\nWill update the score after looking at the results of Weakness.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The experiments are well-detailed, with clear metrics for comparison against baselines (e.g., EAGLE and multiple architectures).\n\nThe method makes sense in the perspective of typical acceptance process in SD.", "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].\n\n2. 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?\n\n[A] Yi et al., 2024. Towards Fast Multilingual LLM Inference: Speculative Decoding and Specialized Drafters, EMNLP 2024-main.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SfG3R0x03X", "reviewer": "Reviewer_6oJb", "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.", "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?\n- 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.\n\nMy 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).", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 3, "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.\n\n- Empirical results are convincing and demonstrate meaningful improvements in terms of acceptance rate and speedups\n- Authors present extensive ablation studies providing additional insights", "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)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gJG3NUfPpL", "reviewer": "Reviewer_ETQd", "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.", "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.\nIf the authors are willing to address our concerns, we are open to raise the scores during the rebuttal stage.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The motivation to improve the current speculative decoding method is clear and well-founded.\n2. The paper is well-written and easy to follow.\n3. 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).\n4. 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.", "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?\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HHlZG4Cy95", "reviewer": "Reviewer_p2CZ", "summary": "This paper proposes HASS, a speculative sampling algorithm. HASS makes two main improvements:\n\n- Training objective. The loss focuses on the k most likely tokens in the target distribution, rather than fitting the complete target distribution.\n\n- 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.\n\nThe 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.", "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)?\n\n- 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.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The research addresses LLM latency reduction, which is highly relevant for LLM applications.\n\n- The experiments are comprehensive, covering various tasks and models, achieving 3-4x speedup compared to vanilla decoding.\n\n- The ablation studies are thorough. They clearly demonstrate the impact of HASS's two improvements and how different parameter settings affect the final results.", "weaknesses": "- HASS builds upon EAGLE but requires higher training overhead compared to EAGLE.\n- The paper lacks comparisons with more related methods, such as vanilla speculative sampling and Medusa.\n\nThank you for the author's response. The additional results are consistent with my expectations, so I have decided to maintain my previous score.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)