$\text{D}_{2}\text{O}$: Dynamic Discriminative Operations for Efficient Long-Context Inference of Large Language Models

Authors: Zhongwei Wan, Xinjian Wu, Yu Zhang, Yi Xin, Chaofan Tao, Zhihong Zhu, Xin Wang, Siqi Luo, Jing Xiong, Longyue Wang, Mi Zhang
Primary area: foundation or frontier models, including LLMs
Keywords: Efficient inference of LLMs, Long context

Abstract:
Efficient generative inference in Large Language Models (LLMs) is impeded by the growing memory demands of Key-Value (KV) cache, especially for longer sequences. Traditional KV Cache eviction strategies, which discard less critical KV-pairs based on attention scores, often degrade generation quality, leading to issues such as context loss or hallucinations. To address this, we introduce **D**ynamic **D**iscriminative **O**perations ($\mathbf{D_2 O}$), a novel method that optimizes KV cache size dynamically and discriminatively at two levels without fine-tuning, while preserving essential context. At **layer-level**, by observing the varying densities of attention weights between shallow and deep layers, we dynamically determine which layers should avoid excessive eviction via our proposed ***dynamic allocation strategy*** to minimize information loss. At **token-level**, for the eviction strategy in each layer, $\mathbf{D_2 O}$ innovatively incorporates a ***compensation mechanism*** that maintains a similarity threshold to re-discriminate the importance of currently discarded tokens, determining whether they should be recalled and merged with similar tokens. Extensive experiments on various benchmarks and LLM architectures have shown that $\mathbf{D_2 O}$ not only achieves significant memory savings and enhances inference throughput by more than 3$\times$ but also maintains high-quality long-text generation.

Decision: accept (poster)
Reviewer scores: 6, 5, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729493578427

thread_type:
official_review

<thread_object>
creator: reviewer_asvx

<content>
summary:
This paper introduces Dynamic Discriminative Operations (D₂O), a method to optimize KV cache management in large language models. At the layer level, D₂O allocates cache sizes dynamically based on attention variance. At the token level, it merges similar tokens and prunes less important ones, preserving essential contextual information. This approach reduces memory usage while improving efficiency and accuracy in long-text generation and inference tasks, outperforming existing cache compression techniques.

strengths:
1. The proposed method outperforms traditional eviction-based approaches, significantly reducing performance loss in sparse long-text inference.
2. Extensive experimental analysis on hyperparameter selection, various models, and datasets strongly validates the method's effectiveness.
3. The method demonstrates excellent generalizability, as it can be applied to a wide range of existed pretrained models without requiring additional post-training or fine-tuning.

weaknesses:
The innovation and novelty of the proposed method are somewhat limited compared to existing sparse attention pre-trained models or eviction-based inference approaches. For instance, the concept of token merging has been discussed in several prior works, such as https://openreview.net/pdf?id=LCTmppB165 and https://arxiv.org/html/2407.08454v1. Additionally, the idea of dynamically allocating different computational loads to different layers has also been explored in other works, such as https://arxiv.org/pdf/2404.02258.

questions:
1. For the Dynamic Allocation Strategy, could it result in a single layer's computational load exceeding that of processing all tokens in a layer? How would this be handled in such cases?
2. In the Dynamic Allocation Strategy, how do you ensure that the computational load is uniform across samples within a batch? Could there be scenarios where some samples in the same batch have significantly higher or lower computational loads, thus reducing inference efficiency? Can this strategy be integrated with existing inference acceleration methods like vLLM?
3. In the Dynamic Allocation Strategy, is there any physical rationale behind the decision to allocate fewer KV caches to layers with high attention score variance?
4. Could the method be compared to a smaller full-size model with similar inference time costs?
5. In Table 9 and Figure 5, D₂O performs better than the full model in some tasks. Could you explain and analyze the reasons for this outcome?
6. In Section A.4, could you provide results for the full model in the multi-turn dialogue examples for comparison?
What are the inference costs in Table 3 for identical batch sizes?
7. Missing citation: https://arxiv.org/pdf/2408.03675 . This paper also focuses on optimizing token eviction methods which introduces a random eviction strategy that retains more token information.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730523782547

thread_type:
official_review

<thread_object>
creator: reviewer_fb7s

<content>
summary:
This paper explores methods for saving the KV cache during inference of large language models. Specifically, the authors investigate ignoring less important tokens in the context during decoding and find that, to prevent excessive performance degradation, the number of tokens discarded should vary across different layers. Additionally, they explore a token compensation mechanism to mitigate the information loss caused by discarding tokens. Experiments demonstrate that their approach accelerates the model while maintaining performance.

strengths:
* The motivation is clear and saving the KV cache of LLMs during inference is significant.
* This paper deeply investigates the difference in token eviction across layers and explores a reasonable token compensation mechanism.
* The experiments are relatively extensive.

weaknesses:
* Some judgments are inaccurate. As noted in lines 181-182: "Employing a uniform eviction strategy such as the ones proposed in H2O...". However, to my knowledge, H2O also adopts a layer-wise token eviction strategy, as evidenced in its open-source code [1].
* The baselines are limited. Exploring the issue of token eviction ratios across different layers is highly meaningful. However, to my knowledge, some relatively classic works in this area, such as PyramidKV [2], were not considered by the authors.
* Some experimental results are quite unusual. In detail,
  * Logically, discarding some tokens would lead to information loss. However, some results in Figure 5 of the paper indicate that, after discarding certain tokens, the performance improves significantly compared to the original model (for instance, on the TruthfulQA task). Is this reasonable?
  * The evaluation results on Longbench differ significantly from those in some previous studies. For example, The results in line 423 of the paper differ by more than **10** points from those in Table 1 of PyramidKV [2] and SnapKV [3]. I believe the authors need to carefully review and revise their findings to provide more reliable results.

[1] https://github.com/FMInference/H2O/blob/main/h2o_hf/utils_real_drop/modify_llama.py#L357

[2] Zhang, Yichi, et al. "PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling." arXiv preprint arXiv:2406.02069 (2024).

[3] Li, Yuhong, et al. "Snapkv: Llm knows what you are looking for before generation." arXiv preprint arXiv:2404.14469 (2024).

questions:
* In formula (5), $S_\ell$ is the cache size for each layer. However, when calculating $\alpha_\ell$, the number of model layers $L$ is multiplied again. Is this a typographical error?
* In formula (6), I cannot understand the calculation of AttnScore during token generation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730622476970

thread_type:
official_review

<thread_object>
creator: reviewer_vhpa

<content>
summary:
This paper focuses on token eviction for large language models and introduces a dynamic allocation strategy at the layer level and a compensation mechanism at the token level.

strengths:
1. The proposed token eviction mechanisms at both the layer level and token level are well-reasoned and contribute to performance improvements in the models.
2. The effectiveness of the proposed methods is validated across multiple model backbones and long-context benchmarks.

weaknesses:
1. The experimental comparisons are not comprehensive enough. The authors should compare their method with recent token eviction approaches, such as SnapKV and SirLLM. Additionally, comparisons with token compression techniques like LoCoco would be beneficial.
2. It is suggested that the authors evaluate the performance of their method on benchmarks with longer contexts, such as InfiniteBench and RULER, which consist of instances up to 128K tokens.

questions:
Please refer to Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730645717263

thread_type:
official_review

<thread_object>
creator: reviewer_ldwy

<content>
summary:
This paper proposes D2O, a method for efficient LLM inference through dynamic KV cache optimization. The key contributions include: (1) layer-wise cache allocation based on attention density patterns, and (2) a token-level strategy combining eviction with dynamic merging using EMA threshold. The method achieves 3x throughput improvement while maintaining generation quality across various tasks.

strengths:
-The paper brings up a major snag in putting LLMs into action, especially the memory requirements of the KV cache during making inferences. This is a crucial issue for real-world applications.

-The authors have conducted extensive experiments across multiple LLM architectures  and tasks, providing a broad evaluation of the method's effectiveness.

 -D2O shows notable improvements in throughput and memory efficiency, achieving up to 3x higher throughput while preserving generation quality, which are important metrics for the deployment of LLMs.

weaknesses:
-The core ideas (layer-wise allocation and token merging) are largely combinations of existing techniques. The EMA threshold mechanism is straightforward and lacks theoretical justification.The dynamic allocation strategy is similar to existing attention-based pruning methods

-There is insufficient analysis. There is no ablation study on the impact of different attention patterns across layers. There is also limited discussion on why the 3:1 ratio for N:M works best. 

-Regarding experimental concerns, the comparison with baselines is not entirely fair as some methods (like H2O) are designed for different scenarios. There is no discussion on the trade-off between computational overhead and memory savings.

questions:
-How does the computational overhead of D2O scale with increasing sequence lengths and model sizes?

-What are the trade-offs between the compression ratio and the quality of the generated output?

-How does D2O handle extreme cases, such as sequences longer than 100k tokens?

-Why is the exponential function chosen for the allocation strategy? Have other functions been considered?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730729926679

thread_type:
official_review

<thread_object>
creator: reviewer_jaqq

<content>
summary:
This paper presents Dynamic Discriminative Operations D2O, an efficient generation method in large language models. D2O uses layer-level to dynamically adjust cache usage and token-level strategies to merge discarded tokens. D2O achieves significant memory savings and improves throughput without sacrificing generation quality, as demonstrated in various benchmarks and LLM architectures.

strengths:
- D2O takes into account the phenomenon of varying attention score densities across different layers in LLMs, dynamically allocating KV cache budgets for each layer. Additionally, for important discarded tokens, D2O uses a weighted merging approach to retain their information, mitigating information loss during token eviction.
- Compared to various KV cache compression methods, D2O demonstrates consistent superiority across a wide range of benchmarks and different models. With a comparable improvement in throughput, D2O achieves significant performance gains in task execution.
- This paper provides extensive ablation analysis, validating the effectiveness of each component within D2O.

weaknesses:
- The Introduction states that "existing methods equally treat all the layers and indiscriminately evict KV pairs at each layer." However, some prior work has already considered allocating different KV cache budgets across layers, such as PyramidInfer[1] and PyramidKV[2]. The allocation method in D2O should be compared with those in previous works.
- The use of mathematical symbols is confusing. For example, the F representing variance in line 207 is a scalar and should not be in bold, as well as u_ij in line 261.
- The experiment setup is not clear enough. For example, in Section 5.3, the maximum length setting for the LongBench test is mentioned.

[1] PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference (Yang et, al.)

[2] PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling (Cai et, al.)

questions:
- As the generation length increases, if a certain KV Cache budget is maintained, the effect of token merging is expected to weaken. Are there any experiments on this?
- The performance testing in the paper separates the benchmarks from the throughput tests. In the LongBench, the prompts are usually long while the outputs are short, so the methods presented in this paper may not save much time on such tasks. Are there more suitable benchmarks where the generated text is longer, allowing for simultaneous evaluation of generation speed and performance?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732155749710

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response Summary and Paper Revision.

comment:
Dear Reviewers,

We sincerely thank you for their valuable feedback, which has significantly improved the quality and rigor of our paper. Based on your 
suggestions, we have revised and updated the paper, highlighting changes and additions in **blue** within the revised PDF. Below, we summarize the key revisions made:

- **1. Addition of Baselines:** In response to comments from reviewers JAQq, VHpa, and FB7s, we have incorporated comparisons with the latest baselines, including SnapKV [1], PyramidKV [2], PyramidInfer [3], SirLLM [4], and LoCoco [5]. These results are presented in the updated Appendix A.10 (highlighted in blue) and **Table 18**.

- **2. Expanded Related Work:** Following suggestions from reviewers aSVX, FB7s, and JAQq, we have added descriptions of additional relevant works, including NaCL [6], SirLLM [4], PyramidKV [2], and PyramidInfer [3], in the Related Work section, now highlighted in blue in the main text.

- **3. Generative Long-Context Evaluation:** Per reviewer JAQq's feedback, we conducted experiments to evaluate both generation speed and performance on generative long-context tasks. These results are included in **Appendix A5 (lines 928–937)** and **Table 15**.

- **4. Inclusion of InfiniteBench and Ruler Results:** As recommended by reviewers LdwY and VHpa, we added evaluations using InfiniteBench [7] and Ruler [8]. The results are detailed in **Appendix A.11** and presented in **Tables 19 and 20**.

- **5. Computational Overhead Analysis:** Addressing reviewer LdwY's suggestion, we evaluated the growth of computational overhead with increasing model sizes and context lengths. These results are included in **Appendix A5 (lines 921–927)** and **Table 14**.

- **6. Needle-in-a-Haystack Results:** Based on reviewer aSVX’s feedback, we expanded the analysis of needle-in-a-haystack results with corresponding inference times, now detailed in **Appendix A5 (lines 934–937)** and **Table 16**.

- **7. Clarity and Typos:** To address comments from reviewers FB7s and JAQq, we corrected and clarified certain expressions and typos. For example: **Line 181-182**: "uniform eviction strategy" was revised to "uniform layer-wise size."  **Line 207 and 261**: Removed bold formatting for scalar variables. These corrections are highlighted in blue in the revised text.

Once again, we deeply appreciate the time and effort you have dedicated to improving this paper. Your insights have made the paper stronger and more comprehensive. Wishing you a Merry Christmas and a Happy New Year!

Bests,

Your friends, 

Team of 5363

**References:**
[1] Li, Yuhong et al. “SnapKV: LLM Knows What You are Looking for Before Generation.” ArXiv abs/2404.14469 (2024).

[2] Cai, Zefan et al. “PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling.” ArXiv abs/2406.02069 (2024).

[3] Yang, Dongjie et al. “PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference.” Annual Meeting of the Association for Computational Linguistics (2024).

[4] Yao, Yao et al. “SirLLM: Streaming Infinite Retentive LLM.” ArXiv abs/2405.12528 (2024).

[5] Cai, Ruisi et al. “LoCoCo: Dropping In Convolutions for Long Context Compression.” ArXiv abs/2406.05317 (2024).

[6] Chen, Yilong et al. “NACL: A General and Effective KV Cache Eviction Framework for LLM at Inference Time.” ArXiv abs/2408.03675 (2024).

[7] Zhang, Xinrong et al. “∞Bench: Extending Long Context Evaluation Beyond 100K Tokens.” ArXiv abs/2402.13718 (2024).

[8] Hsieh, Cheng-Ping et al. “RULER: What's the Real Context Size of Your Long-Context Language Models?” ArXiv abs/2404.06654 (2024).
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces a method for efficient LLM inference, D2O, through dynamic KV cache optimization. The key contributions are (1) layer-wise cache allocation based on attention density patterns, and (2) a token-level strategy combining eviction with dynamic merging using the EMA threshold. 

Reviewers respect the conributions such as task importance, effective apporach, and promising results with extensive experiments while a reviewer raised limited novelty issue.

AC carefully read the paper, reviewer comments, and author response. It seems that the authors provided clear explanation and additional experiments on the issues raised by LdwY. 

Considering the contributions and no additional feedback, AC recommends accepting this paper.