LazyLLM: DYNAMIC TOKEN PRUNING FOR EFFICIENT LONG CONTEXT LLM INFERENCE

Regarding formatting issue

Thanks for pointing it out. We will fix the issue in the camera ready version.

Regarding KV Cache

Yes your understanding is correct – once a token is computed, its KV cache is fixed and we’ll not update until the end of generation. We found this is simple and effective. Empirically we found updating the KV along the generation is very slow (coule offset the benefits of token dropping) and provides marginal performance improvement. Our intuition is that the pruned token generally has relatively smaller attention weights, thus doesn’t have much impact on the following tokens.

Regarding implementation details

We'll detailed pseudocode in the appendix and will release our complete implementation upon acceptance.

Regarding performance improvement over baseline

Your observation about LazyLLM sometimes outperforming the baseline is spot on and raises an interesting point. As demonstrated in "Lost in the Middle: How Language Models Use Long Contexts," irrelevant information in long contexts can actually harm performance. By pruning less relevant tokens, LazyLLM may sometimes improve model focus on pertinent information, leading to better results. This aligns with recent findings about context processing in large language models.

Dear Authors,

Thank you very much for your response, which really helps a lot. After reading the rebuttal, I feel that this work actually suggests an interesting point: calculating each token once is enough for some tasks.

However, I still have one question:

• If possible, could the author explain why such a method could work?
  • Layer i has tokens [T1, T4, T5, T7] and Layer i+1 has tokens [T4, T5].
  • Current generation step, Layer i has tokens [T1, T2, T4, T5], and Layer i+1 has tokens [ T4, T5].
  • Apparently, the last step [T4, T5] and this step [ T4, T5] actually should have different embedding representations, while this work proves that it could work well directly using the last step embedding [T4, T5].
  • Is it because the text token is highly redundant? Or are there any other potential reasons?

I will consider increasing the score or the confidence after the author provides the potential explanation.

Regarding GPU memory usage

We appreciate the opportunity to clarify this important point. As mentioned in line 289, LazyLLM's memory management is quite efficient: tokens are either in KV cache or Aux Cache, never both. Given that hidden state size is smaller than KV cache entries, the total memory footprint is actually smaller than the baseline's KV cache. Furthermore, our method reduces attention map sizes, leading to additional memory footprint savings.

Regarding Alignment of GPU memory costs in experiments

As mentioned in Sec. 5.4, except for the Prompt Compression, the peak memory usage of LazyLLM remains equivalent to that of the baselines.

Regarding comparisons with advanced baselines

While we acknowledge the importance of methods like H2O and SnapKV, they focus primarily on decoding optimization rather than TTFT reduction (as noted in line 133). Our approach is complementary to these KV cache management techniques and could potentially be combined with them for even better performance. We’ll cite these papers and add more detailed discussions in the related work.

We have updated the paper to provide more detailed implementation insights:

1. Added comprehensive pseudocode in the appendix
2. Included a step-by-step example that walks through the token pruning and revival process in the appendix
3. Provided detailed illustrations of Aux and KV cache management across generation steps in the appendix

We hope these additions help clarify both the technical details of our method and its memory efficiency. The example walkthrough particularly demonstrates how LazyLLM handles token pruning and revival across different transformer layers during inference while maintaining a smaller memory footprint than the baseline approach.

Please let us know if it resolves your concern, thanks!

End-to-End Latency and Throughput Comparison

We provide the comparison of the end-to-end latency and throughput comparison with other baselines as below:

The experimental results reveal that LazyLLM achieves near-baseline accuracy (marginal decrease of -0.88%) while significantly improving computational efficiency. Notably, LazyLLM reduces end-to-end latency by 15.5% compared to the baseline and increases throughput by 17.5%. This performance advantage is particularly meaningful when compared to other baselines, which show considerable accuracy degradation despite modest efficiency gains.

We'll add these results and analysis to the revised version.

GPU memory usage comparisons

As we have mentioned in section 5.4 (line 413), all attention-based token pruning methods share the same peak GPU memory usage. This consistency stems from the architectural constraint where peak memory is fundamentally determined by the maximum attention map dimensions. Given that all token pruning approaches, including LazyLLM, process the complete token set in the initial transformer layers, they generate equivalent maximum attention map sizes, resulting in uniform peak memory requirements.

Our analysis extends to prompt compression techniques, which maintain the same peak memory usage as the vanilla baseline. This occurs because these methods necessitate a complete context processing phase to generate the compressed prompt.

The notable exception in our evaluation is the Random Token Drop method, which demonstrates reduced GPU memory utilization through random context pruning prior to LLM processing. With a 10% pruning rate, this approach achieves a modest 5% reduction in GPU memory consumption compared to the vanilla baseline. However, this memory advantage comes at a substantial cost to model performance (as detailed in Table 1), resulting in significantly degraded accuracy.

Finally, we would like to emphasize that LazyLLM's primary innovation lies in optimizing Time-to-First-Token (TTFT) and overall generation efficiency during LLM inference, while not increasing the GPU memory usage. This strategic focus ensures that performance improvements do not come at the cost of increased memory overhead, making LazyLLM suitable for practical applications where both efficiency and resource constraints are critical considerations.

We'll extend the GPU memory discussion in the paper accordingly.

Regarding hyperparameter settings

We appreciate this feedback and have expanded our discussion in the appendix. LazyLLM employs five hyperparameters that control its progressive token pruning strategy:

1. start_prune_layer: The first layer where pruning begins
2. end_prune_layer: The final layer where pruning occurs
3. start_keep_ratio: Initial token retention ratio
4. end_keep_ratio: Final token retention ratio
5. num_prune: Number of pruning layers

To maintain efficiency and controllability, we perform token pruning at num_prune layers that are evenly distributed between start_prune_layer and end_prune_layer. At each pruning layer i, we retain tokens based on a linearly interpolated keep ratio:

keep_ratio(i) = start_keep_ratio - (i - start_prune_layer) / (end_prune_layer - start_prune_layer) * (start_keep_ratio - end_keep_ratio)

This progressive pruning strategy allows for gradual token reduction while maintaining model performance. We determined effective hyperparameter values through a coarse grid search on a small held-out validation set, finding that the method is relatively robust to hyperparameter choices within reasonable ranges. While more sophisticated optimization techniques might yield marginal improvements, our current approach achieves strong and consistent performance across diverse tasks with minimal tuning required.

Regarding token revival strategy

Thank you for this insightful observation. Our decision to append revived tokens rather than maintaining strict sequential order was based on the following findings:

• Updating KV cache entries before and after revived tokens would necessitate repeated computation of the same tokens, effectively negating the performance benefits of our token pruning strategy and introducing additional latency to the generation process.
• While insertion of tokens mid-sequence is theoretically possible, current tensor manipulation implementations in deep learning frameworks (including PyTorch) make this operation computationally expensive, introducing significant overhead.
• Our analysis shows that tokens selected for pruning consistently exhibit low attention weights across layers, suggesting their precise sequential position has minimal impact on the model's understanding of the context. Importantly, we preserve the original positional embeddings of these tokens to maintain relative positional information, ensuring the model retains access to the original sequence structure.
• Extensive empirical validation demonstrates that this approach incurs minimal performance impact - we observe less than 1% relative performance difference compared to strict sequential ordering, while achieving significant computational efficiency gains. Based on these findings, our approach represents an optimal balance between maintaining model performance and achieving computational efficiency. We will add a discussion to the paper to clarify this.

Regarding attention-based pruning and bias

We share the concern about potential bias amplification. Our use of top-percentile pruning (rather than top-k) was specifically chosen to maintain a majority of tokens and minimize bias introduction. While we acknowledge that any pruning strategy may affect model behavior, our approach seeks to balance efficiency gains with maintaining model fidelity. We're actively investigating more sophisticated pruning strategies for future work.

Regarding cross-layer token relevance

Thank you for the thoughtful comment. We acknowledge that tokens may have varying importance across different layers. While allowing each layer to independently choose tokens would be ideal, this would require redundant computation when a token needs to be revived. Our current approach represents a practical compromise between computational efficiency and maintaining important cross-layer relationships. We’ll add more detailed discussion and reference to the revised version.

Regarding Possible memory overhead from the Aux Cache during decoding

First of all, as mentioned line 289, we would like to address that in LazyLLM, a token is either in KV cache or Aux Cache, thus, the total size of KV cache and Aux Cache is small than the KV cache of the baseline (consider the hidden state size of token is smaller than its KV). Yes we agree that further optimization (as mentioned) can be applied to the AuxCache (and also our KV cache), which we’ll mark as the soon future work!

Regarding the generality of the problem and TTFT optimization

We appreciate your insight about different server configurations. While some large organizations can indeed separate prefilling and decoding across servers, we believe our work addresses an important practical need: Many organizations and individual researchers work with more constrained infrastructure where such separation isn't feasible. More importantly, our method is particularly valuable for edge devices and mobile applications where resource optimization is crucial and such separation is impractical. Our contribution is complementary to parallel computation approaches and provides benefits even in distributed settings.

Regarding Figure 4

Thank you for bringing this up. We have added a detailed example of LazyLLM inference in the appendix to improve clarity. (to be upload soon). We will make sure to improve the figure in the final version and would greatly appreciate any specific feedback about the aspects of the figure you found more confusing so we can make targeted improvements. Regarding comparison with parallel computation methods: We acknowledge the importance of parallel computation techniques and want to thank you for providing the reference. Please note that the mentioned paper [2] first appeared online on 25 Sep 2024. Per ICLR Guide: “if a paper was published (i.e., at a peer-reviewed venue) on or after July 1, 2024, authors are not required to compare their own work to that paper”. That said and more importantly, our method focuses on input context optimization, which is orthogonal and complementary to parallel computing approaches. Both techniques could be combined for potentially greater benefits. We will add this reference to the paper upon acceptance.

Regarding AuxCache vs LESS

While both approaches involve caching, they serve fundamentally different purposes: LESS optimizes KV cache storage through low-rank optimization, whereas AuxCache stores hidden states of pruned tokens for efficient revival when needed. We clarify this in the paper.

Regarding performance variation across tasks

The variation in speedup between code completion (1.01×) and single-document QA (1.34×) reflects the inherent nature of code versus natural language. In code completion tasks, nearly every token carries semantic or syntactic significance, with minimal redundancy in the context. This fundamental characteristic of code means we ultimately need to process most tokens to maintain code integrity. Nevertheless, our method still achieves significant TTFT improvements of 1.94× and 3.47× for code completion tasks, demonstrating its effectiveness even in scenarios with dense, meaningful contexts.

We have updated the paper to provide more detailed implementation insights:

1. Added comprehensive pseudocode in the appendix
2. Included a step-by-step example that walks through the token pruning and revival process in the appendix
3. Provided detailed illustrations of KV cache management across generation steps in the appendix

We hope these additions help clarify the technical details of our method. The example walkthrough particularly demonstrates how LazyLLM handles token pruning and revival across different transformer layers during inference.

Please let us know if you would like us to elaborate on any specific aspect of the implementation or if you have additional questions. We're committed to ensuring the method is clear and reproducible.