Learning Harmonized Representations for Speculative Sampling

Q3: "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."

A3: Thanks for your valuable suggestions.

We will replace the colors in Figure 3 with more distinguishable ones.

PLD and Lookahead are free of traning, which respectively use string matched from the prompt and cached n-grams as draft tokens instead of generating draft tokens from a draft model's predicted probability distribution. Therefore, the results of PLD and Lookahead under temperature=1 are not reported in Table 1.

Q2: "The paper lacks comparisons with more related methods, such as vanilla speculative sampling and Medusa."

A2: We choose EAGLE-2 as one of our baselines, because it is considered as one of the fastest speculative sampling methods which consistently outperforms Vanilla Speculative Sampling and Medusa in EAGLE-2's original paper [A].

Thanks for your valuable suggestion and we are working on the results of Vanilla Speculative Sampling and Medusa. We will update once the results are available.

[A] Li et al., 2024. EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees, EMNLP 2024.

A1-b: We show the computational cost (TFLOPs) for training HASS in the following table, and more detailed analyses can be found in Figure 10 Appendix A.8.

As shown from the table, the standard HASS (Align-3) requires approximately 3x computational cost of EAGLE-2.

A1-c: We show the average and peak GPU memory (GB) across the training process of HASS in the following table.

As shown from the table, both the average and peak GPU memory increase with more aligning steps. The GPU memory requirement can be covered by a single NVIDIA H800 GPU even at Align-5 and batch size set to 2.

Thanks for your valuable comments.

Q1: "HASS builds upon EAGLE but requires higher training overhead compared to EAGLE." & "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)?"

A1-a: It is worth mentioning that the harmonized context alignment's acceleration effect converges with 3/4 HASS training steps (refer to HASS's section 4.2.3) and all the experiments in our paper are conducted with 3 aligning steps unless stated otherwise. While the training overhead of HASS increases compared to EAGLE-2, it requires no extra inference overhead, and can be considered as a "free lunch" as mentioned by reviewer ETQd.

To investigate the actual training overhead of HASS, we train draft models on a single NVIDIA H800 GPU with batch size set to 2 and varied aligning steps, and summarize the results of training speed, computational cost, and GPU memory in Appendix A.8 Training Overhead. It is noted that the training overhead of HASS with 1 alignment step is the same as that of EAGLE-2, and HASS's aligning step is always set to 3 (Standard) in all the experiments unless stated otherwise.

We show the actual training speed (batch per second) and time ratio w.r.t. EAGLE-2 of HASS in the following table. 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 the table, the training speed decreases with more aligning steps. However, the actual training time of standard HASS (Align-3) is only 66.34% more than EAGLE-2 averaging over four target LLMs, and the highest extra time cost compared to EAGLE-2 is just 91.47% (on LLaMA3-Instruct 8B). The training overhead of HASS is totally affordable, while HASS achieves superior performance and requires unchanged inference overhead.

Q3: "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."

A3: The vertical axis of the line charts in Figures 5 and 6 represents the acceptance rate $\alpha$, which has been transformed into percentage format. Since tree drafts sample multiple tokens per location with only one accepted, we follow EAGLE and utilize chain drafts without tree attention when measuring this metric, which is aligned with Speculative Sampling and Distillspec.

Q2: "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."

A2: Recall that the context inconsistency between training and decoding comes from the inaccuracies in features generated by the draft model compared to the features from the target LLM. The proposed method, i.e. harmonized context alignment, aims to address this inconsistency by adapting the inaccurate features in previous timesteps. It is achieved by first taking the inaccurate feature from the last draft model as query, and then considering the inaccuracy accumulation in key-value part of the transformer block. Formally, in the HASS training step $j$, given input token sequence $x_1, x_2, .., x_T$, we optimize the draft model $\mathcal{M}^{(s)}$ with the objective function $\min_{\mathcal{M}^{(s)}} \sum_{t=1}^{T-1} [\text{CrossEntropy}(P^{(l)}(x_{t+1}|x_1,\dots,x_t), P^{(s)}(x_{t+1}|x_1, \dots, x_t)) + \text{Aux-loss}],$ where

and $P^{(l)}$ is the auto-regressive probability distribution provided by the target LLM, and the $\text{Aux-loss}$ consists of the proposed Top-K loss and the feature regression loss (following EAGLE). When training tokens in the entire sequence in parallel, the above formulation adapts the inaccurate features in previous $j - 1$ steps for all positions except the first $j-1$ positions.

To be more specifically, HASS training step 1 takes the target LLM's feature as input, and focuses on generating draft tokens in the draft model's first auto-regressive speculation during decoding; HASS training step 2 takes the target LLM's feature as well as the last step draft model's feature as input, and focuses on generating draft tokens in the draft model's second auto-regressive speculation during decoding; HASS training step $j$ takes the target LLM's feature as well as features from draft models in previous training steps as input, and focuses on generating draft tokens in the draft model's $j$-th auto-regressive speculation during decoding. Therefore, HASS adapts the inaccuracy accumulation from previous speculations during decoding.

Thanks for your valuable comments.

Q1: "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?"

A1: In the attention operation, we add an extra dimension into the query and key matrices to adapt them to more than one piece of feature as input, i.e., the query and key matrices are shaped as (1, batch_size, num_head, seq_len, head_dim) and (feat_num, batch_size, num_head, seq_len, head_dim), respectively. Thus, we can compute all attention scores through a single matrix multiplication. You can refer to line 311 in https://github.com/HArmonizedSS/HASS/blob/main/model/cnets_hass.py for the real implementation.

To investigate the training overhead of HASS, we train draft models on a single NVIDIA H800 GPU with batch size set to 2 and varied aligning steps, and summarize the results in Appendix A.8 Training Overhead. It is worth mentioning that the training overhead of HASS with 1 alignment step is the same as that of EAGLE-2, and HASS's aligning step is always set to 3 (Standard) in all the experiments unless stated otherwise.

We show the actual training speed (batch per second) and time ratio w.r.t. EAGLE-2 of HASS in the following table. 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 the table, the training speed decreases with more aligning steps. However, the actual training time of standard HASS (Align-3) is only 66.34% more than EAGLE-2 averaging over four target LLMs, and the highest extra time cost compared to EAGLE-2 is just 91.47% (on LLaMA3-Instruct 8B). The training overhead of HASS is totally affordable, while HASS achieves superior performance and requires unchanged inference overhead.

Thank you for the detailed response, which effectively addresses my concerns. I have updated my score to 8 accordingly.

Q4: "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."

A4: We train EAGLE-2 on the model-generated data and report its comparison with HASS in Table 8 Appendix A.4 Self-Distillation. We show the speedup ratios and acceptance lengths $\tau$ of HASS and EAGLE-2 trained with fixed and model-generated data averaged across MT-bench, HumanEval, and GSM8K in the following table.

As shown from the table, HASS outperforms EAGLE-2 on either fixed training data or model-generated training data. It is noted that HASS trained on the fixed dataset even achieves better performance than EAGLE-2 trained on the model-generated data consistently. With model-generated data for training, HASS consistently achieves more improvement or less degeneration in terms of the acceptance length compared with EAGLE-2. Therefore, context misalignment is not mitigated by training on model-generated data.

Q3: "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?"

A3: In all the experiments, we use the fixed ShareGPT dataset with 68,000 dialogues for training draft models by following [B, C, D], unless stated otherwise.

We evaluate HASS and EAGLE-2 trained from ShareGPT on dialogue (MT-bench), code generation (HumanEval), and mathematical reasoning (GSM8K) tasks in the main text. We also evaluate HASS and EAGLE-2 trained from ShareGPT on translation tasks following [A] in Appendix A.7 Evaluation on Translation Tasks.

We show the speedup ratios and acceptance lengths $\tau$ of HASS and EAGLE-2 averaged across MT-bench, HumanEval, and GSM8K (MT+HE+GSM) and averaged across five translation tasks (Translation) in the following table.

Both HASS and EAGLE-2 demonstrate generalization ability across tasks different from the training task, while HASS consistently outperforms EAGLE-2 with a promising improvement.

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

[B] Li et al., 2024. EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty, ICML 2024.

[C] Cai et al., 2024. MEDUSA: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads, ICML 2024.

[D] Du et al., 2024. GLIDE with a CAPE: A Low-Hassle Method to Accelerate Speculative Decoding, ICML 2024.

Thanks for your valuable comments.

Q1: "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."

A1: Recall that the context inconsistency between training and decoding comes from the inaccuracies in features generated by the draft model compared to the features from the target LLM. The proposed method, i.e. harmonized context alignment, aims to address this inconsistency by adapting the inaccurate features in previous timesteps. It is achieved by first taking the inaccurate feature from the last draft model as query, and then considering the inaccuracy accumulation in key-value part of the transformer block. Formally, in the HASS training step $j$, given input token sequence $x_1, x_2, .., x_T$, we optimize the draft model $\mathcal{M}^{(s)}$ with the objective function

$\min_{\mathcal{M}^{(s)}} \sum_{t=1}^{T-1} [\text{CrossEntropy}(P^{(l)}(x_{t+1}|x_1,\dots,x_t), P^{(s)}(x_{t+1}|x_1, \dots, x_t)) + \text{Aux-loss}],$

where

and $P^{(l)}$ is the auto-regressive probability distribution provided by the target LLM, and the $\text{Aux-loss}$ consists of the proposed Top-K loss and the feature regression loss (following EAGLE). When training tokens in the entire sequence in parallel, the above formulation adapts the inaccurate features in previous $j - 1$ steps for all positions except the first $j-1$ positions.

Q2: "I am also confused by Figure 3 - should Training Step 2 have $f_{t-2}^{(l)}$ and $f_{t-1}^{(l)}$ in the bottom row and only use superscript $(s_1)$ in the bottom right entry $f_{t}^{(s_1)}$ (analogously for steps 3 and 4)?"

A2: In HASS training step 2, if we only use superscript $(s_1)$ in the bottom right entry $f_t^{(s_1)}$ and use superscript $(l)$ for the others, merely the last position of the input sequence can adapt to the inaccurate feature generated by last draft, while the rest positions degenerate to HASS training step 1. Oppositely, if we use superscript $(s_1)$ for the last three entries in the bottom row, the entire sequence can adapt to the inaccurate features generated by last draft except for the first position.

Q3: "The paper lacks experiments analyzing how the framework performs across different ... task types [Referring to A]. Adding these analyses could provide a more comprehensive understanding of the proposed method's scalability and robustness. Are there certain task types where the proposed method excels or struggles compared to baselines?"

A3: In the main text, HASS is trained on the fixed ShareGPT dataset with 68,000 dialogues and evaluated on three different task types, i.e., dialogue (MT-bench), code generation (HumanEval), and mathematical reasoning (GSM8K), where the proposed HASS consistently outperforms other baselines.

In Appendix A.7 Evaluation on Translation Tasks, we further evaluate HASS and EAGLE-2 trained from the fixed ShareGPT dataset on five translation tasks (https://github.com/Kthyeon/Multilingual-SpecBench/tree/main/dataset/eval) by following [A].

We show the speedup ratios and acceptance lengths $\tau$ of HASS and EAGLE-2 averaged across five translation tasks, i.e., De$\rightarrow$En, Fr$\rightarrow$En, Ja$\rightarrow$En, Ru$\rightarrow$En, Zh$\rightarrow$En, in the following table. It is noted that De, Fr, Ja, Ru, Zh, and En stand for German, French, Japanese, Russian, Chinese, and English, respectively. Please refer to Table 11 in Appendix A.7 for results on each translation task.

As shown from the table, HASS consistently outperforms EAGLE-2 under all settings. HASS achieves 2.44x-3.28x wall-clock time speedup ratio averaging across five translation tasks, surpassing EAGLE-2 by 17%-24%. In terms of acceptance length, HASS achieves 8%-14% improvement over EAGLE-2. In consistent with results on MT-bench, HumanEval, and GSM8K, HASS shows promising improvements over EAGLE-2 on translation tasks, reflecting its robustness across different task types.

Q2: "The paper lacks experiments analyzing how the framework performs across different token counts ... [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?"

A2: To investigate the influences of training token numbers as inspired by [A], we randomly sample different proportions of the training dataset, i.e., the ShareGPT dataset with 68,000 dialogues. In specific, we train draft models of HASS and EAGLE-2 with 1/8, 1/4, 1/2 and the entire ShareGPT dataset and summarize the results in Appendix A.6 Number of Training Tokens.

We show the speedup ratios and acceptance lengths $\tau$ of HASS and EAGLE-2 averaged across MT-bench, HumanEval, and GSM8K with different proportions of ShareGPT as training dataset in the following table.

As shown from the table, HASS consistently outperforms EAGLE-2 under different proportions of training dataset with temperature $\in \\{ 0,1 \\}$. 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 on objectives and contexts between training and decoding. The speedup ratio and acceptance length of HASS and EAGLE-2 are approximately logarithmically proportional to the scale of training data, which is consistent with the finding in [A], and so we are happy to cite [A] in our paper.

Thanks for your valuable comments.

Q1: "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]."

A1: It is worth mentioning that we only use the fixed ShareGPT dataset (https://huggingface.co/datasets/Aeala/ShareGPT_Vicuna_unfiltered) to train draft models for all the experiments unless stated otherwise, which is aligned with the settings of [B, C, D]. We also contruct self-distilled ShareGPT dataset by feeding the prompts from ShareGPT into the target LLM recursively with temperature set to 0 and collecting the responses as multi-turn conversations following [B], and summarize the results of HASS and EAGLE-2 trained on the self-distillation dataset in Appendix A.4 Self-Distillation.

We show the speedup ratios and acceptance lengths $\tau$ of HASS and EAGLE-2 averaged across MT-bench, HumanEval, and GSM8K with fixed and self-distilled ShareGPT as training dataset in the following table, where F and MG stand for Fixed and Model-Generated, respectively.

Self-distillation consistently brings improvements for HASS and EAGLE-2 on LLaMA2-Chat 7B, whille it only achieves marginally better or comparable results on LLaMA2-Chat 13B, which is consistent with the observation from [B] ("data from the target LLM marginally improves performance" in [B] section 4.3.3). This may be due to the code generation dataset HumanEval and the mathematical reasoning dataset GSM8K being less similar to the training dataset ShareGPT compared with MT-bench. There could be more ways to construct the self-distillation dataset such as varying the temperature during generation, but they are out of the scope of this paper. Our main focus is addressing the disharmonies between the draft model's training and decoding on their objectives and contexts, which yields an auto-regressive draft model having better alignment with the target LLM. With affordable training overhead, HASS requires no extra inference overhead compared to EAGLE-2 and consistently outperforms exisiting baselines.

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

[B] Li et al., 2024. EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty, ICML 2024.

[C] Cai et al., 2024. MEDUSA: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads, ICML 2024.

[D] Du et al., 2024. GLIDE with a CAPE: A Low-Hassle Method to Accelerate Speculative Decoding, ICML 2024.