MiniCache: KV Cache Compression in Depth Dimension for Large Language Models

Thanks to the reviewer for the valuable comments.

Q1: Ad-hoc design choices (mid-layer merging, only two layers) need improvement.

A1: It is worth noting that MiniCache is the pioneer work to explore KV cache compression along the depth dimension (Lines 48-50). This insight has been highly recognized by Reviewer TyMG ("...novel perspective of compressing across layers"), Reviewer fj43 (“The idea ... introduces a new perspective”) and Reviewer ZDdv (“This novel perspective, previously unexplored”).

With this new motivation, the design of MiniCache is initially based on strong empirical evidence as shown in Figure 1(a), where KV cache states exhibit high similarity between the adjacent layers in the middle-to-deep portion of LLMs. Thus,

1. Merging after the mid-layer is a reasonable and good starting point for our early exploration as it would ideally keep most performance of LLMs.
2. More importantly, we have already shown that MiniCache can perform cache-merging across more layers in Figure 4, e.g. from shallow to deep layers. We have also discussed cross-layer merging beyond two layers using an advanced merging algorithm termed Spherical Cubic Interpolation in further work (Lines 339-342).

Furthermore, we agree that our current approach could be slightly simple as an early effort. However, at this stage, we have already conducted comprehensive experiments to demonstrate the effectiveness of our main idea of cross-layer KV cache merging. Notably, this has been recognized by all reviewers. Thus, we will leave more advanced design for future works, as also recognized by Reviewer fj43 ("could inspire further research in inter-layer redundancy exploitation.").

Q2: Results of merging more than two layers?

A2: We have considered multiple-layer merging functions; however, in the current version of MiniCache, we employ SLERP (Spherical Linear Interpolation) as our merging function. This choice is due to SLERP's inherent ability to merge only two vectors.

To extend this capability to more than two layers, please refer to general response Q3.

As discussed in Q1, we will explore merging across more than two layers in future work.

Q3: Merge based on cosine similarity with a threshold?

A3: In our approach, we used angular distance as the threshold for merging the KV cache. Mathematically, cosine similarity $\cos(\theta)$ can be converted to angular distance $d$ using the following relationship: $d = \frac{1}{\pi}\arccos(\cos(\theta))$ Thus, merging based on angular distance is equivalent to merging based on cosine similarity.

Q4: Compare with early layer-exiting methods?

A4: Both early-exiting and MiniCache aim to accelerate inference by addressing the depth dimension. Early-exiting is more effective for compute-bound operations, whereas MiniCache is advantageous for memory-bound operations. Furthermore, early-exiting methods typically require extensive training, whereas MiniCache is training-free, making it broadly applicable to a wide range of off-the-shelf models and scenarios. In future work, we will explore combining these two paradigms. The relevant reference [C] will be included in the revised paper.

Q5: Is the retention index defined per layer in Eq (4)? How to measure d_i, d_{min}, and d_{max}?

A5: In Eq. (4), the retention index is shared for the paired KV states by position between adjacent layers. Specifically, $d_i$ represents the angular distance between the paired tokens, with $d_{min}$ and $d_{max}$ denoting the minimum and maximum distances across all paired tokens between two adjacent layers, respectively. The rationale behind this approach is to retain paired tokens when they exhibit low similarity, which corresponds to a large distance $d_i$ exceeding the retention threshold.

Q6: Why do some tokens have low cosine similarity in Fig 2(b)? Are there any common characteristics?

A6: We conducted a thorough analysis and observed that the salient tokens often include the initial token and punctuation, as illustrated in Figure I and Table IX of the rebuttal PDF.

We have observed that the numerical values of these salient tokens are relatively large compared to other tokens. After passing through the attention layers, these salient tokens demonstrate significant numerical differences, as they receive strong attention despite not being semantically important. Similar phenomena have been observed in previous studies [D][E]. Consequently, the significant numerical differences between these tokens result in low similarity, rendering them unmergeable in our algorithm. We will include the analysis in the revised paper.

Q7: Merging per attention head or entire key/value vectors?

A7: We apply compression to the entire key and value vectors.

Q8: Efficiency comparison without quantization is needed to avoid misleading readers.

A8: Table VII in the rebuttal PDF demonstrates that MiniCache's cross-layer merging without quantization achieves almost lossless performance compared to the FP16 baseline, with a compression ratio of 1.53. In contrast, 2-bit quantization causes performance degradation. For instance, on the GSM8K dataset, performance drops from 0.159 to 0.127, but the compression ratio improves to 3.95.

Our MiniCache method, focusing on the depth dimension, can complement any quantization and existing KV cache compression methods. The results indicate that combining cross-layer merging with 4-bit quantization achieves the optimal balance between performance and efficiency.

To avoid confusion, we will revise the description as ``On the ShareGPT dataset, LLaMA-2-7B with cross-layer merging achieves a compression ratio of 1.53. Additionally, since MiniCache is orthogonal to existing quantization techniques, it can achieve a compression ratio of up to 5.02x when combined with the 4-bit quantization technique.”

Note: Please refer to the reference list in the general response.

Dear Reviewer GYW9,

We would like to thank you for your valuable feedback to improve our work. We are wondering whether our response has addressed your questions and can improve your opinion of our work. We are committed to ensuring that our work meets the expectations of our esteemed reviewers.

Kindly let us know if you have any other concerns, and we will do our best to address them.

Best regards, Authors of #1613

Thanks to the reviewer for the valuable comments.

Q1 : Concern about computational overhead during reparametrization and restoration stages. Does this compression strategy increase computational overhead?

A1: Reparametrization-based compression involves computing magnitude and direction vectors, and applying the SLERP algorithm through simple matrix manipulation. Additionally, we compute pairwise distances to identify salient tokens, as detailed in Section 4.2.

For restoration, we utilize an in-place index-based token replacement operation, where the computational overhead is managed by the PyTorch index selection procedure, as mentioned in line 216.

As shown in Table V of the rebuttal PDF, the reparametrization process, including the computation of magnitude and direction, takes 0.093 ms (0.031 ms + 0.062 ms), and the restoration process, including the computation of the distance matrix and token replacement, takes 0.116 ms (0.061 ms + 0.055 ms). These times indicate a negligible computational overhead compared to the overall attention computation, which takes 9.756 ms. The running time is measured in milliseconds per attention layer with a batch size of 1 on LLaMA-2-7B.

Q2: The authors selected the middle layers for cross-layer merging compression. Have the authors considered any principled methods to further increase the compression ratio by compressing the shallow layers in future work?

A2: Yes, we have considered and proposed methods to further improve the compression ratio.

Our observations indicate that certain shallow layers exhibit low similarities, suggesting the presence of differences and residual errors. To address this, we suggest employing a meta-learning-based approach to approximate these differences in the shallow layers. Once these approximations are made, we can bridge the gaps in the shallow layers, enabling their effective merging and compression. The possibility of merging shallow layers has been demonstrated in concurrent work [B], though it trains from scratch.

References:

[B] You only cache once: Decoder-decoder architectures for language models. Arxiv 2024.

Thanks to the reviewer for the valuable comments.

Q1: Lack of Baseline Comparisons. In Table 1, the authors only include quantization methods for performance comparison, neglecting other KV cache eviction methods such as those proposed in [14] and [15].

A1: We have included the baseline comparison with H2O [14] in Appendix Section A, along with additional experimental results. We further add the Attention Sink [15] benchmark, as shown in Table III in the rebuttal PDF.

According to this table, MiniCache consistently outperforms H2O [14] and Attention Sink [15]. Additionally, MiniCache explores the KV Cache in a novel depth perspective, making it orthogonal to existing quantization and sparsity techniques for further improvements.

Q2: Lack of Evaluation for Instruction-following Benchmarks. Given that instruction-tuned models are more generalizable for downstream applications, it is essential to evaluate how MiniCache performs on instruction-following benchmarks such as MT-Bench [A].

A2: We benchmark the MiniCache using the MT-Bench dataset, which focuses on instruction-following tasks, as shown in Table VIII in the rebuttal PDF. The models employed for this benchmarking include LLaMA-2-7B-Chat, Mistral-7B-Instruct, and LLaMA-3-8B-Instruct. Our MiniCache-4bit exhibits a reasonably slight performance reduction caused by 4-bit quantization. However, MiniCache can achieve almost lossless performance via standalone cross-layer merging.

Q3: Latency Benchmark. Beyond memory compression, what is the impact of MiniCache on latency?

A3: We benchmark the latency of LlaMA-2-7B on an NVIDIA A100 GPU using different sequence lengths ranging from 1024 to 4096 with a batch size of 16, as shown in Table IV in the rebuttal PDF. We compare it with H2O, which requires calculating full attention scores to estimate token importance. In contrast, MiniCache performs reparameterization-based merging and token restoration using simple matrix manipulations, resulting in more lightweight computations and lower latency. Specifically, when the sequence length is 4096, MiniCache shows a 36.83% reduction in latency compared to H2O.

References:

[A] Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena, NeurIPS 2023.

Thanks to the reviewer for the valuable comments.

Q1: Distance-Based Threshold for Retention (Line 226). Why did you choose a distance-based threshold instead of the merging ratio of overall tokens as the control for retention? Can you show the effect on accuracy and efficiency as the ratio of retention tokens varies, highlighting this trade-off?

A1: According to our Observation 2, token pairs with lower similarity scores are more critical for restoration, and different adjacent layers exhibit distinct patterns.

Our proposed dynamic distance-based threshold allows us to selectively retain salient tokens according to low similarity scores. Setting the retention ratio of the overall tokens as a hyperparameter is impractical because determining the optimal retention ratio for each layer manually is challenging due to the dynamic nature across different layers.

To illustrate the trade-off regarding accuracy and efficiency, we conducted experiments by varying the ratio of retention tokens, as shown in Table I in the rebuttal PDF. We observe that as the overall token retention ratio increases, accuracy improves up to a certain point before plateauing. Retaining 20% of the tokens is necessary to ensure all salient tokens are preserved without performance degradation. In contrast, using a dynamic distance-based threshold, as proposed in our paper, we only need to retain the top 5% of the salient tokens. Note that efficiency decreases as more tokens are retained. This demonstrates that our distance-based approach better balances performance and efficiency than the fixed retention ratio counterpart.

Q2: Ratio of Merged Tokens in Ablation Study. In the second ablation study, can you show the ratio of merged tokens and the efficiency trade-offs for different settings?

A2: We are adding another column in terms of compression ratio to represent the efficiency trade-offs, as shown in Table II in the rebuttal PDF. The results suggest that setting γ to 0.05 achieves the best balance between performance and efficiency.

Q3: No Implementation Source Code Provided. Can you include the implementation source code?

A3: We will definitely release our source code upon acceptance.

To ensure reproducibility, we have included comprehensive pseudocode in the Appendix, specifically in Algorithm 1 and Algorithm 2.

• Algorithm 1: This algorithm outlines the overall logic of the MiniCache Inference process, covering both the prefilling and decoding stages. It provides a step-by-step breakdown of the core processes involved in our approach.
• Algorithm 2: This algorithm details the compression and restoration processes. It elaborates on how the KV cache is compressed and subsequently restored, highlighting the technical intricacies and optimizations employed in our method.

We believe that the inclusion, along with the detailed comments, will significantly aid the community in grasping the logic and mechanics of our implementation.

Q4: Justification for SLERP. Can you justify your choice of SLERP over other interpolation methods?

A4:

1. Motivation for Using SLERP: We conceptualize the KV Cache as activations, which consist of two factors: magnitude and direction in the spherical feature space. SLERP allows us to use a more compact form to represent the original KV Cache by merging states effectively. We can compute the overall magnitude vector and store it in a channel with a dimension equal to 1, significantly reducing the memory overhead, as shown in section 4.2.

2. Performance Metrics and Empirical Evidence: In our initial experiments, we considered average merging as the preliminary method; however, the performance in terms of accuracy was not promising, as shown in Figure 2(a). Subsequently, we conducted experiments using maximum norm-preserving interpolation, but it demonstrated lower accuracy compared to SLERP, as detailed in Appendix Section C (Line 579). Ultimately, we selected SLERP [64] as our final solution due to its multiple advantages in key performance metrics, including compression ratio, accuracy, and computational efficiency.

To further substantiate our claim, we conducted additional ablation studies comparing SLERP with average merging and maximum norm-preserving interpolation. The results, as shown in Table VI in the rebuttal PDF, reaffirm that SLERP provides superior performance in information preservation.