MatryoshkaKV: Adaptive KV Compression via Trainable Orthogonal Projection

Authors: Bokai Lin, Zihao Zeng, Zipeng Xiao, Siqi Kou, TianQi Hou, Xiaofeng Gao, Hao Zhang, Zhijie Deng
Primary area: foundation or frontier models, including LLMs
Keywords: Inference Optimization, KV Cache Compression, Low-rank Projection

Abstract:
KV cache has become a *de facto* technique for the inference of large language models (LLMs), where tensors of shape (layer number, head number, sequence length, feature dimension) are introduced to cache historical information for self-attention. 
As the size of the model and data grows, the KV cache can, yet, quickly become a bottleneck within the system in both storage and memory transfer.
To address this, prior studies usually focus on the first three axes of the cache tensors for compression.  
This paper supplements them, focusing on the feature dimension axis, 
by utilizing low-rank projection matrices to transform the cache features into spaces with reduced dimensions. 
We begin by investigating the canonical orthogonal projection method for data compression through principal component analysis (PCA). 
We identify the drawback of PCA projection that model performance degrades rapidly under relatively low compression rates (less than 60%).
This phenomenon is elucidated by insights derived from the principles of attention mechanisms.
To bridge the gap, we propose to directly tune the orthogonal projection matrix on the continual pre-training or supervised fine-tuning datasets with an elaborate Matryoshka learning strategy.
Thanks to such a strategy, we can adaptively search for the optimal compression rates for various layers and heads given varying compression budgets. 
Compared to Multi-head Latent Attention (MLA), our method can easily embrace pre-trained LLMs and hold a smooth tradeoff between performance and compression rate. 
We witness the high data efficiency of our training procedure and find that our method can sustain over 90\% performance with an average KV cache compression rate of 60% (and up to 75% in certain extreme scenarios) for popular LLMs like LLaMA2 and Mistral.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730590944941

thread_type:
official_review

<thread_object>
creator: reviewer_9usd

<content>
summary:
This work proposes a KV cache compression technique for efficient LLM inference through low-rank projections along the feature dimension.  Using principal component analysis (PCA) over key and value matrices to obtain low-rank counterparts leads to performance loss, especially at high compression ratios. To tackle this, the authors propose to tune the projection matrices for each layer and attention head with a distillation loss. Results are presented for Llama-2 (7b) and Mistral-v0.3 (7b) models.

strengths:
1. This work targets KV cache compression, a crucial problem for efficient LLM inference at large sequence lengths. Results show extensive improvements over vanilla PCA for a variety of downstream tasks.

weaknesses:
1. The training to obtain orthogonal projection matrices involves a KL divergence loss, which ensures that the KV compressed model performance stays close to the original model. However, this makes the strategy task-dependent by using a form of calibration dataset for the downstream task itself, leading to the necessity of training every time one needs a compressed KV cache for performing inference on certain task(s). Additionally, there is no clarification on the calibration dataset used for continual pretraining experiments in Section 5.1.
2. For results in Section 5.2, the authors propose a 2-stage training pipeline: 1. LoRA (standard fine-tuning) 2. Updating projection matrices and LoRA parameters jointly. While the first stage is generally employed to improve the performance of LLMs on downstream tasks, the second stage is the associated overhead with this form of KV cache compression. Additionally, the joint update in this stage implies the need to compress the KV cache specifically for each downstream task. It would be ideal to have a task-agnostic KV cache compression scheme.
3. Comparisons with baselines are missing and/or somewhat ambiguous. Is the PCA baseline in Table 1 the same as Eigen Attention [1]? Another missing potential baseline is ASVD [2], which also involves training-free low-rank projection to reduce the KV cache footprint. The authors clarify that a variety of works compress the KV cache by targeting the sequence length or channel dimension, but don't demonstrate the possibility of using their approach concurrently with such techniques [3,4].
4. Results on the more recent family of Llama models (Llama 3/3.1) are missing but crucial to establish the effectiveness of this approach.

[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.

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

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

questions:
1. Can the authors clarify how they achieve performance benefits with heterogeneous ranks across the head dimension? Different dimensions across the heads may require some form of padding before concatenating them during actual inference, so it would be great to see some hardware performance numbers as well.
2. Which calibration dataset was used for the results presented in Table 1? 
3. For the SFT setup described in Section 5.2, it would be interesting to see if the second stage of training can just be limited to projection tuning instead of the proposed joint tuning with LoRA parameters as well.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691366943

thread_type:
official_review

<thread_object>
creator: reviewer_76bc

<content>
summary:
The paper presents a novel method to efficiently compress Key-Value (KV) cache for large language models (LLMs) like LLaMA2-7B and Mistral-7B, which can be a bottleneck in terms of storage and memory transfer. The authors propose using low-rank orthogonal projection matrices, initialized with PCA and further fine-tuned with a distillation objective, to compress the feature dimension of the KV cache. The novel Matryoshka training strategy allows adaptive selection of compression levels for different layers and heads, balancing model performance and compression rate. Experiments demonstrate high data efficiency, with the method achieving over 90% performance while compressing the KV cache by 60% on average.

strengths:
- The paper introduces a novel method to compress KV cache by focusing on the feature dimension. By employing low-rank projection matrices combined with orthogonality constraints, the authors efficiently reduce the KV cache size without requiring retraining from scratch, allowing the compression mechanism to be integrated directly into pre-trained models.

- The proposed training strategy to fine-tune orthogonal projection matrices effectively preserves model performance while allowing adaptive compression rates, providing a flexible approach to balance resource usage.

- The use of heterogeneous compression rates across different layers and heads is well-motivated and effectively demonstrated

weaknesses:
- Although the paper briefly mentions other KV cache compression methods, such as sharing KV headers across layers and merging redundant tokens, it lacks a detailed comparison to highlight the advantages of the feature-dimension based compression. Including experimental comparisons or a more thorough discussion of the advantages and disadvantages of each approach would strengthen the contribution and clarify the unique benefits of the proposed method.

- The justification for using predefined schedules for Matryoshka strategy and the heterogeneous compression rates with greedy search could be made stronger with more theoretical backing or detailed analysis. For example, are the results sensitive to different schedules? Is the greedy search algo. deterministic and with the grantee to converge, and how's the scalability? 

- The Figure 4 seems to indicate the Matryoshka strategy is much more important than greedy search (yellow and green lines are closed in right), and Orthogonal Constraint has less effect when Cache Utilization<0.5. More discussion and analysis on these findings are encouraged.

questions:
see weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730693154454

thread_type:
official_review

<thread_object>
creator: reviewer_uec1

<content>
summary:
This paper proposes MatryoshkaKV, a method to compress the key-value (KV) cache in large language models (LLMs) along the feature dimension using trainable orthogonal projection matrices. As LLMs grow in size, the KV cache can become a bottleneck in storage and memory transfer. Previous approaches have focused on compressing the cache along the layer, head, and sequence length dimensions. This work explores compressing along the feature dimension.

The authors first investigate using PCA to obtain orthogonal projection matrices for dimensionality reduction of the keys and values in each attention head. While this works well at moderate compression levels without needing training, performance degrades quickly at higher compression.

To improve on this, they propose MatryoshkaKV which tunes the orthogonal projection matrices end-to-end using a knowledge distillation objective and a special "Matryoshka" training strategy that enables adaptively searching for optimal compression rates per layer and head at inference time. The orthogonality of the projections is enforced using a Cayley parameterization.

strengths:
The paper tackles the problem of KV cache compression in LLMs from a new angle by focusing on the feature dimension. While prior work has explored compressing along the layer, head, and sequence length dimensions, this work shows that significant compression gains can also be achieved along the feature axis. This opens up a promising new direction for efficient LLM inference.

The MatryoshkaKV method demonstrates impressive performance in experiments. It can compress KV caches by 60-75% on average while retaining over 90% of the full model's accuracy. This is a significant improvement over the PCA baseline, especially at high compression rates. The results hold across both continual pre-training and supervised fine-tuning settings, showing the approach is robust and widely applicable.

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

While they provide some error analysis in Appendix A, it's relatively brief and doesn't fully explain the theoretical underpinnings of their method's superior performance

Limited evaluation on very long sequence tasks where KV cache compression would be most valuable

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

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

Why is PCA initialization critical for convergence? Could other initialization strategies work?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730706783982

thread_type:
official_review

<thread_object>
creator: reviewer_xwvw

<content>
summary:
The authors study the problem of KV cache compression. While existing work focuses on compression along the layer number, the head number, and the sequence length, the authors work on the feature dimension. While PCA is the most intuitive approach, it does not provide good enough performance. Instead, the authors propose to directly tune the orthogonal projection matrices with a distillation objective using an elaborate Matryoshka training strategy.

strengths:
1. The paper is easy to follow.
2. Much stronger performance than the PCA baseline when the compression ratio is low.

weaknesses:
1. It is unclear whether the novelty of the paper is significant.
2. The paper does not compare with the methods that compress the other dimensions. Thus, it is unclear whether the proposed method is more effective. It is also unclear whether the proposed method can be combined with the others while maintaining its effectiveness.

questions:
1. Can the author provide more insight so that the novelty is more than just the direct application of the Matryoshka training strategy proposed by the other paper?
2. Can the author compare with one to two baselines that compress the other dimensions to show compressing the feature dimension is more effective? Or can the author show that they can be combined together to further improve the performance?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730729851767

thread_type:
official_review

<thread_object>
creator: reviewer_lgqe

<content>
summary:
The paper introduces MatryoshkaKV, a method for compressing the Key-Value (KV) cache in large language models (LLMs) to reduce memory during inference. The method begins with PCA for initial dimensionality reduction but addresses PCA’s limitations by tuning projection matrices through knowledge distillation and applying Matryoshka training strategy to enable adaptive compression, allowing the model to balance performance and compression. Furthermore, this paper demonstrates effectiveness with high compression rates while maintaining relatively high accuracy across various LLMs on both CPT and SFT tasks.

strengths:
1. The proposed Matryoshka training strategy effectively preserves hierarchical structures in orthogonal matrices inherited from PCA at various compression levels, ensuring robust performance across dimensions.

2. Greedy search algorithm effectively adapts to differing sparsity in each 
$𝑊_𝑘$ and $𝑊_𝑣$ matrix, showcasing flexibility in compression rates across layers.

3. There are comprehensive MKV evaluations across cache budgets, which reveals substantial improvements, particularly under extremely low cache budget.

weaknesses:
1. Lack of Runtime Evaluation: The absence of runtime metrics makes it challenging to assess the practical benefits of this method fully (see Questions).

2. Missing State-of-the-Art Comparisons: Unusually, the paper doesn’t thoroughly compare to existing state-of-the-art methods. Although it mentioned the other methods may collapse under  60% cache budget (lines 126-131), a comparison with Eigen-Attention and HeadKV at different cache budgets and tasks in terms of both performance and runtime would strengthen the evaluation.

questions:
1. Although the paper mentions it only needs processing 2 million training tokens (line 104), it does not clarify the runtime for each base model and task. Please provide the runtime details for the KV compression process, including PCA initialization, the greedy search for compression level selection, and fine-tuning in both CPT and SFT tasks. Notably, since both the greedy search and compression levels rely on outputs from the original model, this could potentially double the training and inference times.

2. How is $\Delta𝑟$ determined in the greedy search algorithm? And what values are used in the experiments? 

3. In Fig.4, MKV seems to work well with uniform compression rates. Does this always apply to all tasks? 

3. Line 96: k→r?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732115144922

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response (Part 4/4)

comment:
- [Schedule Choice @LgQe, @uec1]: In Appendix G, we include experiments that demonstrate how our choice of predefined compression rate schedules during training impacts the results, which are also listed here:

Table 4: Accuracy of our MatryoshkaKV after CPT on six benchmarks. We use uniform compression levels for inference here for simplicity. Different hyper-parameters are compared. In the table we donate  the schedule $ M_{1}  =  \frac{i}{4} d (i = 1,2,3,4) $ , and the schedule $ M_{2}  =  \frac{i}{8} d (i = 1,2, ...,8) $. We use uniform compression levels for inference here for simplicity.
| **Model**         | **Budget** | **Method** | **HLSG** | **ARC-C** | **ARC-E** | **PIQA** | **WG** | **CSQA** | **Avg.** |
|-------------------|------------|------------|----------|-----------|-----------|----------|--------|----------|----------|
| LLaMA2 7B-base    | 100.0%     | $\mathcal{M}_{1}$ | 72.03    | 36.61     | 52.56     | 76.71    | 61.64  | 67.16    | 62.07    |
|                   |            | $\mathcal{M}_{2}$ | 72.05    | 37.29     | 52.38     | 76.66    | 61.72  | 67.32    | 61.24    |
|                   | 87.5%      | $\mathcal{M}_{1}$ | 72.03    | 37.29     | 53.09     | 76.28    | 62.75  | 65.77    | 62.18    |
|                   |            | $\mathcal{M}_{2}$ | 72.22    | 35.93     | 52.20     | 76.28    | 62.12  | 65.27    | 60.67    |
|                   | 75.0%      | $\mathcal{M}_{1}$ | 70.79    | 34.92     | 53.62     | 76.88    | 60.54  | 65.03    | 61.31    |
|                   |            | $\mathcal{M}_{2}$ | 70.98    | 34.58     | 55.20     | 76.77    | 61.56  | 63.64    | 60.46    |
|                   | 62.5%      | $\mathcal{M}_{1}$ | 69.03    | 32.88     | 52.91     | 74.86    | 59.19  | 64.54    | 59.69    |
|                   |            | $\mathcal{M}_{2}$ | 69.22    | 37.29     | 55.73     | 75.22    | 59.35  | 64.21    | 60.17    |
|                   | 50.0%      | $\mathcal{M}_{1}$ | 66.34    | 32.88     | 53.09     | 74.97    | 58.25  | 62.49    | 58.59    |
|                   |            | $\mathcal{M}_{2}$ | 66.62    | 34.24     | 52.91     | 75.46    | 58.41  | 62.00    | 58.27    |
|                   | 37.5%      | $\mathcal{M}_{1}$ | 61.55    | 31.19     | 49.91     | 73.83    | 56.27  | 52.09    | 53.78    |
|                   |            | $\mathcal{M}_{2}$ | 62.38    | 32.20     | 50.26     | 73.34    | 56.67  | 55.28    | 55.02    |
|                   | 25.0%      | $\mathcal{M}_{1}$ | 50.91    | 26.10     | 44.97     | 68.39    | 52.72  | 38.33    | 46.38    |
|                   |            | $\mathcal{M}_{2}$ | 51.91    | 27.46     | 44.44     | 69.64    | 54.54  | 44.39    | 48.73    |

- [Inference speed @LgQe, @9USD]: We evaluate the inference speed of our LLM equipped with MatryoshkaKV. During the inference process with a batch size of 32, our current implementation consumes a slightly faster time than the baseline full-KV model.

Table 5: The inference speed of our MatryoshkaKV under the uniform compression rate during the inference process with a batch size of 32.
|      | LLaMA2-7B-base | 100% | 87.5% | 75% | 62.5% | 50% | 37.5% | 25% |
|------|------|------|------|------|------|------|------|------|
| Tokens per second | 33.65 | 34.12 | 34.08 | 34.90 | 35.27 | 36.42 | 36.75 | 37.22 |

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

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

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

[4]: Yuan et al., "ASVD: Activation-aware Singular Value Decomposition for Compressing Large Language Models", ArXiv 2024
</content>
</thread_object>

</thread>
<thread>
creation date:
1732115182254

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response (Part 3/4)

comment:
- [Combination with KIVI @76BC, @XwVW, @9USD]: In Appendix E, we add the results of the combination of our method with KIVI [3], also presented below:


Table 2: Results of Combination of Distilled MatryoshkaKV Projections and KIVI (2bit KV cache quantization) on Six Benchmarks. We use uniform compression levels for inference here for simplicity.
| **Model**         | **Budget** | **Method** | **HLSG** | **ARC-C** | **ARC-E** | **PIQA** | **WG**  | **CSQA** | **Avg.** |
|-------------------|------------|------------|----------|-----------|-----------|----------|---------|----------|----------|
| **LLaMA2 7B-base** | 100.0%     | MKV        | 70.89    | 36.95     | 53.26     | 76.39    | 61.56   | 67.08    | 60.98    |
|                   |            | MKV+KIVI   | 69.76    | 35.93     | 51.98     | 76.55    | 61.48   | 66.26    | 60.49    |
|                   | 87.5%      | MKV        | 70.87    | 36.95     | 51.15     | 76.17    | 61.80   | 64.95    | 60.47    |
|                   |            | MKV+KIVI   | 70.45    | 36.61     | 50.79     | 76.44    | 61.01   | 64.13    | 59.91    |
|                   | 75.0%      | MKV        | 69.30    | 33.90     | 54.67     | 75.90    | 61.09   | 63.23    | 60.08    |
|                   |            | MKV+KIVI   | 68.62    | 32.20     | 54.85     | 76.06    | 60.30   | 63.23    | 59.68    |
|                   | 62.5%      | MKV        | 67.25    | 36.27     | 53.62     | 75.52    | 59.27   | 64.46    | 59.39    |
|                   |            | MKV+KIVI   | 66.56    | 35.25     | 51.68     | 75.41    | 59.43   | 61.59    | 58.33    |
|                   | 50.0%      | MKV        | 65.08    | 33.56     | 52.03     | 74.81    | 57.54   | 60.36    | 56.98    |
|                   |            | MKV+KIVI   | 63.25    | 32.54     | 51.15     | 74.43    | 57.38   | 59.46    | 56.35    |
|                   | 37.5%      | MKV        | 61.02    | 29.83     | 49.21     | 73.45    | 55.64   | 55.36    | 54.09    |
|                   |            | MKV+KIVI   | 57.11    | 28.81     | 48.85     | 71.71    | 55.64   | 50.37    | 52.08    |
|                   | 25.0%      | MKV        | 50.61    | 25.76     | 45.33     | 69.64    | 54.30   | 43.90    | 47.96    |
|                   |            | MKV+KIVI   | 48.12    | 27.80     | 42.86     | 67.85    | 53.59   | 40.54    | 46.78    |

The results show that our MatryohskaKV can be easily combined with KV quantization techniques and achieve a higher compression rate.

- [Baseline @LgQe, @9USD]: In Appendix F, we add some baselines such as ASVD [4] and compare with our MatryoshkaKV， with the results listed below:

Table 3: Comparison between our MatryoshkaKV and baseline ASVD. We use uniform compression levels for inference here for simplicity.
| **Model**              | **Budget** | **Method** | **HLSG** | **ARC-C** | **ARC-E** | **PIQA** | **WG** | **CSQA** | **Avg.** |
|-----------------------|------------|------------|----------|-----------|-----------|----------|--------|----------|----------|
| LLaMA2                | 100.0%     | baseline   | 74.00    | 35.93     | 50.97     | 78.50    | 61.64  | 65.93    | 61.16    |
|                       | 95%        | ASVD       | 71.12    | 36.95     | 52.20     | 76.28    | 62.35  | 66.67    | 60.92    |
|                       |            | MKV        | 72.59    | 36.27     | 53.09     | 76.44    | 62.43  | 66.75    | 61.25    |
|                       | 90%        | ASVD       | 70.45    | 34.92     | 52.03     | 75.63    | 61.72  | 64.70    | 60.06    |
|                       |            | MKV        | 72.30    | 36.61     | 54.50     | 76.50    | 62.90  | 65.93    | 62.03    |
|                       | 85%        | ASVD       | 67.23    | 35.93     | 50.26     | 74.86    | 60.38  | 62.16    | 59.29    |
|                       |            | MKV        | 72.33    | 35.93     | 53.26     | 76.33    | 61.80  | 64.78    | 61.13    |
</content>
</thread_object>

</thread>
<thread>
creation date:
1732115211188

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response (Part 2/4)

comment:
**Simultaneously, we add some experiments requested by the reviewers:**
- [Combination with H2O @76BC, @uec1, @XwVW, @9USD]: In Appendix E, we add the results of the combination of our method with H20 [2] and show perplexity on LongBench, which are also listed below.

Table 1: Results of Combination of Distilled MatryoshkaKV Projections and H2O across Seven Benchmarks. We use uniform compression levels for inference here for simplicity. The first and second columns indicate the individual compression rates along two axes. If H2O uses 20% cache on the sequence length axis and MatryoshkaKV uses 50% cache on the feature dimension axis, the overall cache utilization is 10%.
| **H₂O** | **MKV** | **LongBench** | **HLSG** | **ARC-C** | **ARC-E** | **PIQA** | **WG** | **CSQA** | **Avg.** |
|---------|---------|---------------|----------|-----------|-----------|----------|--------|----------|----------|
| 100%    | 100%    | 4.17          | 72.05    | 37.29     | 52.38     | 76.66    | 61.72  | 67.32    | 61.24    |
|         | 87.5%   | 4.44          | 72.22    | 35.93     | 52.20     | 76.28    | 62.12  | 65.27    | 60.67    |
|         | 75.0%   | 4.57          | 70.98    | 34.58     | 55.20     | 76.77    | 61.56  | 63.64    | 60.46    |
|         | 62.5%   | 4.70          | 69.22    | 37.29     | 55.73     | 75.22    | 59.35  | 64.21    | 60.17    |
|         | 50.0%   | 4.93          | 66.62    | 34.24     | 52.91     | 75.46    | 58.41  | 62.00    | 58.27    |
|         | 37.5%   | 5.47          | 62.38    | 32.20     | 50.26     | 73.34    | 56.67  | 55.28    | 55.02    |
|         | 25.0%   | 7.66          | 51.91    | 27.46     | 44.44     | 69.64    | 54.54  | 44.39    | 48.73    |
| 75%     | 100%    | 4.18          | 70.71    | 36.61     | 52.38     | 76.55    | 60.54  | 66.50    | 60.55    |
|         | 87.5%   | 4.44          | 71.42    | 35.25     | 53.09     | 76.33    | 59.91  | 64.62    | 60.74    |
|         | 75.0%   | 4.57          | 70.31    | 34.34     | 54.14     | 76.39    | 59.27  | 62.90    | 59.94    |
|         | 62.5%   | 4.70          | 68.47    | 36.27     | 54.32     | 75.41    | 58.48  | 63.96    | 59.89    |
|         | 50.0%   | 4.94          | 66.00    | 32.54     | 51.50     | 75.63    | 57.30  | 61.43    | 57.46    |
|         | 37.5%   | 5.47          | 61.50    | 32.88     | 49.21     | 73.01    | 55.09  | 55.12    | 54.63    |
|         | 25.0%   | 7.67          | 51.32    | 27.80     | 44.09     | 69.37    | 53.59  | 44.55    | 48.47    |
| 50%     | 100%    | 4.20          | 68.72    | 33.22     | 52.20     | 76.12    | 56.67  | 64.78    | 58.62    |
|         | 87.5%   | 4.46          | 67.89    | 34.58     | 51.85     | 76.28    | 55.88  | 62.00    | 58.13    |
|         | 75.0%   | 4.59          | 66.01    | 35.59     | 53.79     | 75.41    | 54.54  | 62.00    | 58.05    |
|         | 62.5%   | 4.73          | 63.59    | 34.92     | 51.32     | 75.68    | 55.25  | 60.52    | 57.04    |
|         | 50.0%   | 4.96          | 61.33    | 36.10     | 50.74     | 73.67    | 55.57  | 57.67    | 55.85    |
|         | 37.5%   | 5.50          | 59.26    | 29.83     | 49.91     | 73.61    | 53.04  | 54.14    | 53.29    |
|         | 25.0%   | 7.71          | 49.44    | 26.44     | 41.80     | 68.72    | 52.96  | 43.24    | 46.94    |
| 20%     | 100%    | 4.40          | 61.55    | 25.76     | 41.27     | 73.29    | 53.28  | 47.01    | 49.98    |
|         | 87.5%   | 4.65          | 61.36    | 30.51     | 39.86     | 73.72    | 52.09  | 49.06    | 50.94    |
|         | 75.0%   | 4.79          | 60.29    | 28.47     | 38.62     | 72.75    | 53.12  | 50.45    | 50.62    |
|         | 62.5%   | 4.93          | 58.77    | 26.78     | 39.86     | 70.84    | 52.72  | 49.30    | 50.58    |
|         | 50.0%   | 5.19          | 56.39    | 26.78     | 38.10     | 71.22    | 51.62  | 49.16    | 49.66    |
|         | 37.5%   | 5.74          | 52.12    | 23.39     | 34.22     | 68.50    | 52.17  | 41.44    | 44.82    |
|         | 25.0%   | 8.01          | 43.22    | 21.02     | 31.92     | 63.93    | 51.38  | 33.09    | 40.75    |


According to the results, by concurrently using MatryoshkaKV and 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732115237719

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response (Part 1/4)

comment:
We thank the reviewers for their thoughtful feedback. We are encouraged that the reviewers found our MatryoshkaKV has comprehensive evaluations (@LgQe, @76BC) and robust performance (@LgQe, @XwVW, @uec1, @9USD), being flexible (@LgQe, @XwVW, @76BC) and novel for KV cache compression (@uec1, @76BC).
**We have made a revision of our previous paper and highlighted the modified parts in red**. The following are the modifications we have made to our paper:

In the experiment part of Section 5, specifically in the subsection of CPT, **we add numerous experimental details**, including:

- [Clarification @9USD]: The training dataset used in CPT and its source.
- [Consumption @LgQe]: The runtime metrics for the experimental part.
- [Clarification @LgQe]: The $\Delta r$ used in the greedy search for adaptive compression rates before the inference stage.

**We provide a more detailed explanation of our method:**
- [Clarification @LgQe]: In Section 4, we clarify the whole pipeline of our method.
- [Analysis @uec1]: In Appendix A, we refine some mathematical formulae.
</content>
</thread_object>

</thread>


Metareview:
metareview:
Summary:
MatryoshkaKV introduces a novel method for compressing Key-Value (KV) cache in large language models along the feature dimension using trainable orthogonal projection matrices. The method combines PCA initialization with knowledge distillation and a Matryoshka training strategy to enable adaptive compression rates across different layers and attention heads. The approach achieves significant compression (up to 95%) while maintaining over 90% of the original model's accuracy.

Main Strengths:

- Novel approach to KV cache compression focusing on the previously unexplored feature dimension
- Strong empirical results showing significant compression with minimal performance degradation
- Compatibility with other compression techniques (demonstrated with GQA, H2O, and KIVI)
- Practical implementation requiring minimal hyperparameter tuning
- Comprehensive ablation studies validating key components

Main Weaknesses:

- Limited theoretical analysis explaining why the method outperforms PCA-based approaches
- Runtime evaluation could be more comprehensive
- Current implementation of heterogeneous ranks may not provide runtime benefits without system-level optimizations
- Task-specific calibration requirement for optimal compression rates