SkipDecode: Autoregressive Skip Decoding with Batching and Caching for Efficient LLM Inference

We sincerely appreciate the insightful feedback and thoughtful inquiries provided by the reviewer. In this response, we aim to comprehensively address the concerns and questions raised, hoping to clarify and enhance the understanding of our research.

OPT choice and decoder architecture:

Please refer to the general answer above under "OPT choice and decoder architecture" here as this point was raised by another reviewer.

Compared only with one baseline

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.

Q.1. How are the min_exit_layer and max_exit_layer chosen?

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.

Q.2. What is the degradation of quality using batch-wise exit-point function instead of a per example-wise?

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.

Q.3. 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?

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.

Q.4. 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?

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). 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).

Q.5. Do the speedups include the prompt computation time?

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.

KV Cache in Early Exit:

Regarding the specific question about computing the cache for early exit methods. In CALM, and similar early exit methods, generating the K,V cache for tokens exiting at different layers presents some challenges. This is primarily due to the architectural deviations these methods introduce in the standard Transformer model. For example, if a token exits beyond the layer reached by previous tokens, CALM attends to the last computed layer for those tokens. While it is technically feasible to generate the cache in such scenarios, this process necessitates additional logic and careful handling to ensure a closer alignment with the original Transformer architecture.

In contrast, SkipDecode maintains the standard layer-by-layer attention mechanism of the Transformer. This adherence to the conventional architecture means that K,V caching can be utilized straightforwardly, without necessitating low-level implementation adjustments or architectural decisions.

The core advantage here is that SkipDecode retains the efficiency of the Transformer's attention mechanism, avoiding the computational overhead of on-the-fly cache recomputation that early exit methods might incur if they are to adhere to the Transformers original architecture (Figure 2).

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.

Is the loss per token a valid proxy for "hardness"?

Our intuition is that, the loss per token, which measures the discrepancy between the predicted and actual outputs, can be indicative of the difficulty the model encounters in accurately predicting the correct answer at each token position. This interpretation aligns with the principle of loss functions where higher loss values typically reflect greater challenges in model prediction (and therefore require greater adjustments).

The context within which the model operates plays a significant role. For instance, in sequence-based tasks, predictions made early in a sequence may be more challenging due to limited preceding context, as opposed to later predictions where the model has more contextual information. This observation is consistent with the findings in the CALM paper (specifically Figure 2.a), where the impact of perturbations at different sequence positions was analyzed. The model found it more challenging to recover from perturbations occurring early in the sequence than those towards the end.

In our work, we use loss per token as an indicative measure, but we also acknowledge the complexity of interpreting this metric as a standalone indicator of 'hardness'. Future research could further refine this approach, perhaps by integrating additional metrics or considering the contextual factors that influence model predictions.

We sincerely appreciate the insightful feedback and thoughtful inquiries provided by the reviewer. In this response, we aim to comprehensively address the concerns and questions raised, hoping to clarify and enhance the understanding of our research.

Position-wise schedule:

You raise a valid point about the potential benefits of a context-based approach over a fixed schedule.

In an ideal scenario, as implemented in methods like CALM, decisions for early exits would indeed be made on a token-by-token basis. However, while theoretically advantageous, this approach faces practical limitations in batched scenarios. The computational theoretical gains presented in methods like CALM are only realized when the batch size is limited to one.

Our method seeks to address this limitation by implementing position-based exits. While a context-driven approach at the batch level is conceptually appealing, it introduces significant complications. Each token in a batch might require a different exit point due to its unique context. For example, a token in the first element of a batch might need an early exit, whereas a token in the last element might benefit from a later exit. This variability can negate computational gains and introduce uncertainty in terms of computational cost.

To ensure efficiency and predictability, our approach focuses on having the entire batch exit at the same position. This exit strategy across the batch ensures that the computational gains are both effective and predictable. In this context, fine-tuning becomes a crucial aspect of our method as it allows us to optimize the exit points for different positions, balancing the trade-off between computational efficiency and the contextual nuances of each token.

KV cache analysis:

Generative inference in LLMs often employs a KV Cache mechanism to enhance generation speed. In the causal decoder used by these models, attention for a token is calculated based on preceding tokens, leading to repetitive recalculations for each token generation. KV cache mitigates this by storing and reusing previously computed Key/Value vectors, thereby avoiding redundant computations at the cost of increased memory usage. Without KV caching, for each new token all KV representations for previous tokens, should be recomputed every time.

[1] conducted an experiment with GPT-2 on a Tesla T4 GPU (details available in [1]) to quantify the impact of KV caching. The results were compelling:

With KV Caching: 11.885 ± 0.272 seconds (for the generation of 1000 new tokens) Without KV Caching: 56.197 ± 1.855 seconds

This demonstrates the significant efficiency gains offered by KV caching. Our simple method, via skipping, allows the use of KV caching in early exit scenarios and, even more, as it progressively reduces layer usage, it decreases both KV requirements and overall memory footprint.

The memory consumption of KV cache increases rapidly as the model size and generation length increase, drastically increasing the pressure of on-device memory [2], in some scenarios even surpassing the model size [4]. When memory usage exceeds GPU capacity, the generative inference of LLMs typically resort to offloading and the KV cache loading ends up dominating the inference cost. According to [3] inference with multihead attention incurs significant memory capacity and bandwidth costs to store and load the KV cache, and these costs can dominate the rest of the inference at large batches or long context lengths.

Our approach aims to improve both sides of the coin, utilizing KV cache to speed up the process in the context of early exit while minimizing its memory footprint. As [2] and [3] state, dealing with KV caching is critical for advancing efficient generative inference in LLMs.

[1]: "Transformers KV Caching Explained, Lages, 2022" ( https://medium.com/@joaolages/kv-caching-explained-276520203249)

[2]: Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs, Ge et. al., 2023

[3]: Efficiently Scaling Transformer Inference, Pope et. al., 2023.

[4]: Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time, Liu et. al, 2023

Role of Fine-Tuning in SkipDecode:

Please refer to the general answer above under "Role of Fine-Tuning in SkipDecode" title as this point was raised by another reviewer.

We sincerely appreciate the insightful feedback and thoughtful inquiries provided by the reviewer. In this response, we aim to comprehensively address the concerns and questions raised, hoping to clarify and enhance the understanding of our research.

Role of Fine-Tuning in SkipDecode:

Please refer to the general answer above under "Role of Fine-Tuning in SkipDecode" title as this point was raised by another reviewer.

Missing angle on efficient LLM inference:

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

E2E Dataset:

Our approach is particularly effective in scenarios where generation tasks are extensive and the scope is well-defined. This makes it well-suited for applications such as summarization and translation.

The E2E dataset is also an ideal fit. It focuses on transforming structured data into coherent verbalizations, usually involving concise prompts. For instance, a typical prompt in the E2E dataset might be: {"prompt": "name[The Dumpling Tree], eatType[restaurant], food[Italian], priceRange[high]"}, with the corresponding response: "The Dumpling Tree is an Italian restaurant with high prices." This dataset requires the model to convert structured, attribute-based information into a fluid narrative form – a task that aligns seamlessly with the strengths of our method.

The reason our method shows significant performance improvements with the E2E dataset can be attributed to several factors:

1. The E2E dataset's well-defined nature and focused scope enable our model to specialize effectively via fine-tuning.
2. Our method demonstrates its strength in maintaining quality up to the point of hidden state saturation. In the case of the E2E dataset, this threshold lies within the manageable size of the model. This indicates that our method can effectively handle the dataset's tasks without requiring a very large network. The good numbers in all speedups indicate that the method is able to avoid degradation up to the hidden state saturation point.

In summary, the E2E nature of its tasks play into the strengths of our method, allowing for more efficient language generation. This is reflected in the performance metrics we observed.

We sincerely appreciate the insightful feedback and thoughtful inquiries provided by the reviewer. In this response, we aim to comprehensively address the concerns and questions raised, hoping to clarify and enhance the understanding of our research.

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.

More Diverse Datasets:

Please refer to the general answer above under "More Diverse Datasets" title as this point was raised by another reviewer.

OPT choice and decoder architecture:

Please refer to the general answer above under "OPT choice and decoder architecture" title as this point was raised by another reviewer.

Role of Fine-Tuning in SkipDecode:

Please refer to the general answer above under "Role of Fine-Tuning in SkipDecode" 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"). 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).

• Thank you very much for the typos we will fix it in the reviewed version