We thank the reviewer for their feedback and suggestions. We appreciate that you highlight the modularity of AQUA-KV and address your concerns below.

My main concern is that in Table 2, where the authors compare their method to other quantization approaches across various LLMs, their method shows a 0.3 improvement over HIGGS on the LongBench average score for Llama 3.X 8B and 70B in the 3-bit quantization setting but requires an additional 0.5 GB of cache memory. Similarly, for 2-bit quantization, it achieves a 0.4-0.6 improvement but demands 0.6 GB more memory on the LongBench average score for Llama 3.X 8B and 70B. This suggests that the experiment may not be fully controlled, making it unclear whether the observed gains in LongBench performance are significant or if any improvement would persist under a similar cache size constraint. Can the authors have a more controlled experiment, or explain why the existing experiment is controlled enough ?

We agree that our comparison can be improved by controlling strictly for the cache size. To address this, we conducted additional evaluations with raw HIGGS (our strongest baseline), where the quantizer was given more quantization clusters to increase the average bitwidth. To recall, a 2.x bit HIGGS splits the data into $d{=}2$-dimensional vectors and rounds each vector to one of $n{=}2^4{=}16$ clusters, with an additional 16-bit scale per $g{=}1024$ quantized values (the default configuration from [1]). This yields $\log_2 n / d + 16/g \approx 2.0156$ bits per parameter.

[1] https://arxiv.org/abs/2411.17525

To offset AQUA-KV predictors, we evaluate HIGGS with a lattice of $n{=}18$ clusters instead of 16, with an average bitwidth of $\approx 2.1$, resulting in a slightly larger overall cache size than for AQUA-KV.

Table: WikiText-2 Perplexity (non-Instruct models, setup from Section 4.2)

Table: Average LongBench scores (Instruct models, setup from Section 4.2)

As we can see, HIGGS with additional clusters does indeed perform better at the cost of a greater memory footprint, but AQUA-KV still outperforms it.

Please also note that the AQUA-KV memory overhead can be reduced by quantizing the predictor weights. We report this setup in Table 1 (see “GPTQ”, L359 for 4-bit quantization). We discuss this in L346-358 (right) and report more detailed results in Table 4 (Appendix).

We hope that these new results alleviate the reviewer’s concern and will add them to Section 4.1 and perform additional evaluations in Appendix.

The other question I have is whether the authors compared with KV cache compression techniques based on token selection and pruning, such as StreamingLLM, SnapKV since the authors have compared their method with H2O.

In our work, we chose the token pruning proposed in the H$_2$O paper as it was a middle ground between StreamingLLM [1] and more recent methods such as SnapKV [2]. In principle, AQUA-KV can also be combined with other token pruning strategies, including the two you proposed. We agree that exploring these combinations can further strengthen our paper, but we need additional time to incorporate them into our codebase and ensure that our experiment setup uses these pruning strategies properly. We thank the reviewer for the suggestion and will add these comparisons in Section 4.3 in the final version of the paper.

[1] https://arxiv.org/abs/2309.17453 Xiao et al, 2023. Efficient Streaming Language Models with Attention Sinks

[2] https://arxiv.org/abs/2404.14469 Li et al, 2024, SnapKV: LLM Knows What You are Looking for Before Generation

We thank the reviewer for their valuable feedback. Overall, the review appreciates the efficacy and simplicity of AQUA-KV, but suggests additional evaluations on extra benchmarks and asks follow-up questions about implementation details. We do our best to address these below.

Additional evaluations

In Table 2 of the paper, there is only a small difference between the performance of HIGGS on Llama 3.1 70B compared to AQUA-KV with HIGGS. Much more rigorous evaluation is required to see if AQUA-KV is effective for large models as well as smaller models.

First, we would like to note that the quantized 70B model already has very high quality, meaning that any absolute changes in score will be small. Thus, we ask the reviewer to take into account the relative difference: for 2.x bits per value, the 70B model perplexity with 16-bit cache is 2.54. AQUA-KV @ 2-bit increases only PPL by +0.08, whereas raw HIGGS increases it by +0.23 (almost 3x). Likewise, the LongBench score drops by 0.13 for AQUA-KV and 0.74 for the nearest 2-bit baseline (>5x error increase).

To fully address the remark concerning larger-scale evaluations, we evaluate the 72B Qwen 2.5 model in the same setup as in Section 4.2:

We also report additional benchmarks for Llama 3.1 70B below and in our response to Reviewer AgDD. To further explore the scalability, we will run the remaining evaluations for 70B+ models in the final version of the paper.

Language generation tasks by instruction-tuned models should be evaluated rigorously as these are closest to those used for actual LLM production. MMLU, GSM8K, HumanEval, and IFEval are frequently used.

We agree and evaluate GSM8K and IFEVAL across different models, including 70B. We prioritize 2.x bit evaluations due to time constraints and since this is where augmenting quantizers makes the most sense.

GSM8K accuracy (%) for Instruct models in the same setup as Section 4.2

IFEval accuracy (%) for Instruct models in the same setup as Section 4.2

The results show a similar trend to our LongBench evaluations, with AQUA-KV being substantially closer to the uncompressed baseline than raw HIGGS. We will include these results in the final version of the paper and conduct additional experiments with the other two benchmarks (MMLU and HumanEval).

Although long sequence evaluation is also important, evaluations on shorter sequences should also be conducted.

We hope that this concern can be alleviated with the results we reported above, since those benchmarks have shorter sequences (e.g. GSM8K question plus answer takes up, on average, 198 tokens for Llama-3.1/3.2 tokenizer). We also report perplexity with shorter sequence length (see the first table in our response to Reviewer AgDD). We do note, however, that KV-cache compression is most effective in the long-context regime.

Questions about overlapping AQUA-KV with model inference

Could the prediction for the keys be overlapped with the calculation of the previous FFN layer?

Thank you for this suggestion. It is indeed possible to overlap FFN/MLP computation with computing the next layer AQUA-KV cache. Furthermore, since the value predictor is linear, we can overlap half of its computation (from previous values) as well, then add the other half after the fact.

Could the benefit of applying the keys to reconstruct the values be investigated more thoroughly? This creates a dependency that prevents overlapping.

We have investigated this question via ablation analysis in Table 4, section “predictor inputs”, on L809 (w/o $K_{rec}$ → V). To summarize, the key predictor does improve perplexity and LongBench scores, but only slightly (in the last digit). This component can indeed be removed in cases where one cares about better overlap. We will discuss this trade-off in Section 4.1.

I am willing to change my rating if these concerns are addressed.

We hope that the additional evaluations and discussions we provided can alleviate the reviewer’s concerns. If you add any follow-up suggestions in the next discussion phase, we will address them in the final version of the paper.

We thank the reviewer for the thoughtful feedback and valuable suggestions. We agree that improving clarity and expanding evaluations would strengthen the paper, and we address these points below:

What does "high-compression mechanisms for internal network states" mean in the abstract? Is it an observation?

We meant to say that there exist methods that quantize KV states to low-bitwidth with small quality degradation. In the final version of the paper it could be rephrased to “In this work, we aim to improve Key & Value compression by exploiting two observations: … 2) the existence of high-compression methods for internal network states (e.g. attention Keys & Values).”

In L59-60 in the paper, the author mentioned "vector quantization". However, I can not find any details.

By “vector quantization” we referred to HIGGS, which is a vector quantization method, as described in L161. We meant that using HIGGS with predictors provides the best quality compared to other methods for KV quantization, as discussed further in Table 2. We will clarify that by adding “and the more advanced vector quantization scheme HIGGS” in L272.

L196-197 needs to be elaborated.

These lines are indeed somewhat convoluted. We meant the following:

1. we noticed that using 1- and 2-bit quantizers (e.g. HIGGS) can ‘explain’ ~0.75 and ~0.89 of the variance respectively. In other words, they have a ~0.25 and ~0.11 relative quantization error.
2. If a probe can predict keys/values with the same error as a 1-bit quantizer, we found that we can use 1 less bit for quantization (e.g. 3-bit instead of 4-bit) after the residual with, on average, the same accuracy (e.g. see Table 2).

We will clarify this in the revised paper.

The experiments are only conducted on Wiki2PPL and some LongBench tasks. It lacks evaluations on more challenging tasks such as math and code tasks with CoT prompts.

We agree that AQUA-KV can benefit from additional evaluations on larger models. As requested, we have conducted evaluations on GSM8k (CoT) and IFEval and report them to in the tables below.

GSM8k accuracy (%) for Instruct models in the same setup as Section 4.2

IFEval accuracy (%) for Instruct models in the same setup as Section 4.2

The results generally align with the trends observed in LongBench evaluations, with AQUA-KV being substantially closer to the uncompressed baseline than raw HIGGS. We will add these evaluations to the final version of the paper and conduct additional experiments for code benchmarks (e.g. HumanEval).

We appreciate the reviewer’s insightful comments, which have helped us improve the paper’s clarity and experimental scope. The additional evaluations confirm AQUA-KV’s consistent performance across domains, as noted in our response. We hope these revisions address all raised concerns. If you add any follow-up suggestions in the next discussion phase, we will address them in the final version of the paper.

We thank the reviewer for their thoughtful feedback and are glad that they appreciate our method's design and empirical results. Below, we provide detailed answers to the questions posed in the review:

If you train your predictors post-ROPE for 8192 sequences, do you see any deterioration in inference on sequences beyond 8192 tokens?

In short, we found that AQUA-KV is not sensitive to the training sequence length: specifically, we did not see a significant impact on longer sequences present in our LongBench evaluation (with some tasks in excess of 100k tokens [1]). We attribute this to the fact that AQUA-KV only trains simple linear predictors.

[1] https://github.com/THUDM/LongBench

However, we agree that it is important to analyze the effect of training sequence length. To that end, we evaluate AQUA-KV for the Llama 3.2 3B model with varying training sequence length and measure the impact on perplexity.

Table 1. WikiText-2 PPL for different training sequence lengths.

Note that in these evaluations, we control for the total number of tokens: every time we halve the sequence length, we also double the number of sequences in the calibration dataset. Otherwise, training with 64-token sequences would overfit because of insufficient training dataset size. For convenience, we report additional sequence lengths in https://anonymous.4open.science/r/rebuttal-pics-24A3/.

How do you think about combining multiple compression techniques together -- low-rank, quantization, predictor-based, pruning,etc.

Our approach is indeed compatible with different cache compression techniques applied simultaneously. Specifically, we combined AQUA-KV with H$_2$O pruning in Section 4.3 (detailed results provided in Appendix E) and show that AQUA-KV can augment H$_2$O to further improve its size-to-accuracy trade-offs by combining three techniques together: pruning (H$_2$O), predictors (AQUA-KV) and quantization (HIGGS). We also explored low-rank predictors in Appendix C (Table 4), which can further reduce the memory footprint at the cost of degraded performance. In case you are interested in other specific combinations, we will consider them and add to the final version of the paper.

What in your opinion is the lower-bound (bits / token) on compression for reasonable performance. Like can we go below 2 bits?

In our work, we focused on ~2 bits per value because this setup can achieve favorable quality-to-size trade-offs for practitioners. However, it is indeed interesting to evaluate AQUA-KV in the extreme sub 2-bit setup. To that end, we evaluate AQUA-KV on top of a 1-bit HIGGS variant with d=8 group dimension and n=256 clusters, with the rest of hyperparameters matching our setup from Section 4.1. This results in circa 1 bit per stored value. We report our results for Llama 3.2 3B and 3.1 8B below.

Table 2. WikiText-2 perplexity evaluation of AQUA-KV for 1 bit compression.

Table 3. LongBench evaluation of AQUA-KV for 1 bit compression.

To summarize, AQUA-KV with 1-bit HIGGS quantization can achieve substantially better quality than raw 1-bit quantization, but both methods show higher quality deterioration. Still, this is an interesting setup and a potential frontier for future research. We will add these and additional sub 2-bit evaluations to the final version of the paper.