KV-Dict: Sparse KV Cache Compression with Universal Dictionaries

We thank the reviewer for recognizing the novelty of our work in exploring the use of dictionary learning in KV cache compression. We address each concern below.

W1: Presentation regarding motivation and baseline methods

We respectfully disagree with the reviewer’s concern that the problem highlighted in our introduction is not well-answered. As the reviewer noted, we emphasize that prior methods are not off-the-shelf and often discard the majority of tokens. Our method, KV-Dict, directly addresses these issues by being readily usable off-the-shelf and retaining all tokens without any discarding.

To clarify any confusion, our choice of KIVI as a baseline may have led to misunderstanding, as KIVI already partially addresses these limitations. To resolve this, we have made two updates:

• Revised introduction (L42-48): We explicitly outline the limitations of both eviction-based and quantization-based methods. Specifically, eviction methods can degrade long-context accuracy, while quantization methods have strict upper bounds on compression rates.
• Expanded baselines: Our experiments now include multiple additional baselines, including PyramidKV and SnapKV for eviction-based methods, and ZipCache and per-token quantization for quantization-based methods.

W2: Using GSM8K as a KV cache compression benchmark

We respectfully disagree with the notion that GSM8K is not appropriate for measuring KV cache compression. In terms of its use as a benchmark, GSM8K is widely recognized in the literature as a standard benchmark for KV cache compression (e.g., [3,4,5,7]), particularly because reasoning and long-generation tasks are highly sensitive to even minimal degradation caused by compression. We also evaluate GSM8K with 5-shot prompting which is longer than just one GSM8K question. Additionally, we complement GSM8K with LongBench, which evaluates long-context tasks such as summarization, information retrieval, and document QA. These benchmarks collectively capture diverse dimensions of KV cache compression performance—GSM8K tests reasoning robustness during generation, while LongBench evaluates token preservation for long-context understanding. Together, they provide a comprehensive assessment necessary for modern LLMs.

W3&4: Generalizability and value cache issue

We believe the reviewer's concern stems from a misunderstanding of our setting. The elements in Figure 2 have been sorted by similarity to highlight the clustering behavior, as stated in L159-160. This sorting is not reflective of sequence order but rather aims to visualize the inherent redundancy and clustering among the key and value states.

Regarding the generalizability concerns (W3), the clustering behavior observed does not imply that the number of clusters will scale proportionally with the sequence length. Instead, it highlights that key and value states naturally form clusters, irrespective of their positions within the sequence. This property is independent of sequence length and supports the universal nature of the method.

Regarding the value cache concerns (W4), our statement is that both the key and value states exhibit clustering behavior, irrespective of their sequence positions, rather than being similar solely to neighboring states in the sequence. While the clustering effect is more pronounced in key states, a similar behavior is also evident in value states. This finding supports the universality of our method, though we acknowledge that clustering in value states is not as strong as in key states.

Thank you for your valuable feedback on our paper. We have provided our responses to your comments and would greatly appreciate it if you could kindly review them. Please let us know if there are any additional points we can clarify to assist in your evaluation.

We thank the reviewer for recognizing the strengths of our work and for providing helpful suggestions and comments.

W3: Varying model sizes

We have conducted additional experiments on 1B and 3B models and present our findings in our General Response "Q2: Additional experiments comparing more compression methods, model sizes, and tasks". Our results demonstrate that KV-Dict remains effective across diverse model sizes, and its universal property holds consistently. This is shown through comparisons with five KV cache compression methods: Per-Token Quantization, KIVI, ZipCache, SnapKV, and PyramidKV.

W1&2,Q1&2: Latency and throughput measurements

We thank the reviewer for highlighting concerns about latency and recognize its importance in certain deployment scenarios. However, we would like to clarify that the primary focus of our method is addressing highly memory-constrained environments, which are increasingly relevant for real-world deployments of large language models (LLMs). This is reflected in the growing interest in compression techniques [2, 3] and open-source tools such as BitNet. In such scenarios, even a modest batch size of 1 can surpass memory limitations with existing baseline compression methods.

By contrast, KV-Dict avoids OOM errors by achieving significantly more efficient KV cache compression. This advantage is substantiated by our additional experiments spanning diverse model sizes, tasks, and baselines, demonstrating the method’s robustness and adaptability (see Figure 4 and 5). While this comes with increased latency, we believe the ability to ensure reliable performance under stringent memory constraints is a critical contribution. Notably, latency trade-off is common in KV cache compression, similarly to the classic space-time trade-offs in computer science. For instance, KIVI also exhibits increased latency, as shown in Figure 4 in [1].

We acknowledge the reviewer’s concern that higher latency may be a limitation for latency-critical use cases. While our focus has been on optimizing for memory efficiency rather than latency, we include end-to-end latency measurements for prefilling and decoding: with an input length of 1000 tokens and a batch size of 1, the standard implementation takes 0.14 seconds to prefill and 48.39 ms/token for decoding; KIVI takes 0.15 seconds to prefill and 59.66 ms/token for decoding; and KV-Dict takes 6.98 seconds to prefill and 96.93 ms/token for decoding. Note that implementing more sophisticated approximation algorithms or hardware optimizations (as KIVI does) may yield significant improvement in latency.

In summary, our method addresses a critical challenge in memory-constrained LLM deployments with limitations of increased prefilling latency, offering a valuable contribution to pushing the pareto frontier in highly memory-constrained settings.

References

[1] Liu, Z., Yuan, J., Jin, H., Zhong, S., Xu, Z., Braverman, V., Chen, B. and Hu, X., KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache. In Forty-first International Conference on Machine Learning.

[2] Luohe, S., Zhang, H., Yao, Y. and Li, Z., Keep the Cost Down: A Review on Methods to Optimize LLM’s KV-Cache Consumption. In First Conference on Language Modeling.

[3] Ma, S., Wang, H., Ma, L., Wang, L., Wang, W., Huang, S., Dong, L., Wang, R., Xue, J. and Wei, F., 2024. The era of 1-bit llms: All large language models are in 1.58 bits. arXiv preprint arXiv:2402.17764.

Thank you for your valuable feedback on our paper. We have provided our responses to your comments and would greatly appreciate it if you could kindly review them. Please let us know if there are any additional points we can clarify to assist in your evaluation.

W1: Presentation is hard to follow

We thank the reviewer for pointing out ways to improve our presentation. We have revised our manuscript to include the following edits to address the reviewer's presentation concerns:

• L80-94: In our updated introduction, we have included a summary of the end-to-end KV-Dict pipeline in three steps. Here, we clarify that we train our dictionary using WikiText-103 and apply it to all tasks, which was ambiguous to a couple reviewers. With this summary, we clearly separate our method and contributions, which follow in L96-107.
• L82-83: The method summary explains that we use both sparsity-based compression and quantization. However, the 8-bit quantization is a minor detail that boosts our performance by reducing memory storage of sparse coefficients.
• L289-291: We include omitted details of dictionary pretraining in the Experiments section.
• Figure 4&5: For clearer result presentation, we include multiple plots that demonstrate KV-Dict versus other baseline methods. We plot KV cache size vs. benchmark performance and demonstrate the Pareto frontier of each method.
• We have fixed some grammatical errors as well, though we did not highlight these in red for the sake of readability.

W2: Qualitative results of output samples

As requested, below we provide an example generation for qualitative analysis. We present the original GSM8K problem and solution below, along with outputs from Llama-3.1-8B-Instruct with full KV cache (FP16), KV-Dict (sparsity $s=14$), and KIVI-2. In the example, we see that KV-Dict and KIVI-2 first performs a useless calculation that 3 German Shepherds and 2 bulldogs are 5 dogs. However, KV-Dict gets back on track and generates the correct solution, while KIVI-2 continues on to incorrect reasoning.

We have attached the example in page 16-17, Appendix D.

Problem: The Kennel house keeps 3 German Shepherds and 2 Bulldogs. If a German Shepherd consumes 5 kilograms of dog food and a bulldog consumes 3 kilograms of dog food per day. How many kilograms of dog food will they need in a week?

Solution: The German Shepherds consume 3 x 5 = <<3*5=15>>15 kilograms of dog food per day. The Bulldogs consume 2 x 3 = <<2*3=6>>6 kilograms of dog food per day. The dogs can consume a total of 15 + 6 = <<15+6=21>>21 kilograms of dog food per day.

Therefore, they need 21 x 7 = <<21*7=147>>147 kilograms of dog food in a week.

Final Answer: 147

Llama 8B with Full Cache:

First, find the total amount of food the German Shepherds consume in a day: 3 dogs * 5 kg/dog = 15 kg

Then find the total amount of food the bulldogs consume in a day: 2 dogs * 3 kg/dog = 6 kg

Then add the two amounts together to find the total amount of food the dogs consume in a day: 15 kg + 6 kg = 21 kg

Then multiply that number by 7 to find the total amount of food the dogs consume in a week: 21 kg/day * 7 days = 147 kg

Llama 8B with KV-Dict (s=14):

The kennel has 3 German Shepherds and 2 bulldogs, so they have 3 + 2 = 5 dogs.

The German Shepherds consume 5 * 3 = 15 kilograms of dog food per day.

The bulldogs consume 2 * 3 = 6 kilograms of dog food per day.

So, the kennel needs 15 + 6 = 21 kilograms of dog food per day.

In a week, they will need 21 * 7 = 147 kilograms of dog food.

The answer is 147.

Llama 8B with KIVI-2:

The Kennel house has 3 German Shepherds and 2 Bulldogs, so they have a total of 3 + 2 = <<3+2=5>>5 dogs.

Each German Shepherd consumes 5 kilograms of dog food per day, so the total amount of dog food consumed by the German Shepherds is 5 * 5 = <<5*5=25>>25 kilograms per day.

Each bulldog consumes 3 kilograms of dog food per day, so the total amount of dog food consumed by the bulldogs is 2 * 3 = <<2*3=6>>6 kilograms per day.

The total amount of dog food consumed per day is 25 + 6 = <<25+6=31>>31 kilograms.

The Kennel house will need 31 * 7 = <<31*7=217>>217 kilograms of dog food in a week.

We thank the reviewer for recognizing the contribution of our work, and for the helpful suggestions and comments. Below we address the suggestions and questions by the reviewer ArJZ.

W1: Additional comparison with more KV cache compression methods

Please refer to our General Response, where we address this question in our response to "Q2: Additional experiments comparing more compression methods, model sizes, and tasks".

W2: Broader range of tasks and input distributions

First, we highlight that we originally evaluated LongBench, which covers 8 different tasks ranging from summarization of news articles (MultiNews) to code completion (LCC, RepoBench-P). This already includes a diverse set of benchmarks and word distributions, demonstrating KV-Dict's generalizability.

Additionally, we conducted experiments on the MMLU-Pro dataset (using our universal distionary trained on WikiText-103) to test its performance on an even broader variety of input distributions. Specifically, we chose to evaluate two subjects with very different word distributions: Engineering and Law. Please refer to our General Response "Q2: Additional experiments comparing more compression methods, model sizes, and tasks" for our results and analysis. The results substantiate our claim that our dictionary is generalizable across input distribution.

W3: Does not introduce new dimensions to KV cache compression

We disagree with the reviewer regarding novelty for two reasons. We believe that not only is our method new in its implementation, but also that our results are quantitatively novel.

1. KV-Dict is the first approach to decompose KV cache into a basis and incorporate dictionary learning, to the best of our knowledge. Furthermore, our hybrid algorithm to train the dictionary using OMP as an "encoder" and perform gradient descent on the decoder linear layer (i.e., dictionary) differentiates our method from prior sparse autoencoder literature.
2. Quantitatively, we venture into low-memory regimes in which previous works have been unable to tread; mainly the fact that KV-Dict can employ higher compression rates than 2-bit quantization can offer (L101-104). Both our original and additional experiments support this claim by outperforming quantization and eviction methods across model sizes and multiple benchmarks.

Q2: Supplemental experiments comparing more compression methods, model sizes, and tasks (R1, R2, R3, R4)

We appreciate the reviewers' suggestions to strengthen our claims. In response, we provide three dimensions in which we strengthen our results: comparing against more compression methods, testing more model sizes, and evaluating on additional tasks.

Experimental Setup

Compression baselines. We conduct extensive experiments and provide comparisons with four new KV cache compression methods: Per-Token Quantization, ZipCache [2], SnapKV [3], and PyramidKV [1]. Originally, we have only compared KV-Dict against KIVI.

• Per-Token Quantization refers to the conventional token-wise quantization approach.
• KIVI is a recent asymmetric quantization method, which quantizes keys per-channel and values per-token in groups of size $g$.
• ZipCache employs mixed-precision quantization and achieves state-of-the-art performance in the category of quantization-based compression.
• SnapKV and PyramidKV are state-of-the-art eviction-based compression methods.

Since these methods achieve varying compression ratios, we present our results using a Pareto curve. The x-axis represents the KV cache size, while the y-axis represents benchmark performance. This visualization enables a clear comparison of trade-offs across methods (e.g., see Figure 4 in our revised submission).

Model sizes. We also add two new model sizes: Llama-3.2-1B-Instruct and Llama-3.2-3B-Instruct. Previously, we have examined the models Llama-3-8B, Llama-3.1-8B-Instruct, and Mistral-7B-Instruct-v0.3.

Additional tasks. For all new baselines and new models, we test GSM8K, MMLU-Pro Engineering, and MMLU-Pro Law across multiple KV cache compression rates. We choose the MMLU-Pro subjects Engineering and Law due to their difficulty, as explained in [4]. We choose more difficult subjects because we tend to see stronger differentiation and complex reasoning is often sensitive KV cache compression loss.

Results and Analysis

Our results demonstrate that KV-Dict often achieves stronger performance than all other baseline methods on GSM8K and MMLU-Pro. For full comprehensive results, please refer to Figures 4 and 5 in our revised paper for Pareto curves for each model and task.

GSM8K Results across model sizes

Not only have we originally shown KV-Dict excelling on GSM8K in Table 3, but we have now painted a more complete picture by running experiments across a spectrum of KV cache sizes and baselines. The advantage holds true for all across model sizes, as seen in our results with Llama-3.2-1B-Instruct and Llama-3.2-3B-Instruct on GSM8K.

We also observe that eviction-based methods significantly underperform. This is because PyramidKV and SnapKV are not compatible with Grouped Query Attention (GQA), an architectural component that uses fewer key-value attention heads and frequently used in recent models (e.g., Llama 3, Mistral v0.3, Qwen 2).

MMLU-Pro Results with Llama-3.1-8B-Instruct

When evaluating KV-Dict on MMLU-Pro, we observe similar gains, especially compared to eviction-based methods PyramidKV and SnapKV. One caveat is that we see our performance is slightly worse than ZipCache and KIVI when KV cache size is around 25%. However, in low memory regime below 20% cache, our method still outperforms any other baseline. This highlights that KV-Dict supports a wide range of compression ratios quite effectively and that our dictionary is generalizable across input distribution.