Kangaroo: Lossless Self-Speculative Decoding for Accelerating LLMs via Double Early Exiting

We thank the reviewer for the detailed comments and valuable suggestions. We answer each of these comments below.

More explanation of the extension to tree decoding for Kangaroo

Thank you for your valuable suggestions. We will provide a more detailed formal expression and description of this extension in the final version. Kangaroo uses a static threshold to determine the timing of the second early stopping, based on the observation that the confidence of draft tokens in the small model is strongly correlated with their acceptance by the large model. Therefore, we approximate the probability of a token being accepted by the large model using the top-1 confidence of each token in the token tree. Considering the contextual dependency of tokens in speculative decoding, we model the probability of a token being accepted by the large model as the product of the top-1 confidences from the root node to that token. This approach is inspired by similar ideas in [1]

Sequoia: Scalable, Robust, and Hardware-aware Speculative Decoding.

The optimal settings for the early exit layer and the dynamic threshold $\eta$ may require careful tuning

Please refer to section 1.1 in the global response.

More experiments on a wider range of model sizes.

To demonstrate the generalizability of Kangaroo’s approach (reusing shallow parameters + learning a lightweight adapter), we conducted experiments on Llama2 and Llama3 as suggested by the reviewer. The results are as follows:

Speedup comparison of various speculative decoding methods on Spec-Bench [22] for Llama2-13B-Chat. Values outside the parentheses indicate speedup, while those inside parentheses indicate the compression rate (CR). $^*$ denotes reproduction result. SpS takes Llama-68M as the draft model.

We thank the reviewer for the detailed comments and valuable suggestions. We answer these comments below.

How did you determine the architecture and its suitability across different model architectures.

The iterative process of designing the adapter network in Kangaroo is reflected in Table 2 of our paper, as follows:

1. Initial Approach with MLP: We initially considered using a simple MLP to map shallow features to the LM head. However, this approach did not yield satisfactory results as it failed to leverage token context effectively.
2. Transformer Layer: Next, we replaced the MLP with a single transformer layer to enhance expressive power. This improved performance but still fell short of our expectations.
3. Removing FFN Module: Observing that the FFN module in the LLM Decoder significantly contributes to latency without leveraging token context, we removed the FFN module from the transformer and introduced a new LM Head to maintain parameter count. This strategy increased the speedup from 1.37 to 1.44.
4. Reusing LLM’s Language Head: Finally, we found that reusing the LLM’s own language head was highly effective, significantly boosting acceleration performance.

To demonstrate the generalizability of Kangaroo’s approach (reusing shallow parameters + learning a lightweight adapter), we conducted experiments on Llama2 and Llama3 as suggested by the reviewer. The results are as follows:

Speedup comparison of various speculative decoding methods on Spec-Bench [22] for Llama2-13B-Chat. Values outside the parentheses indicate speedup, while those inside parentheses indicate the compression rate (CR). $^*$ denotes reproduction result. SpS takes Llama-68M as the draft model.

Provide a more rational basis or empirical analysis for the chosen size of the shallow network.

Please refer to section 1.1 in the global response, where we conducted multiple comparative experiments across models with different architectures and sizes. It can be seen that deeper models (with depth $N$) have larger optimal early exit layers $\ell$, but the ratio $\ell / N$ is relatively consistent. This implies that the latency cost for the self-speculative small model compared to the full model is fairly stable. Therefore, we recommend setting $\ell / N$ between 1/16 and 1/10 in practical applications.

Regarding whether different tasks should have different optimal early exit layers $\ell$, we summarized the changes in speedup across different datasets with varying early exit layers for Kangaroo@Vicuna-7B, as shown in the table below:

It can be seen that the optimal early exit layers for different subtasks are not exactly the same but are generally close.

Verifying the speedup over multiple GPUs and diverse hardware setups.

We conducted inference on Vicuna-33B using 4 NVIDIA V100 GPUs, achieving a speed of approximately 7.6 tokens per second. With Kangaroo, using early exit at the 5th layer and training the adapter, the inference speed increased to 11.4 tokens per second, resulting in a speedup of about 1.5x.

Sincerely,

Paper 2124 Authors

Thank you for the detailed rebuttal, and I appreciate the effort in addressing the comments. However, I still have concerns and a few follow-up questions and suggestions:

1. Tree-Attention Experiments:

• I noticed that your added results primarily focus on the non-tree attention variant of Kangaroo. Could you clarify why the tree attention variant was not included in the comparisons? My hypothesis is that adding tree attention might shift the computation from being memory-bound to more compute-bound, potentially diminishing some of Kangaroo's advantages. Could you provide insights or results regarding this?

2. Temperature Sampling Experiments:

• The lack of experiments involving temperature > 0 during inference remains a concern. In real-world applications, varying temperature settings are commonly used, and understanding how Kangaroo performs under these conditions is crucial. Could you elaborate on why these experiments were not conducted or share any preliminary results if available?

3. Interpretation of Ratio $l/N$:

• You mentioned that the optimal early exit layers for different subtasks are "generally close." However, looking at the results, it seems that other methodologies might actually perform better in certain scenarios. Doesn't this indicate that the differences are more significant than "generally close"? Could you clarify this point?

4. Task-Specific Performance:

• The rebuttal tables highlight Kangaroo's strengths effectively, but I'm interested in understanding if there are scenarios or tasks where Kangaroo does not perform as well. Acknowledging these limitations is important, as no single method can excel in every task. It would be valuable to have a deeper discussion on where Kangaroo excels and where it might need further improvement. This would provide a more balanced and comprehensive perspective on the method's contributions.

I will definitely raise your score once I receive answers to these questions. However, for now, I will lower your score because I still have additional doubts and concerns.

Interpretation of Ratio $\frac{\ell}{N}$

By generally close, we refer to cases where the early exit layer falls within our recommended robust range of $[1/16, 1/10]$. For the 32-layer Vicuna-7B, only the second and third layers meet this criterion. On Vicuna-7B, the best early exit layers for Summarization and Math tasks fall on the third layer, but their optimal speedup performance (1.53x and 1.62x) is generally close to that of the speedup performance using the second layer for early exit (1.50x and 1.61x). It's unfair to compare the performance of other methodologies against Kangaroo when it exits early at the fifth layer on Vicuna-7B.

Task-Specific Performance

After a thorough analysis, we could explain why Medusa performs better on translation tasks, while Kangaroo excels in mathematical reasoning, summarization, and retrieval-augmented generation from two perspectives:

1. The usage of an auto-regressive adapter. It should be noted that the draft tokens in Kangaroo are generated in an auto-regressive manner, while Medusa uses time-independent FFN heads. Since Kangaroo can better leverage long-distance contextual relationships, it excels in subtasks with high similarities between input and output, such as summarization, retrieval-augmented generation, and mathematical reasoning. For the translation subtask, the input and output are the least similar.

2. The depth of the feature used for drafting. Medusa uses second-to-top-layer features, while Kangaroo performs early exits at shallow layers and trains a lightweight adapter to generate draft tokens. As Reviewer Tuhh also pointed out [3,4], the shallow layers of LLM may learn relatively simple tasks, while deep-layer features learn more abstract capabilities. Therefore, for the more abstract translation task, Medusa shows higher accuracy, reflected in its compression rate (CR) being higher than Kangaroo's. On the other hand, Medusa's compression rate is lower than Kangaroo's in all other subtasks.

[1] Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding

[2] Benjamin Spector and Chris Re. 2023. Accelerating LLM inference with staged speculative decoding

[3] The Unreasonable Ineffectiveness of the Deeper Layers

[4] Exploring Concept Depth: How Large Language Models Acquire Knowledge at Different Layers?

We would like to express our sincere appreciation again for the valuable suggestions from the reviewer, which have greatly enriched the quality of this manuscript.

Sincerely,

Paper 2124 Authors

We have conducted experiments for the case where temperature > 0 and explored a comparative analysis of the performance of Kangaroo and other methods across various tasks. We hope that the following responses can address your concerns.

Tree-Attention Experiments

In most of the rebuttal tables, we only reported the non-tree attention variant of Kangaroo. This is because most of the reviewers’ concerns in the rebuttal focused on the robustness of Kangaroo's two hyperparameters (i.e., $\eta$ and $\ell$), which are not closely related to the use of tree attention. Therefore, in the comparison table for Llama2-13B-Chat, we compared Kangaroo and three other algorithms (Draft & Verify[1], SpS, Medusa-1) that also do not use tree attention. It's indeed true that adding tree attention might shift the computation from being memory-bound to more compute-bound. However, there is no evidence that tree attention might potentially diminish some of Kangaroo's advantages. To alleviate your concerns, we have summarized the performance of Kangaroo with and without tree attention across different model architectures:

The second column Tree denotes whether tree attention is used. Values outside the parentheses indicate speedup, while those inside parentheses indicate the compression rate (CR).

Temperature Sampling Experiments

In line 118 of the main paper, we claimed that in this work we focus on greedy decoding while our methodology can be easily extended to the rejection sampling case. To further address your concern, we followed the setting in Draft & Verify and set temperature = 0.2 to conduct two sets of comparative experiments on Kangaroo (without tree attention). One set used the original top-1 confidence as the variable to determine early stopping, while the other used adjusted top-1 confidence (softmax of the adjusted logits divided by the temperature). The results are shown in the table below:

From the analysis of the table, it can be seen that Kangaroo still achieves a speedup effect very similar to that of greedy decoding when $T = 0.2$. Moreover, the top-1 confidence used for the second early stopping mechanism should remain unadjusted because when using temperatures less than 1, the adjusted top-1 confidence tends to be overestimated, making it more difficult to trigger the early stopping mechanism (even if we use a large threshold of $\eta = 0.96$ for the adjusted confidence). Besides, the acceleration effect of Kangaroo decreases with an increase in sampling temperature. This is attributed to the increased computational complexity of the speculative sampling criterion at higher temperatures, as revealed in prior research[2].

Dear Reviewer J1nq,

We greatly appreciate your ongoing efforts in helping us improve the quality of this manuscript. We will add more baseline methods and conduct a comprehensive comparison for temperature sampling in the final version as suggested.

We would like to point out that Medusa uses typical sampling when the temperature > 0, which is a lossy method that does not guarantee consistency with the original model’s sampling distribution. In contrast, Kangaroo is a lossless speculative decoding approach that exhibits a larger speedup ratio than Medusa. As claimed in Section 2.3.1 of Medusa[1], ``We ascertain that it is typically unnecessary to match the distribution of the original model. Thus, we propose employing a typical acceptance scheme to select plausible candidates rather than using rejection sampling.`

[1] MEDUSA: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

By the way, as you mentioned, could you please adjust the score to 5 in your review?

Sincerely,

Paper 2124 Authors

We thank the reviewer for the detailed comments and answer each of these comments below.

Evaluate more LLM families other than Vicuna

To validate the robustness of Kangaroo across different architectures, we additionally trained adapter networks on Llama2-13B-Chat and Llama-3-8B-Instruct and evaluated the performance on Spec-Bench. The results are shown in the tables below:

Speedup comparison of various speculative decoding methods on Spec-Bench [22] for Llama2-13B-Chat. Values outside the parentheses indicate speedup, while those inside parentheses indicate the compression rate (CR). $^*$ denotes reproduction result. SpS takes Llama-68M as the draft model.

How frequently does the top-1 probability of small model goes below the predefine threshold? Do you have some observations on the efficacy of double early exiting?

The frequency with which the small model's top-1 confidence falls below the predefined threshold varies depending on the model and the evaluation dataset. For instance, in Figure 4 of the main paper, about 40% of tokens have a confidence lower than the threshold $\eta$, but most of these tokens are "hard tokens" that would not be accepted by the large model. More confidence distribution graphs are available in the Figures in the global response PDF.

The purpose of the second early exit mechanism is to reduce the time spent on difficult tokens, allowing the model to speculate further on simpler token sequences. As shown in Figure 6(b), when the stride $\gamma = 6$ but the dynamic early exit mechanism is not used ($\eta = 0$), the speedup is 1.2. When $\eta = 0.6$, the speedup increases to 1.5, indicating that the second early exit mechanism can provide approximately a 25% performance improvement.

Requires some hyperparameters search (e.g., exit-layer l)

We argue that in practical applications, the early exit layer $\ell$ can be set based on an empirical ratio derived from the target model's depth. To determine the optimal depth for shallow sub-networks relative to the full model depth, we conducted multiple comparative experiments across models with different architectures and sizes. The table below records the average speedup achieved on Spec-Bench with early exits at various depths, and the final column shows the optimal ratio of the early exit layer to the total model depth.

It can be seen that deeper models (with depth $N$) have larger optimal early exit layers $\ell$, but the ratio $\ell / N$ is relatively consistent. This implies that the latency cost for the self-speculative small model compared to the full model is fairly stable. Therefore, we recommend setting $\ell / N$ between 1/16 and 1/10 in practical applications.

Dear Reviewer oUCN,

With the discussion phase nearing the end, we would like to know whether the responses have addressed your concerns.

Should this be the case, we are encouraged that you raise the final rating to reflect this.

If there are any remaining concerns, please let us know. We are more than willing to engage in further discussion and address any remaining concerns to the best of our abilities.

We are looking forward to your reply. Thank you for your efforts in this paper.

Best regards,

Authors

I appreciate the authors' feedback and acknowledge the generalizability of the method as well as the impact of its various hyperparameters. I will maintain my original rating for now but am open to revising the score if necessary during the reviewer-AC discussion phase.

We thank the reviewer for the detailed comments and answer each of these comments below.

The main difference between Kangaroo and Draft & Verify

Both Draft & Verify and Kangaroo are self-speculative decoding algorithms that construct a small model by reusing the parameters of the original model. However, the starting points differ: Draft & Verify aims to approximate the full model by removing redundant middle layers, leaving a small model with about half the layers of the original model. This creates a trade-off between the small model's accuracy and inference latency. In contrast, Kangaroo is designed to minimize the layers of shared parameters while maintaining small model accuracy. Therefore, Kangaroo reuses only the initial shallow layers of the LLM, which are crucial for feature extraction [3]. Additionally, the lightweight adapter in Kangaroo is designed to enhance the small model's accuracy without significantly impacting latency.

[3] Tang Y, Liu F, Ni Y, et al. Rethinking optimization and architecture for tiny language models[J]. ICML, 2024.

For the benefits of reusing the initial layers, please also refer to the question "Why starting from shallow layers works better?" raised by Reviewer NQcw.

Is the optimal layer dataset specific or domain specific?

Thank you for pointing out this insightful question and for the valuable references. According to [2], LLMs learn different difficulty levels of tasks at different layers. This raises the question of whether the optimal early exit layer varies across datasets (domains). To address this, we summarized the changes in speedup across different datasets with varying early exit layers for Kangaroo@Vicuna-7B, as shown in the table below:

It can be seen that the optimal early exit layers for different subtasks are not exactly the same but are generally close.

Is it also possible that the early exit layer also becomes dynamic?

Early exiting methods often exit at different layers based on the difficulty of the samples [4]. These methods save inference costs by exiting earlier for easier samples. However, in self-speculative inference, the small model only needs to align with the large model on simple tokens. Pursuing higher accuracy by exiting at deeper layers increases the small model's inference cost. Therefore, we chose to exit at a fixed shallow layer to balance inference cost and framework simplicity. Nevertheless, considering different exit layers within the shallow sub-network for different tokens is a valuable direction for future research.

[4] Teerapittayanon S, McDanel B, Kung H T. Branchynet: Fast inference via early exiting from deep neural networks[C]//2016 23rd international conference on pattern recognition (ICPR). IEEE, 2016: 2464-2469.

Sincerely,

Paper 2124 Authors

Dear Reviewer Tuhh,

With the discussion phase nearing the end, we would like to know whether the responses have addressed your concerns.

Should this be the case, we are encouraged that you raise the final rating to reflect this.

If there are any remaining concerns, please let us know. We are more than willing to engage in further discussion and address any remaining concerns to the best of our abilities.

We are looking forward to your reply. Thank you for your efforts in this paper.

Best regards,

Authors

We thank the reviewer for the detailed comments and answer each of these comments below.

1. Insufficient explanation and comparison to existing early-exiting methods.

We highlight the fundamental differences between Kangaroo and existing early-exiting methods in the following table:

SPEED specifically targets parameter-sharing (block-reuse) architectures, combining speculative decoding with pipelined execution for inference acceleration. Parameter sharing allows deeper decoders to enhance model performance. The combination of speculative decoding and pipelined execution offers SPEED a speed advantage over larger models without parameter sharing. However, Draft & Verify and Kangaroo are architecture-agnostic and can be plugged into mainstream large language models (LLMs). Draft & Verify employs Bayesian optimization to pre-select less critical layers in the target LLM for skipping, while Kangaroo implements early exits in the continuous shallow layers of the target LLM and uses a lightweight adapter to bridge the expressiveness gap between the self-speculative small model and the large model. Due to SPEED’s limited applicability and there is no released code, we focus on a detailed performance comparison between Kangaroo and Draft & Verify.

We used the official code from the Draft & Verify repository to optimize the redundant layers for Vicuna-7B and used the official optimized layers for Llama2-13B-Chat. Values outside the parentheses indicate speedup, while those inside parentheses indicate the compression rate (CR).

While Draft & Verify achieves a higher compression rate, indicating the retained layers in the draft model better approximate the full model, the high latency of the small model limits its end-to-end speedup.

2. Why starting from shallow layers works better?

Theoretical analysis of speculative inference [3] highlights two key factors: the accuracy of the small model and its inference latency. Although Draft & Verify benefits from higher accuracy by leveraging deeper features, its small model incurs a high cost, almost half that of the large model (compared to about 1/10 in Kangaroo), limiting its overall speedup.

The reasons for the effectiveness of shallow layers in Kangaroo may include:

• The initial layers directly connect token embeddings to subsequent decoder layers, playing a critical role in feature representation. References provided by Reviewer Tuhh [4, 5] suggest that the shallow layers of LLMs are sufficient for understanding simple tasks, aligning with the requirement for speculative decoding where the draft model needs only to align with the large model on simple tokens.
• The autoregressive adapter used in Kangaroo enhances the expressiveness of shallow layers, and its lightweight design reduces inference latency.

[3] Leviathan Y, Kalman M, Matias Y. Fast inference from transformers via speculative decoding[C]. ICML 2023.

[4] The Unreasonable Ineffectiveness of the Deeper Layers

[5] Exploring Concept Depth: How Large Language Models Acquire Knowledge at Different Layers?

3. Do different down-streaming tasks have different favorable thresholds and distribution?

As shown in Figure 4 in the main paper, "Accept" and "Reject" draft tokens cannot be perfectly separated by a fixed threshold. However, selecting the intersection of these distributions as the threshold allows us to retain a high proportion of tokens likely to be accepted while minimizing the proportion of tokens likely to be rejected. To explore the sensitivity of this threshold, we visualized the conditional distributions of the small model's Top-1 confidence across different model architectures and subtasks (see Figures in the global response PDF). The confidence distributions across different subtasks suggest that selecting $\eta \in [0.6, 0.8]$ is a robust choice.

4. Robustness and sensitivity of Kangaroo's hyperparameters

Please refer to the common concerns in the global response.

Reproducibility

We will make our training and evaluation code publicly available in the camera-ready version.

Dear Reviewer NQcw,

With the discussion phase nearing the end, we would like to know whether the responses have addressed your concerns.

Should this be the case, we are encouraged that you raise the final rating to reflect this.

If there are any remaining concerns, please let us know. We are more than willing to engage in further discussion and address any remaining concerns to the best of our abilities.

We are looking forward to your reply. Thank you for your efforts in this paper.

Best regards,

Authors