Keep the Cost Down: A Review on Methods to Optimize LLM’s KV-Cache Consumption

W1: The papers is in some places too condensed

Thank you for your careful review. Due to page limitations, we were unable to provide more detailed explanations in some sections. However, we plan to offer more extensive discussions and elaborations in the appendix in our future work. This will ensure that readers can access detailed information without compromising the overall structure and conciseness of the main paper.

Q1: Some acronyms are not explained: GQA (page 3), RAG (page 4), please consider explaining

Here GQA refers to Grouped-query Attention, and RAG refers to Retrieval-augmented Generation.

Q2&Q3: Typos and Presentation

Thank you for your suggestion. We will revise and optimize the expressions in the subsequent versions and conduct a more thorough proofreading to correct any typos.

Q4: Appendix A: "Mixture of Experts model" - please explain what is that

In the context of transformer models, a MoE(Mixture of Experts) consists of two main elements:

• Sparse MoE layers are used instead of dense feed-forward network (FFN) layers. MoE layers have a certain number of “experts”, where each expert is a neural network. In practice, the experts are FFNs.
• A gate network or router, that determines which tokens are sent to which expert. Please refers to paper (https://arxiv.org/pdf/2101.03961) for more detailed explanations.

Q5: Appendix C 'Other methods' deserves to be a section in the body of the paper, and deserves more elaboration. If you get more pages for the paper, consider moving this section

Thank you for your invalueable suggestion. We will try to integrate the content from Appendix C into the main body of the paper for a more comprehensive discussion.

Q6：The following papers in the References section have no venue/publication data: Arora et al., Beltagy et al., Han et al. Please check completeness of the references

We will double check and revise these references in updated version.

W1: In depth discussion and review is missing since it is a survey paper. Work could be supplemented by experimental benchmarks

Thank you for your feedback. We have included an in-depth discussion review of the different methods explored in the Key Takeaways section. Besides, we have also provided our insightful perspective on this critical problem and pointed out promising potential future directions for the research.

Moreover, we've conducted a needle-in-a-haystack test on the methods mentioned in the Post-Training Optimization section, as most of them did not test this in their paper. We also planned to add a section of the pros and cons of various methods in section 5. Results are as follows.

✓ and ✗ refers to suitable and not suitable. ○ refers to partially suitable or depends on the implementation. * refers to results reported within the paper. The ‘Merge’ in the table refers to a type of method that emerged in early April 2024, which merges the KV Cache of different tokens. This is a part that we added when we continued to maintain this work.

W2: The authors didn't discuss financial implications of different approaches while the title says 'Keep the Cost Down'

Thank you for suggesting another possible interpretation of our title. But here, our ‘cost’ is more inclined towards the consumption of performance, specifically in terms of memory, I/O bandwidth, and computational resources. While the KV Cache does indeed impact the financial cost of the model, the financial cost is also influenced by many factors unrelated to language models and subject to rapid changes. These are irrelevant to the issues we are concerned with, so we have not discussed this aspect.

W1: This paper is mainly a literature review without empirical experiments exploring which methods are most effective in which scenarios

Thank you for your careful review. We would like to clarify that in the first paragraph of the second page, we divided the methods into pre-training, post-training, and deployment stages, and discussed the appropriate scenarios for each method. For those with sufficient computational resources, methods like MQA and GQA are recommended. As listed in Appendix A, many organizations that can train a large language model (LLM) from scratch have successfully used GQA and MQA. For those with limited resources, methods in the post-training and deployment stages are more suitable. During the rebuttal period, we also created supplementary tables to help readers gain a deeper understanding of the different KV Cache optimization methods in next section.

W2: It would be nice to provide a table listing the advantages and disadvantages of different methods with suggestions on their best usage

Thank you for your suggestion. During the rebuttal period, we have added a supplementary table to compare different methods, highlighting their advantages and disadvantages along with suggestions for their best usage.

✓ and ✗ refers to suitable and not suitable. ○ refers to partially suitable or depends on the implementation. * refers to results reported within the paper. The ‘Merge’ in the table refers to a type of method that emerged in early April 2024, which merges the KV Cache of different tokens. This is a part that we added when we continued to maintain this work.

W3: Prompt compression techniques are quite related to this topic and related work should be discussed

Please refer to Appendix C where we have included a brief introduction of prompt, or embedding, compression techniques and related work. Besides, we will integrate the content from Appendix C into the main body of the paper for a more comprehensive discussion.

W1: Lack of connections

Thank you for your suggestion. In fact, in Section 5.1, paragraph 5, we have discussed in detail how different works estimate the future significance of a specific token based on the historical information of attention weights, exploring their similarities and differences. This paper focuses on categorizing and summarizing different methods in detail. For eviction methods, we have decomposed them into two lines of approaches: static policies and dynamic policies, based on whether they use manual or dynamic methods to evict unnecessary tokens.

We will include the discussion of Attention Sink and Repetitive Attention Pattern in future versions and establish connections between different works within each category to enhance clarity and coherence. Besides, we will seek for more connections between methods. For example, methods like GQA and Eviction both focused on how to reduce the number of stored variables, making the information more condensed.

W2: The evaluation section falls short

Thank you for your suggestion. We have already listed the mainstream methods for testing ultra-long context capabilities in Section 6.1. Nevertheless, our current chapter layout indeed occupies too much space on common metrics that are less related to big KV Cache scenarios. In the next version, we will focus more on supplementing benches focused on long-text, since they are more relevant to evaluate the KV Cache compression. We will also add discussions on the current shortcomings of benches. The main content is an unresolved contradictions, that is, evaluation methods with long texts are often simple in form and task, due to the sparsity of critical information; evaluation methods that are complex in form and difficult in task are difficult to extend to sufficient lengths.

Moreover, for methods mentioned in section 5, we provide a table describes the pros and cons of them.

✓ and ✗ refers to suitable or not. ○ refers to partially suitable or depends on the implementation. * means result reported by the paper. The ‘Merge’ in the table refers to a type of method that emerged in early April 2024, which merges different tokens' cache.

W1: The non-core sections are substantially weaker. There's a noticeable disconnect at the 5-6 section border

Thank you for your valuable feedback. The primary goal of Section 6 was to provide beginners with an understanding of the commonly used datasets and metrics for KV Cache optimization. In the next version of our paper, we will restructure the content by reducing and moving Section 6 to the appendix. Additionally, we will offer a more detailed discussion of the results table and integrate the content from Appendix C into the main body of the paper. We will also enhance the logical flow between sections.

W2&Q1: Does $p_i$ denote a distribution or a single probability

As stated in Section 2, $p_i \in \mathbb{R}^V$ denotes a distribution, not a single probability.

Q2: how do $𝑛_ℎ$ combination weights result in $𝑛_ℎ$ combinations?

As explained in section 3 in page 3, in MQA, all heads share the same Key and Value matrices, while each head retains its own set of Query parameters. Therefore, for $n_h$ heads, there are $n_h*1=n_h$ different QKV combinations.

Q3: In Section 3, what's the correct formula?

We apologize for the typo. The correct formula is:

$$\eta = \frac 12 \cdot \left( 1 + \frac{n_g}{n_h} \right)$$

where the weight of the queries and the output projection layer remains unchanged, but the Key and Value matrices are reduced to $\frac{n_g}{n_h}$ of their original size.

W3 :Some illustrations is confusing and inconsistent

Thank you for your feedback. The purpose of Figure 1 is to provide a roadmap to help readers gain a clearer understanding of the overall structure of the paper. Our paper unfolds in a chronological order of LLM: (1) In the training stage, we introduce the KV Cache compression methods that can be used during model pre-training. These methods are usually the most effective, but are not suitable for modifying existing models with low computational power. (2) In the deployment stage, we discuss the use of different frameworks to optimize the use of KV Cache which do not require significant modifications to the KV Cache itself but can significantly enhance its efficiency. (3) In the post-training stage, we introduce a large number of on-time optimization methods for KV Cache.

We will improve the illustrations to provide more clarity and detailed explanations.

W4: Grammar

We will thoroughly optimize and proofread the presentation to address the grammatical errors.