Paper: SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference (ICLR-2024)

Reviewer comment (Experiments):
The SkipDecode algorithm is compared to a single algorithm and the provided results are limited. While the SkipDecode algorithm was evaluated on three different datasets with three different evaluation metrics on each, the comparison evaluation includes only two different datasets with a single metric for each dataset. Adding more comparison results would help.

Author response:
We do not know other work comparable to ours in terms of early exit in the decoding stage. CALM is the closest we are aware of.

Our work is one of the first to focus on early exit strategies in the decoding stage of autoregresive neural networks. There are limited studies that align closely with our approach. 

We identified CALM as the closest relevant work to our approach. However, CALM is designed for encoder-decoder architectures. To facilitate a meaningful comparison, we implemented a decoding-only version of CALM, which we believe is the most appropriate baseline given the specific focus of our research. The scarcity of directly comparable works in terms of early exit strategies in decoding stages has limited our ability to conduct a broader comparative analysis with a range of baselines.

The adaptation of CALM to a decoding-only context was necessary to provide at least one viable point of comparison. We acknowledge that this adaptation might not fully capture the nuances of other potential baselines designed explicitly for decoding-only architectures.

While Skipdecode is implemented with the Metaseq framework and therefore we can handle a wide range of large datasets, our implementation of CALM is not as optimized as Metaseq, that's why we limited the comparison to fewer datasets and sizes.

Reviewer comment (Experiments):
How are the min_exit_layer and max_exit_layer chosen?

Author response:
Our approach to selecting min_exit_layer and max_exit_layer is driven by the target speedup, such as 2x. A given target speedup will give us different min_exit_layer and max_exit_layer configurations based on Figure 3 (Section 2.3). We determine these configurations through hyperparameter search on the E2E dataset and then apply the same configurations across other datasets. We touched on this in Section 3.1 (second paragraph) of our manuscript, **we will revise this section to provide a clearer and more explicit explanation of our method for choosing these layers.**

Reviewer comment (Experiments):
What is the degradation of quality of using batch-wise exit-point function instead of a per example-wise?

Author response:
In our case there is no degradation per se as it is a position wise exit. In the SkipDecode, we employ a position-wise exit strategy, meaning that tokens at the same position in different elements of the batch exit simultaneously. Our method maintains consistency across batch elements, regardless of the size of the batch.

Reviewer comment (Theory):
What is the gain of using Skipping instead of Early-Termination? If the function is monotonic decreasing wouldn't that solve the problem for Early-Termination?

Author response:
The key benefit of using Skipping over Early-Termination, even with a monotonic decreasing function, lies in how the transformer architecture processes information layer by layer. In Early-Termination, later tokens might only attend to intermediate hidden states, potentially missing out on crucial information processed in subsequent layers. Skipping ensures that the full computation performed by earlier tokens is utilized. This concept is illustrated in the 'Computational Efficiency' section of Figure 2 in our paper.

Reviewer comment (Experiments):
How is the speed-up and quality affected by the sequence-length hyperparameter? Why was is chosen to be the median? How would this work in a general chat-bot fashion where the generated sequence length can largely vary?

Author response:
The sequence-length hyperparameter, when other variables are held constant, influences the slope of the decaying function (see Figure 3), but it doesn't significantly impact the overall speedup. The choice of sequence length depends on the dataset, as our framework focuses on Fine Tuning and specialization (refer to the [general answer](https://openreview.net/forum?id=bcHty5VvkQ&noteId=PYRP8JGGUt)). Therefore, it is not ideally suited for general chat-bot applications. The sequence length and prompt size will ultimately depend on the dataset training set statistics. We selected the median prompt size during training to allow the model to adapt naturally to the typical prompt sizes found within the dataset's variance (Section 2.3, last paragraph; Section 3.1, 4th paragraph).

Reviewer comment (Evaluation):
Do the speedups include the prompt computation time?

Author response:
No, the speedups reported in our study do not include the prompt computation time, only the decoding. **We will clarify it more. We will ensure to make this clearer in the revised version of the manuscript.**

Reviewer comment (Presentation):
I am having a hard-time understanding Table-4. I have spent close to 15 minutes trying to understand the Table and reading text around. There are no lables, no descriptions what the numbers mean and how do you go about fixing speedup. Please improve the presentation.

Author response:
**Table 4:**

Table 4 compares SkipDecode with 2 baselines derived from the CALM paper at different levels of Speedup. Speedup reflects the various sizes of the network according to each approach. Multilayer: reflects the number of layers used monolithically, while CALM-DEC reflects the exit points at each token level. The numbers reflect the accuracy in each dataset.

**Sorry for the inconvenience, we will expand the section to include proper descriptions in the revised manuscript.**

Thank you for your feedback regarding Table 4. We apologize for any confusion caused by the presentation of the data.

Table 4 is designed to compare the performance of our proposed system, SkipDecode, against two baseline approaches adapted from the CALM framework. These baselines are a multi-layer exit network and CALM-DEC, a variation tailored for decoder-only models like ours.

The 'Speedup' column in the table represents different levels of computational efficiency achieved by each method. For the Multi-layer approach, 'Speedup' indicates the number of layers used in a monolithic fashion, essentially a measure of how early the model exits. In the case of CALM-DEC, it reflects the dynamic exit points at each token level, determined by the network's hidden state saturation.

The subsequent columns show the performance of each approach on two datasets, E2E and Reddit-TLDR, at various speedup levels. The performance is measured in terms of accuracy, with the numbers provided being percentage scores.

**Upon reviewing your feedback, we realize the need for a more comprehensive explanation of these concepts in the manuscript. **

We hope this clarification addresses your concerns. We will ensure the revised manuscript provides a clearer and more detailed explanation of Table 4.

Reviewer comment (Novelty):
The authors only discussed established directions such as distillation and quantization, but no recent works on compressing LLM nor efficient inference of LLM in Section 4. Relevant recent works include:  
Lin, Ji, et al. "AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration." arXiv preprint arXiv:2306.00978 (2023).  
Xiao, Guangxuan, et al. "Smoothquant: Accurate and efficient post-training quantization for large language models." International Conference on Machine Learning. PMLR, 2023.  
Frantar, Elias, et al. "Gptq: Accurate post-training quantization for generative pre-trained transformers." arXiv preprint arXiv:2210.17323 (2022).  
Sheng, Ying, et al. "FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU." (2023).  
Liu, Zichang, et al. "Deja vu: Contextual sparsity for efficient LLMs at inference time." International Conference on Machine Learning. PMLR, 2023.  
Zhang, Zhenyu, et al. "H_2 O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models." arXiv preprint arXiv:2306.14048 (2023).  
Liu, Zichang, et al. "Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time." arXiv preprint arXiv:2305.17118 (2023).

Author response:
The reviewer is right and we are thankful for your insightful feedback regarding the coverage of recent works on LLM compression and efficient inference in our manuscript. We acknowledge that incorporating recent advancements in this field would significantly strengthen our discussion and argument, especially concerning KV cache management during LLM inference.

In light of your suggestions, we realize the importance of emphasizing our method's ability to manage KV caches efficiently. Our approach not only facilitates KV caching during early exit scenarios but also effectively reduces the KV cache size, thereby diminishing the memory footprint. This point, although present in our initial submission, was not highlighted sufficiently. We appreciate that your comment has brought this to our attention, and even [1] suggested by the reviewer helped us argue in favor of efficiently dealing with KV caches in other reviewer response.

**We will incorporate a more detailed discussion of these recent works in the revised version of our paper. Specifically, we plan to:** 
1. Discuss the relevance of these recent methods in the context of our work, particularly focusing on how they complement and contrast with our approach;
2. Explicitly highlight how our method addresses some of the challenges posed by LLM compression and efficient inference;
3. Strengthen the section on KV cache management, underscoring its significance in efficient LLM inference, as your suggestion and the reference [1] have prompted us to consider.

We believe that these additions will improve the overall quality and comprehensiveness of our manuscript.

Thank you once again for your constructive feedback.

[1] *Specializing Smaller Language Models towards Multi-Step Reasoning, Fu et. al., 2023*

Reviewer comment (Experiments):
For a work that claims practical performance speedups of deep learning inference, it should be important to comprehensively evaluate the real-world increase in speed compared to the baselines. However, I have found that part of the experiments to be missing multiple crucial details, for example, the type of hardware used to run the experiments, the batch size during generation, and the metric that was used to obtain the numbers for true speedup (was it latency, throughput, or something else?). Also, apparently there is no real-world speedup comparison between SkipDecode and CALM-DEC.

Author response:
**Speedup measurement and Hardware details:**

**We will add hardware information and real time for the SkipDecode experiments.**

As hardware we use a single node with 16xV100 GPUs.

With respect to the baseline, real time might not be the best metric as our implementation is by no means optimized as it is the OPT Metaseq implementation. 

However, our aim here is to compare different configurations (2x, 3x, 4x, 5x) against the full network cost (1x) both in performance and speedup. 

**We will include detailed information about the hardware used for running our experiments. Additionally, we will provide real-time measurements for the SkipDecode experiments to offer a more tangible understanding of the speedups achieved.**

While we understand the importance of real-time metrics, we would like to clarify that our baseline CALM-DEC implementation is not fully optimized. SkipDecode and CALM-DEC are implemented on different frameworks. Therefore, comparing real-time metrics directly might not accurately reflect the efficiency of the underlying method. 

We followed the metric originally proposed by CALM which is the number of layers used, which is a good proxy for time and resources needed. Same as CALM, the True Speedup in table 3 is measured as the Total Number of Layers in the network by the Average Number of Layers used per token. Difference with CALM is that we can do batching so this measure includes a batched setting (i.e., batches of arbitrary size), while CALM restricts the analysis to a scenario with a batch size equal to 1.

However, as our aim here is to compare different configurations (2x, 3x, 4x, 5x) against the full network cost (1x) both in performance and speedup. We can provide real-time metrics for different configurations of SkipDecode.

We believe these additions and clarifications will comprehensively address the concerns raised and provide a clearer picture of the practical performance speedups achievable by SkipDecode. We appreciate the opportunity to enhance our work based on your feedback.

Reviewer comment (Writing):
There are quite a few typos in the submission: for example, "figure 2" -> "Figure 2" at the bottom of page 3, "e2e dataset" -> "E2E dataset" and "figure 3" -> "Figure 3" on page 6, "SkipDecodemodels" -> "SkipDecode models" on page 8.

Author response:
* Thank you very much for the typos we will fix it in the reviewed version

Reviewer comment (Experiments):
If I understood correctly, each of the experiments you ran involved finetuning the model on a target dataset. However, in practice, model providers might often serve a single model for many different applications: for example, this means that the average prompt/response lengths can vary dramatically. Is it possible to extend SkipDecode to such a scenario?

Author response:
**Role of Fine-Tuning in SkipDecode:**

Please refer to the general answer above under ["Role of Fine-Tuning in SkipDecode"](https://openreview.net/forum?id=bcHty5VvkQ&noteId=PYRP8JGGUt) title as this point was raised by another reviewer.

**Adapting parameters to different datasets:**

We indeed fine-tune for each dataset (see ["Role of Fine-Tuning in SkipDecode"](https://openreview.net/forum?id=bcHty5VvkQ&noteId=PYRP8JGGUt)). To adapt for each dataset, we use the statistics of the training set. To handle varying prompt lengths we train the network using the median prompt size (Section 3.1) which has worked well in our experiments.

Regarding the prompt length again we set the max_len based on our training set, and if a datapoint in test is longer than that it will be assigned the min_layer for each token beyond the max_len parameter (Section 3.1).