MatryoshkaKV: Adaptive KV Compression via Trainable Orthogonal Projection

W2 (Q3): It would be ideal to have a task-agnostic KV cache compression scheme. If the second stage of training can just be limited to projection tuning？

We clarify our experiments on continual pretraining are what you desired. We add the SFT experiments to only prove that our method can effectively handle certain specific or more challenging tasks (such as GSM8K). Generally, LLMs also require task-specific SFT for handling such tasks.

Moreover, the advantage of our 2-stage training pipeline is that, compared to CPT, it has lower costs and requires less time. For example, for our SFT on GSM8K, the first stage only consumes 0.5 A100(40G) GPU $\times$ hour, and the second stage only consumes 1.5 A100(40G) GPU $\times$ hour. This is much less than the cost of the CPT scenario, which consumes 30 A100 GPU $\times$ hour. Despite the lower cost, satisfactory performance can be achieved on more difficult tasks (such as GSM8K). Compared with the normally LoRA-finetuned Llama2-7b (Lora rank is set as 8), we can achieve approximately 90% accuracy with only 50% of the KV cache.

We have attempted to tune merely orthogonal projections, but we find that this would significantly affect the convergence of the model. We believe that for some difficult downstream tasks such as GSM8K, perhaps the model requires a higher degree of freedom to fit this data distribution.

W3: Baselines & Using the approach concurrently with other techniques.

Thanks for your advice on supplementing baselines.

We first clarify that our PCA baseline is roughly identical to Eigen-Attention [2], and the reproduction details are described in Section 3.2). The results of the PCA baseline in our paper are similar to those in Eigen-Attention [2] (see Table 2 in [2] and Table 1 in ours). We will make it clear in the upcoming revisions.

Besides, we have incorporated ASVD [3] as our baseline, the results are supplemented in Table 3 in General Reply.

As ASVD only furnishes checkpoints for three cache budgets, namely 85%, 90%, and 95%, we have thus compared our method with ASVD within these three specific budget scenarios. The results exhibit that our MatyoshkaKV performs better than ASVD at all three cache budget levels.

In Section 5, we apply MatryoshkaKV to Mistral-v0.3-7B-base, and it evidently demonstrates enhanced compression capabilities when used together with Group Query Attention (GQA). Moreover, we also combine our methods with H2O(token eviction and merging) and KIVI (a KV quantization technique)

For the combination with H2O, our results are supplemented in Table 1 in General Reply. According to the results, when MatryoshkaKV and H2O are used concurrently, the perplexity on long contexts only increases by 1.02 at a 10% KV cache budget. Additionally, if we compress by 50% on both the sequence length and feature dimension axes (with an actual cache usage rate of 25%), we can achieve an average accuracy of 55.85 on 6 benchmarks, which is 91.32% of the baseline. This is much better than the effect of only using 25% separately on these two individual axes.

For the combination with KIVI, our results are supplemented in Table 2 in General Reply. The combination of MatryoshkaKV and KIVI doesn't lead to a significant decline in accuracy on six benchmarks mentioned in Section 5, showing that MatryoshkaKV can be integrated with KV quantization to achieve a higher compression ratio.

W4: Results on the more recent family of Llama models (Llama 3/3.1) are missing.

Thank you for the comment. First of all, we would like to clarify that we have covered multiple architectures, including LLaMA2-7B-base and Mistral-7B-v0.3-base. The difference between LlaMA3 and LLaMA2 is not significant, which mainly lies in whether GQA is used. We've already demonstrated that our method can be compatible with GQA by applying our MatryoshkaKV to Mistral-7B-v0.3-base. Therefore, it can be presumed that our method basically works for LLaMA3 as well. Additionally, related works such as MiniCache [4] have also conducted experiments on LlaMA2 and Mistral. We will attempt to include the results of Llama3 in the final version.

[1]: Roberts et al., "TensorNetwork: A Library for Physics and Machine Learning", arXiv 2019

[2]: Saxena et al., "Eigen Attention: Attention in Low-Rank Space for KV Cache Compression.", ArXiv 2024.

[3]: Yuan et al., "ASVD: Activation-aware Singular Value Decomposition for Compressing Large Language Models", ArXiv 2024

[4]: Liu, A., Liu, J., Pan, Z., He, Y., Haffari, G., & Zhuang, B. "MiniCache: KV Cache Compression in Depth Dimension for Large Language Models", NeurIPS 2024.

Thank you for your careful reading and many pertinent suggestions. We are also very grateful for your advice on supplementary experiments. We have revised the unclear parts in the paper as quickly as possible and presented the results according to your experimental suggestions.

W1 (Q2): Which calibration dataset was used for the results presented in Table 1? The necessity of training every time?

Firstly, we would like to clarify that in the continual pretraining scenario, our approach encompasses two datasets: a training set, Redpajama, which is task-agnostic and utilized for training our orthogonal matrices via single-time continue pre-training (CPT), and a calibration set which is task-dependent and employed solely for the greedy search to attain adaptive compression rates (training-free) before inference on the specific task.

Secondly, the calibration set is randomly sampled from the downstream dataset and composed of merely 5-10 samples. We will have a more clear description in the revised version.

Q1: Can the authors clarify how they achieve performance benefits with heterogeneous ranks across the head dimension?

Indeed, what we have done based on padding currently results in an actual time that is even slightly worse than that of the uniform approach. However, this can be resolved through system-level optimizations (e.g., flattening the caches of different heads for storage, see AdaKV), which is not the main technical focus of this paper. The upper bound of the theoretical acceleration of our approach exists.

On the other hand, in practice, our uniform strategy is also quite powerful. It can also achieve an average accuracy of 55.02 on 6 benchmarks with a 37.5% cache budget, which is 89.96% of the baseline. The uniform strategy doesn't require those system implementations and takes about the same amount of time as the original attention in real-world scenarios as listed below.

The inference speed of our MatryoshkaKV under the uniform compression rate during the inference process with a batch size of 32.

We clarify our paper mainly focuses on reducing storage and memory transfer costs instead of runtime, and we have proven that our approach indeed consumes significantly less storage cost than the baseline. We will make these points clear in the revised revision.

[1]: Roberts et al., "TensorNetwork: A Library for Physics and Machine Learning", arXiv 2019.

[2]: Saxena et al., "Eigen Attention: Attention in Low-Rank Space for KV Cache Compression.", ArXiv 2024.

[3]: Yuan et al., "ASVD: Activation-aware Singular Value Decomposition for Compressing Large Language Models", ArXiv 2024

[4]: Liu, A., Liu, J., Pan, Z., He, Y., Haffari, G., & Zhuang, B. "MiniCache: KV Cache Compression in Depth Dimension for Large Language Models", NeurIPS 2024.

The authors have resolved most of my concerns, so I am happy to increase my score.

I have one question related to the task-dependent greedy search: does it mean that whenever a model is used for inference on a new task, it needs to go through the greedy search for adaptive compression rates? Can the authors clarify the increase in inference latency/run-time due to this step?

We sincerely thank the reviewer for the time to read our paper. We are glad you thought our paper was easy to follow.

Q1 (W1): It is unclear whether the novelty of paper is significant.

Thanks for the comment. We would like to clarify that solving the KV cache compression problem with the proposed trainable orthogonal projection and Matryoshka training strategy is non-trivial. To achieve KV compression, many previous works are based on token eviction and merging. In contrast, we focus on compression in the feature dimension, which is completely orthogonal and compatible with them. The corresponding results are displayed in our response to Q2(W2).

By tuning the orthogonal projection with the Matryoshka strategy, we have addressed the issue of training-free PCA to lead to sub-optimal solutions. The Matryoshka strategy ensures the hierarchy of the columns in the orthogonal matrices during training, enabling users to arbitrarily choose KV cache compression ratios for tasks of different difficulties. We have also experimentally demonstrated the importance of this strategy in the Ablation part of Section 5.3.

We kindly ask the reviewer to re-evaluate the contribution of this paper and welcome more specific comments on this.

Q2 (W2): Can author show MatryoshkaKV can be combined to further improve performance?

Thanks for the kind suggestion. In particular, the application of MatryoshkaKV to Mistral-v0.3-7B-base in Section 5 can serve as a combination of our method with other dimensions of KV compression because Mistral-v0.3-7B-base uses the Group Query Attention (GQA), which has already compressed the KV in head-number axis.

Additionally, we combine our methods with H2O, which operates on the sequence axis, and the results are supplemented in Table 1 in General Reply. We see our approach can be effectively combined with H2O, the perplexity on long contexts increases by merely 1.02 at 10% KV cache budget. Additionally, if we compress by 50% on both the sequence length and feature dimension axes (with an actual cache usage rate of 25%), we can achieve an average accuracy of 55.85 on 6 benchmarks, which is 91.32% of the baseline. This is much better than the effect of only using 25% separately on these two individual axes.

Moreover, the results of the combination of MatryoshkaKV and KIVI (a KV quantization technique) are also supplemented in Table 2 in General Reply.

The results also reflect our MatryoshkaKV can be used in synergy with KV cache techniques on other dimensions.

Dear reviewer,

The revision deadline is approaching, and we sincerely hope to receive any concerns or questions you may have beforehand. If there are any issues, please let us know at your earliest convenience so we can address them promptly.

Best regards,

The Authors

We would like to express our sincere gratitude for your careful review and the valuable concern you raised regarding the time consumption of our work. Your feedback has provided us with an opportunity to further clarify and enhance the description of our method.

W1 (Q1): Lack of Runtime Evaluation.

We first clarify that in the continual pretraining (CPT) experiments, the process of our approach is (1) obtaining the PCA initialization based on a small subset of a general corpus, (2) training our model on the CPT corpus, (3) searching for the heterogeneous compression levels for various heads with a small calibration dataset (5 - 10 samples) on the specific task of concern, and (4) performing inference on that task given the identified compression levels. The time for steps (1-3) can be substantially amortized by multiple downstream tasks, and thus should not be a concern. Besides, the training time for our approach on the CPT corpus is 8 hours using 4 A100 GPUs. Regarding the wall-clock time for per inference step, our current implementation consumes a similar time as the baseline full-KV model:

The inference speed of our MatryoshkaKV under the uniform compression rate during the inference process with a batch size of 32.

We clarify our paper mainly focuses on reducing storage and memory transfer costs instead of runtime, and we have proven that our approach indeed consumes significantly less storage cost than the baseline. We will make these points clear in the revised revision.

W2: Missing State-of-the-Art Comparisons.

Sorry for the missing comparisons. We first clarify that our PCA baseline is roughly identical to Eigen-Attention [1] (the details are described in Section 3.2) and has a similar performance as Eigen-Attention [1] (see Table 2 in [1] and Table 1 in ours), so we take our PCA baseline as a surrogate of Eigen-Attention [1]. Our existing results have reflected the superiority of our method.

Regarding HeadKV, we first clarify that it has not been open-sourced. It uses low-rank orthogonal matrices and merges multiple heads into one. There are certain similarities between HeadKV and our method: both use orthogonal matrices and have basically the same parameterization. However, the dimensions in which we and HeadKV perform compression are inherently different, and HeadKV also does not further optimize the orthogonal matrices, hence still suffers from suboptimal performance due to the confounding issues. We have added these discussions to the revision and promise to add HeadKV to the empirical comparison once HeadKV is open-sourced.

Besides, we have compared to another related work, ASVD [2], the results are supplemented in Table 3 in General Reply.

ASVD notices the low-rank property of LLM parameters and compresses the KV cache while compressing the model parameters. Since it only releases checkpoints for three cache budgets of 85%, 90%, and 95%, we mainly compare our method with them in these three budgets. The results exhibit that our MatyoshkaKV performs better than ASVD at all three cache budget levels.

Q2: How is $\Delta r$ determined in the greedy search algorithm? And what values are used in experiments?

We simply choose $\Delta r = d / 8$ in the greedy search algorithm, because our predefined compression rate schedule is { $\frac{i}{8} (i=1,2, ...,8)$ }, These values can be further tuned for better empirical results, left as future work.

Q3: MKV seems to work well with uniform compression rates. Does this always apply to all tasks?

We clarify uniform compression rates do not always work well for all tasks. As shown in Table 3 of Appendix D, on some challenging tasks such as CSQA and ARCC, using a uniform compression rate may lead to an accuracy drop of more than 2 percentage points. Specifically, at a 50% cache budget on ARCC, after applying the greedy search algorithm, the accuracy increases from 34.34 to 36.61. Thus, the flexible identification of the compression levels for various heads can be necessary in practice. And, we would like to comment that the greedy search-based identification algorithm can be efficient enough in practice, using a calibration set of only 5 - 10 samples.

Q4:: k→r?

We sincerely apologize for this error. We have already corrected it.

[1]: Saxena et al., "Eigen Attention: Attention in Low-Rank Space for KV Cache Compression.", ArXiv 2024.

[2]: Yuan et al., "ASVD: Activation-aware Singular Value Decomposition for Compressing Large Language Models", ArXiv 2024

W3 (part 1): Figure 4 seems to indicate the Matryoshka strategy is much more important than the greedy search

Thanks for your comments. Matryoshka strategy employs randomly sampled compression ratios on each attention head during the training process. It ensures the hierarchization of the orthogonal matrices and decouples the compression ratios of different heads during training, enabling users to arbitrarily choose KV cache compression ratios for tasks of different difficulties. As shown in Figure 4 right, without the Matryoshka strategy, our orthogonal projections fail to adapt and perform effectively at compression ratios that are not encountered during the training phase.

Our greedy search algorithm is to set different ratios for different heads, which aligns with the observed phenomenon of anisotropy in different heads as reported in some related works we've discussed in Section 5.4.

It is noted that without the Matryoshka strategy, it is impossible for us to find the optimal ratio on each attention head through greedy search. Therefore, the two need to be used in combination. On the one hand, each head may only perform well at the trained static compression ratio. On the other hand, the compression ratios between each head are also mutually coupled. This makes it impossible for us to use the greedy search algorithm to find the independent compression ratio on each attention head.

Uniform compression rates indeed work well in some cases, but not for all tasks. As demonstrated in Table 3 of Appendix D, on certain challenging tasks like CSQA and ARCC, employing a uniform compression rate can sometimes result in a significant drop in accuracy. For instance, on ARCC at a 50% cache budget, after applying the greedy search algorithm, the accuracy rises from 34.34 to 36.61. This indicates that in practice, the flexible identification of compression levels for different heads may be necessary.

W3 (part 2): Orthogonal Constraint has less effect when Cache Utilization<0.5.

When orthogonality is not maintained, the model can still learn primary and secondary information and transform the cache features into spaces with reduced dimensions.

However, without orthogonality between each column vector within the orthogonal matrices, the space spanned by the cache feature may not be fully utilized due to a lower rank. This can lead to a noticeable degradation in performance, especially at higher cache ratios.

To sum up, orthogonality ensures that the full-rank space is effectively utilized, supporting consistent performance across different cache budgets.

[1]: Zhang et al., "H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models", NeurIPS 2023.

[2]: Liu et al., "KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache", ICML 2024.

[3]: Aditya Kusupati, et al. Matryoshka representation learning. Advances in Neural Information Processing Systems, 2022.

[4]: DeepSeek-AI at al., "DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model", arXiv 2024.

We would like to express our sincere gratitude for your in-depth and insightful comments on our paper. Your feedback has provided us with valuable directions for improvement, and we truly appreciate the time and effort you have dedicated to reviewing our work.

W1: It lacks a detailed comparison to highlight the advantages of the feature-dimension based compression

Thanks for the question. We clarify that there are no specific advantages of feature-dimension based compression compared to works on other dimensions. Nevertheless, the exploration of compression in the feature dimension has been largely uncharted previously due to the inherent difficulties of compressing features of deep models.

The enabling of the combination of feature-dimension compression with works on other dimensions is another important contribution of this work. We have demonstrated the ability to combine our MatryoshkaKV with the Group Query Attention (GQA) method in Section 5, by equipping Mistral-7V-v0.3-base with our trainable projections.

Additionally, we combined our methods with H2O (token eviction and merging)[1] and KIVI (KV quantization)[2], and the results are presented in Table 1 in General Reply and Table 2 in General Reply.

W2: The justification for using predefined schedules for the Matryoshka strategy and the heterogeneous compression rates with greedy search could be made stronger with more theoretical backing or detailed analysis.

Thanks for the comment. We clarify that our training approach is not highly sensitive to the selection of the schedule. We have added experiments using other schedules in Table 4 in General Reply. It can be observed that the results do not change significantly.

As for the greedy search for the adaptive compression levels algorithm, it is only employed after the training process to obtain heterogeneous compression rates. During the inference phase, this rate remains unchanged. Therefore, it has no connection with whether the training converges or not. During the training, we merely randomly mask the feature dimension of each head. We have found that normal convergence can be achieved in this way.

[1]: Zhang et al., "H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models", NeurIPS 2023.

[2]: Liu et al., "KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache", ICML 2024.

[3]: Aditya Kusupati, et al. Matryoshka representation learning. Advances in Neural Information Processing Systems, 2022.

[4]: DeepSeek-AI at al., "DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model", arXiv 2024.

We are grateful to the author for providing us with numerous suggestions regarding theory and analysis.

W1, W2: The paper lacks rigorous theoretical analysis of why their proposed MatryoshkaKV method works better than PCA-based approaches.

We conduct an in-depth analysis of a single head within a single layer in Appendix A. We find that to achieve the minimum error, it is essential to jointly optimize the projection matrices for both K and V. For example, if the principal components of K and V just lie in each other's orthogonal complement, the approximation error would be extremely large. This shows that methods like PCA, which do not consider the interaction between K and V, are suboptimal. Unfortunately, due to the nonlinear relationship (softmax function) in the calculation between K and V within the attention mechanism, certain mathematical methods such as GSVD (Generalized Singular Value Decomposition) are inapplicable for jointly optimizing the projection matrices of both K and V.

Moreover, the optimal solution also varies with the input data distribution. Given the difficulty in modeling the distribution of all corpora globally, we believe that using a data-driven approach for optimization is a reasonable way to minimize the error of the model after KV cache compression on most tasks. Thus, we make these orthogonal matrices trainable to obtain optimal results. By directly fine-tuning the projection matrices to maximize the data likelihood, we can better adapt to different data distributions and task requirements, thereby enhancing the overall performance of the model.

Q1: Why does the Matryoshka training strategy work better than static compression ratios? What's the theoretical justification?

The static strategy has a significant drawback in that it suffers from poor generalization. Low-rank projections fail to adapt and perform effectively at compression ratios that are not encountered during the training phase, as proven by results in Figure 4. In contrast, our proposed strategy, which employs randomly sampled compression ratios on each attention head during the training process, offers several notable advantages:

• The randomly sampled compression ratios during training endow the low-rank projections with the ability to generalize across all possible compression ratios.
• Setting various compression ratios across different attention heads effectively decouples the working ratios of different heads. This means that during the testing phase, we have the flexibility to set different ratios for different heads, which aligns with the observed phenomenon of anisotropy in different heads as reported in some related works we've discussed in Section 5.4.

Q2: How sensitive is the method to the choice of sampling schedule for compression rates during training?

Thanks for the question. We clarify our training approach is not sensitive to the selection of the schedule. We have added experiments using other schedules in Table 4 in General Reply, where the results do not change significantly.

Q3: Why is PCA initialization critical for convergence? Could other initialization strategies work?

Thanks. We identify the importance of PCA initialization empirically. We have attempted initialization with randn_init, kaiming_init, and xaiver_init, yet none of these methods have proven to be effective. This is due to the introduction of the orthogonality constraint upon the projection matrices, which makes the optimization process probably unstable and suffers from a cold start.

W3: Limited evaluation on very long sequence tasks.

Our approach to KV cache compression is on the feature dim axis rather than the sequence length axis, thus, it applies to both short and long texts. We have supplemented experiments on Longbench to support this view. Our results are displayed in Table 1 in General Reply.

Due to GPU memory constraints, we set the maximum length of training samples to 2048 during CPT. Accordingly, we truncate LongBench samples to 2048 to calculate perplexity. On average, the perplexity only increases by 0.40, 0.76, and 1.30 at 75%, 50%, and 37.5% KV cache budget. Meanwhile, the experiments of combining our method with H2O can also prove that we can achieve a higher compression ratio for long context, which further demonstrates that our method can effectively relieve the I/O bottleneck for long context.

We will add the results of 4K length (the maximum length for LLaMA2-7B-base training samples) in the updated version.