MatryoshkaKV: Adaptive KV Compression via Trainable Orthogonal Projection

## Abstract

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 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 observe the issue with PCA projection where significant performance degradation is observed at low compression rates.
To bridge the gap, we propose to directly tune the orthogonal projection matrices with a distillation objective using an elaborate Matryoshka training strategy.
After training, we adaptively search for the optimal compression rates for various layers and heads given varying compression budgets.
Compared to previous works, our method can easily embrace pre-trained LLMs and hold a smooth tradeoff between performance and compression rate.
We empirically 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-7B-base and Mistral-7B-v0.3-base.

## 1 Introduction

Large language models (LLMs) like GPT-4 (OpenAI et al., ) and Claude3 (Enis & Hopkins, ) have shown great promise, finding applications in areas such as text generation (Brown et al., ; Raffel et al., ), code completion (Rozière et al., ), and sentiment analysis (Zhang et al., ).
The Key-Value (KV) cache, which is introduced to cache historical information for self-attention, is essential for maintaining context and accelerating the inference of LLMs.
However, as the size of the model and data continues to grow (Fu et al., ; Ding et al., ; Chen et al., ), the KV cache can swiftly lead to system bottleneck in terms of storage and memory transfer (Shi et al., ).

Considerable efforts have been devoted to addressing such an issue.
Noting that the KV cache contains tensors of shape (layer number, head number, sequence length, feature dimension), existing works have investigated compressing the KV cache from the axes of layer number (Brandon et al., ; Sun et al., ; Goldstein et al., ), head number (Ainslie et al., ; Shazeer, ; Yu et al., ), and sequence length (Wang et al., ; Zhang et al., ; Li et al., ; Xiao et al., ).
Conversely, the exploration of feature dimension for KV cache compression significantly lags behind, partially because of the inherent difficulties of modifying a well-structured feature space.

This paper aims to tackle this with the help of curated low-rank projection matrices, e.g., both the query and key are projected into the same lower-dimensional space wherein the inner product closely approximates that in the original space.
We first identify the necessity to guarantee the orthogonality among the rows of such matrices, and hence attempt to take the principal components of the keys or values in each layer to instantiate the projections, given the prevalence of Principal Component Analysis (PCA) for data compression.
We observe that such projections can be seamlessly plugged into pre-trained LLMs while retaining reliable generation quality at a moderate compression level.
Compared to the low-rank architectures of Multi-head Latent Attention (MLA) (DeepSeek-AI et al., ), the PCA strategy is more approachable due to its training-free nature and also advocated by Saxena et al. ().
Yet, we note that the PCA projections suffer from quickly degraded performance when further increasing the compression level.
This is because, while the principal components are optimal for recovering the keys or values in each individual layer, they may be suboptimal for preserving the global outputs due to the non-linearity and compounding effects in LLM.

To bridge the gap, we propose to jointly adjust all orthogonal projection matrices incorporated into the model with a knowledge distillation objective, enforcing the model output based on the projected keys and values to remain close to the original one.
The orthogonality constraint upon the projection matrices is consistently enforced by a Cayley parameterization.
Besides, we desire a hierarchy over the columns of the projection matrices—as in PCA—so that we can smoothly trade-off between compression level and performance.
To this end, we introduce a Matryoshka training strategy—compute the model output based on the first $r$ columns of the matrices, where $r$ is randomly sampled from a predefined schedule such as $\{4,8,16,...\}$, and ensure its closeness to the original output.
In practice, we sample various $k$ for different layers, heads, and keys/values during training to disentangle the projections in the model.
Doing so enables the search for heterogeneous compression rates for different projection matrices during inference and we develop a greedy algorithm for this.
Heterogeneous compression rates are displayed in Figure 1

Figure: Figure 1: Visualization of the feasible compression level for the keys and values in our MatryoshkaKV model distilled from the LLaMA2-7B-base model. We individually leverage samples in ARC-challenge (ARC-C), ARC-easy (ARC-E) (Clark et al., ), and Winogrande (WG) (Sakaguchi et al., ) to determine the compression level. Lighter colors indicate higher compression levels. As shown, our approach enables the use of various compression strategies for various tasks.
Refer to caption: x1.png

Experiments on both continual pre-training (CPT) and supervised fine-tuning (SFT) exhibit the efficacy of our MatryoshkaKV approach.
For the former, we opt to experiment on LLaMA2-7B-base (Touvron et al., ) with the RedPajama dataset (Computer, ).
To demonstrate compatibility with Group Query Attention (GQA) (Ainslie et al., ), we also apply our approach to the Mistral-v0.3-7B-base (Jiang et al., ) model.
We observe that after rarely processing 2 million training tokens, MatryoshkaKV achieves a 37.5% compression rate while retaining over 90% of the original model’s accuracy.
In the SFT experiments, we train both MatryoshkaKV and LoRA (Hu et al., ) on downstream tasks including OBQA (Mihaylov et al., ), GSM8K (Cobbe et al., ), etc.
The results show that our MatryoshkaKV can utilize less than 40% cache while still achieving over 90% accuracy derived from full cache utilization.
We also perform extensive ablation studies to chase a deep understanding of our approach.

## 2 Related Work

KV Cache Eviction and Merging.
KVMerger (Wang et al., ) and PyramidKV (Cai. et al., ) introduce innovative approaches to reduce KV cache memory consumption along sequence length dimension in long-context tasks.
KVMerger merges KV by Gaussian weights and attention score, while PyramidKV uses a layer-wise approach with recent tokens occupying more weights. CLA (Brandon et al., ), YOCO (Sun et al., ), and GoldFinch (Goldstein et al., ), among others, exploit inter-layer KV cache reuse by sharing KV heads across layers.
This significantly reduces the KV cache size along the head number dimension without compromising model capacity.
GQA (Ainslie et al., ), MQA (Shazeer, ), and HeadKV (Yu et al., ), especially the last one, have demonstrated the effectiveness of compressing KV cache on the axis of head number due to their low-rank properties.

KV cache Compression along the hidden size dimension.
DeepSeekv2 (DeepSeek-AI et al., ) employs MLA techniques to reduce the feature dimension of keys and values within the attention mechanism, but this requires costly retraining from scratch.
Concurrent advancements, however, have addressed this limitation.
InfiniGen (Lee et al., ), Eigen-Attention (Saxena et al., ) and HeadKV (Yu et al., ) achieve a 40% reduction in the KV cache sizes using orthogonal projections parameterized by the SVD of the Q, K, and V matrices derived from a subset of samples.
To mitigate performance degradation, LoRA (Hu et al., ) is employed to fine-tune model parameters.
However, this compression approach on the axis of feature dimension results in a sharp decline in model performance when using less than 60% cache budget.
Furthermore, fine-tuning the base model with LoRA may lead to catastrophic forgetting.
In this paper, our method MatryoshkaKV circumvents these risks and achieves higher compression rate by directly fine-tuning orthogonal projections.

## 3 Preliminary

This section provides a review of the KV cache mechanism and elucidates the implementation of PCA projection for KV cache compression.

### 3.1 KV Cache

Consider the inference of an LLM $p(\bm{\cdot}|{\bm{x}})$ with ${\bm{x}}$ as the prompt.
It is a common practice to deploy the KV cache technique to each self-attention head in the model to store the key and value states for the present context, including both the prompt ${\bm{x}}$ and the tokens that have already been generated.
Given the KV cache for the context of length $L-1$ and dimension $d$ in each head,
the model generates a subsequent new token $y$ with the attention states $\mathrm{softmax}({QK^{\top}}/{\sqrt{d}})V$, where $Q\in\mathbb{R}^{1\times d}$ is the query vector for $y$ and $K,V\in\mathbb{R}^{L\times d}$ denote the concatenation of the KV cache and the KV vectors for $y$.
This way, the computational complexity for one decoding step is reduced from $\mathcal{O}(L)$ to $\mathcal{O}(1)$.

However, the size of the KV cache can grow quickly w.r.t. that of the model and context, often causing system bottlenecks in terms of both storage and memory transfer during the inference phase.
To address this, various KV cache compression techniques have been proposed, e.g., sharing the KV headers across layers inside LLMs (Brandon et al., ; Sun et al., ; Goldstein et al., ), merging heads that require caching KV (Yu et al., ; Ainslie et al., ), evicting or merging redundant tokens (Xiao et al., ; Li et al., ; Cai. et al., ; Zhang et al., ).
This work alternatively focuses on compressing the feature dimension $d$ of the KV cache, exploring a novel axis for KV cache compression that is compatible with existing methodologies.

### 3.2 Traing-free Dimension Reduction Via PCA

A simple way to reduce the dimension of the KV cache is finding some matrices to project $K,V$ as $K^{\prime},V^{\prime}\in\mathbb{R}^{L\times r}$, $(r<d)$.
Then, we can only cache $K^{\prime}$ and $V^{\prime}$, reducing the storage and memory transfer cost from $\mathcal{O}(d)$ to $\mathcal{O}(r)$.
The rank $r$ is desired to be adjustable based on the available compression budget: when the budget is sufficient, caching full KV states helps prevent information loss; in cases of limited budget, caching only the most essential information should be feasible.
To fulfill this, it is reasonable to introduce full-rank projection matrices $U\in\mathbb{R}^{d\times d}$ and demand a hierarchy over the columns of $U$ so that the optimal $r$-rank cache can result from the first $r$ columns of $U$, denoted as $U_{r}\in\mathbb{R}^{d\times r}$.
In practice, $U$ should be distinct for $K$ and $V$ and vary across attention heads and layers within the model, as these states commonly exhibit diverse distributions.

During the forward pass of the model, we should be able to recover the original $K$ and $V$ from the reduced $K^{\prime}$ and $V^{\prime}$.
A natural choice is using $U^{\top}$, the transposition of the projection matrices, where
$U_{r}U_{r}^{\top}\approx I$ needs to be satisfied.
Given that $r$ can vary from $1$ to $d$, we identify that $U$ should be orthogonal matrices.
It is known that the optimal orthogonal projections for compressing a set of high-dimension vectors can be their principal components, so we suggest constructing $U$ based on the PCA results of the key or value states of a long sequence of tokens for each head separately.

Table 1 displays an empirical study of the efficacy of such training-free projections.
As shown, PCA projections exhibit reliable performance at moderate levels of compression budget.
This is remarkable because the PCA strategy does not need costly from-scratch training of the projection matrices, in sharp contrast to the projection mechanisms used by MLA (DeepSeek-AI et al., ).
We note that PCA projection is also advocated by Saxena et al. (); refer to Appendix B for the difference between our attempts and theirs regarding applying projections before or after RoPE (Su et al., ) and whether performing fine-tuning.

Nevertheless, as the table displays, the PCA projections suffer from quickly degraded performance when further increasing the compression level.
This is because despite principal components being optimal for key or value recovery in individual head layers, they may be inadequate for preserving the final output due to the non-linearity and compounding effects of the attention mechanism.

## 4 Methodology

To address the aforementioned issue, we propose to jointly tune the orthogonal projection matrices introduced to the LLM under an elaborate objective, to realize a more robust KV cache compression.
This section provides the training and inference details of our approach.

Figure: Figure 2: Vanilla KV cache vs. the proposed MatryoshkaKV. In particular, we introduce orthogonal projection matrices to reduce the dimension of stored keys and values. We explicitly enforce a hierarchy over the columns of projection matrices so as to concentrate the principal information on the head dimensions and enable the adjustment of compression level according to resource constraints.
Refer to caption: x2.png

### 4.1 Minimize Compression Loss by Knowledge Distillation

Recalling the objective for the compression is that the model outputs based on the compressed states should stay close to the original one.
This implies a knowledge distillation objective (Hinton et al., ), which can be instantiated with the KL divergence:

$$ $\mathcal{L}_{KD}=\displaystyle D_{\mathrm{KL}}(p\left(\bm{\cdot}|{\bm{x}} \right)\|p^{\prime}\left(\bm{\cdot}|{\bm{x}}\right))$ (1) $$

where we abuse $p^{\prime}$ to refer to the LLM equipped with low-rank projection matrices.
As suggested by the literature (Kou et al., ), we also incorporate a language modeling loss to $p^{\prime}$ to prevent the generated text from deviating from the context distribution of the dataset, thereby ensuring high-quality generation.
The tuning process involves only the update of $U$, which ensures that the model performance under the full-rank KV cache is maintained.

Orthogonal constraint.
We initialize the trainable orthogonal projections with the PCA ones due to their effectiveness.
To confine the evolution of the projection matrices within the orthogonal matrix family throughout the tuning process, we employ Cayley parameterization to formulate the orthogonal matrix.
Specifically, there is $U=\left(I+Q\right)\left(I-Q\right)^{-1}$ with $Q$ as a skew-symmetric trainable matrix of size $d\times d$.
Considering that $d$ is usually small (e.g., $64$ or $128$), the complexity of performing such an orthogonal transformation during training is minimal.

### 4.2 Acquire Hierarchical KV Cache by MatryoshkaKV Training

The tuning process can destroy the hierarchical structures present in the orthogonal matrices inherited from the PCA ones because there is no prioritization given to the columns of the matrices $U$ from the training objective.
Consequently, we lose the flexibility to achieve a smooth transition between the level of compression and maintenance of the original performance.

To tackle this challenge, we draw inspiration for Matryoshka representation learning (Kusupati et al., ), introducing a Matryoshka strategy for training the projection matrices $U$.
In particular, for each training iteration, we randomly sample $r$ from a predefined schedule such as $\{4,8,16,...,d/4,d/2,d\}$ and use the first $r$ columns of $U$, i.e., $U_{r}$, to construct the model $p^{\prime}$ for training.
Note that the keys and values at different heads and layers use separately sampled $r$ to avoid the entanglement of the compression effect.
An illustrative explanation of this is given in Figure 2, and our approach is then called MatryoshkaKV for short.

### 4.3 Find Heterogeneous Compression Rates for Various Layers & Heads

The Matryoshka training strategy enables the search for heterogeneous compression rates for various layers and heads in the model given a specific compression budget.
Basically, we can first propose a compression level for the projection matrix at a particular position, assessing the deviation of the model output from the original on a predefined calibration dataset (measured by KL divergence), and determining whether to accept the proposal based on a predefined tolerance threshold for the deviation.
Algorithm 1 exhibits a greedy algorithm for accelerating this based on accepting proposals in parallel.
Note that this greedy algorithm also applies to the PCA projections.

Figure: Algorithm 1 Greedy search for adaptive compression levels in our efficient LLM.

Discussion.
The recent KV cache compression approach on sequence length aspect (Cai. et al., ) also observes that compared to uniformly compressed KV cache using the same rate across all layers (Li et al., ), employing a distinct compression rate for each layer results in improved information utilization.
Furthermore, as observed in (Wu et al., ), certain retrieval heads within an LLM consistently attend to crucial information, regardless of contextual variations.
The indiscriminate compression rates of these heads can lead to significant performance degradation.
These both support the necessity of the proposed heterogeneous KV cache compression approach.

## 5 Experiments

In this section, we conduct experiments on continual pre-training (CPT) and supervised fine-tuning (SFT) scenarios to demonstrate that our MatryoshkaKV can not only preserve the foundation knowledge of a base model but also be compatible with LoRA (Hu et al., ) for downstream tasks.
Ablation studies in Section 5.3 validate the efficacy of our proposed method.

### 5.1 Continual Pre-training

Setup.
We select LLaMA2-7B-base (Touvron et al., ) and Mistral-v0.3-7B-base (Jiang et al., ) as our base models.
We adopt the Matryoshka training strategy detailed in Section 4.2 and fine-tune MatryoshkaKV projections with knowledge distillation loss in Equation 1 and language modeling loss,
applying a 1:3 weighting ratio between the two losses.
The projection rank $r_{k}$ and $r_{v}$ are randomly sampled from a predefined schedule set $\{\frac{i}{8}d\}_{i=1}^{8}$ during training and are chosen dynamically with the greedy search for adaptive compression levels, as detailed in Section 4.3 during inference.
We use Opencompass (Contributors, ) to test performance on several widely-used zero-shot benchmarks: PIQA (Bisk et al., ), ARC-challenge (ARC-C) (Clark et al., ), ARC-easy (ARC-E) (Clark et al., ), WinoGrande (WG) (Sakaguchi et al., ), HellaSwag (HLSG) (Zellers et al., ), and CommonSenseQA (CSQA) (Talmor et al., ).

Results.
Table 1 presents the results of our experiments.
In zero-shot tasks, our MatryoshkaKV cache substantially reduces the cache footprint with minimal impact on performance.
Specifically, our method retains 93.10% of LLaMA2-7B-base’s average accuracy and 92.63% of Mistral-v0.3-7B-base’s average accuracy, while utilizing only 37.5% of the original cache size.
For simpler tasks like PIQA, it achieves 88.71% and 92.00% of the base model’s performance with just a 25% cache budget.
On more challenging tasks such as ARC-C, a larger cache budget is required, with 50% needed to retain 90% of the base model’s performance.
By contrast, PCA projection shows a sharp performance drop when the cache budget is reduced below 62.5%, achieving just 70.42% accuracy of LLaMA2-7B-base and 52.86% of Mistral-v0.3-7B-base.
These results underscore the superior efficiency of our approach compared with PCA.
We attribute PCA’s performance decline to suboptimal projection matrices, whereas our method maintains closer alignment with the base model, thereby mitigating this degradation.

**Table 1: Comparison between our MatryoshkaKV method (donated as MKV in the table) and PCA projection. We use LLaMA2-7B-base and Mistral-v0.3-7B-base as our source models, and their performance is used as a baseline. Accuracy on HellaSwag, ARC-challenge, ARC-easy, PIQA, WinoGrande, and CommonSenseQA is reported, with higher scores indicating superior performance, at seven KV cache budgets. At the same budget, the higher average accuracy is underlined.**
| Model | Budget | Method | HLSG | ARCC | ARCE | PIQA | WG | CSQA | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| LLaMA2<br>7B-base | 100.0% | Baseline | 74.00 | 35.93 | 50.97 | 78.50 | 61.64 | 65.93 | 61.16 |
| PCA | 72.04 | 36.95 | 52.38 | 76.66 | 61.72 | 67.24 | 61.17 |  |  |
| MKV | 72.05 | 37.29 | 52.38 | 76.66 | 61.72 | 67.32 | 61.24 |  |  |
| 87.5% | PCA | 71.91 | 35.93 | 53.97 | 76.66 | 61.40 | 67.65 | 61.25 |  |
| MKV | 71.58 | 37.97 | 53.26 | 75.95 | 62.12 | 69.57 | 61.74 |  |  |
| 75.0% | PCA | 70.99 | 35.59 | 54.14 | 76.22 | 60.06 | 66.99 | 60.67 |  |
| MKV | 71.58 | 38.31 | 55.56 | 76.01 | 61.09 | 66.75 | 61.55 |  |  |
| 62.5% | PCA | 67.16 | 34.24 | 54.85 | 74.76 | 57.77 | 61.10 | 58.31 |  |
| MKV | 68.03 | 37.97 | 56.08 | 75.12 | 60.30 | 65.44 | 60.49 |  |  |
| 50.0% | PCA | 42.11 | 29.83 | 35.10 | 58.16 | 52.57 | 40.62 | 43.07 |  |
| MKV | 66.78 | 36.61 | 55.91 | 74.32 | 59.12 | 61.92 | 59.11 |  |  |
| 37.5% | PCA | 24.24 | 26.44 | 26.63 | 51.25 | 50.36 | 19.90 | 33.14 |  |
| MKV | 63.97 | 33.90 | 51.68 | 74.97 | 57.92 | 59.21 | 56.94 |  |  |
| 25.0% | PCA | 23.98 | 29.49 | 26.28 | 51.20 | 50.36 | 16.22 | 32.92 |  |
| MKV | 51.91 | 27.46 | 44.44 | 69.64 | 54.54 | 44.39 | 48.73 |  |  |
| Mistral-v0.3<br>7B-base | 100.0% | Baseline | 75.50 | 42.03 | 63.14 | 80.25 | 65.43 | 70.68 | 66.17 |
| PCA | 75.46 | 42.03 | 62.96 | 80.25 | 65.35 | 70.27 | 66.05 |  |  |
| MKV | 75.44 | 42.03 | 62.96 | 80.25 | 65.51 | 70.27 | 66.08 |  |  |
| 87.5% | PCA | 73.46 | 42.71 | 63.32 | 79.54 | 63.93 | 70.76 | 65.92 |  |
| MKV | 75.63 | 42.03 | 64.37 | 79.71 | 65.51 | 70.35 | 66.27 |  |  |
| 75.0% | PCA | 70.75 | 37.63 | 61.73 | 78.18 | 62.59 | 68.47 | 63.23 |  |
| MKV | 75.29 | 43.39 | 63.14 | 79.54 | 64.96 | 69.12 | 65.90 |  |  |
| 62.5% | PCA | 63.48 | 34.24 | 55.73 | 75.90 | 60.77 | 62.24 | 58.73 |  |
| MKV | 74.23 | 40.34 | 62.96 | 79.33 | 64.25 | 68.63 | 64.96 |  |  |
| 50.0% | PCA | 28.12 | 22.71 | 28.40 | 58.16 | 49.64 | 22.85 | 34.98 |  |
| MKV | 73.32 | 38.98 | 62.08 | 79.16 | 61.88 | 67.08 | 63.75 |  |  |
| 37.5% | PCA | 25.04 | 22.03 | 28.04 | 53.86 | 49.25 | 21.21 | 33.24 |  |
| MKV | 70.40 | 35.93 | 58.91 | 77.91 | 60.30 | 64.29 | 61.29 |  |  |
| 25.0% | PCA | 24.91 | 26.10 | 25.40 | 52.67 | 48.30 | 19.74 | 32.85 |  |
| MKV | 59.21 | 25.42 | 48.68 | 73.83 | 54.30 | 45.13 | 51.10 |  |  |

### 5.2 Supervised Fine-tuning

Setup.
We use LLaMA2-7B-base (Touvron et al., ) as our base model and verify the efficacy of our method on PIQA (Bisk et al., ), GSM8K (Cobbe et al., ), HellaSwag (Zellers et al., ), and OpenbookQA (OBQA) (Mihaylov et al., ) datasets.
We design a two-stage training strategy to make Matryoshka training strategy compatible with LoRA (Hu et al., ) fine-tuning.
Specifically, LoRA is firstly used to adapt the base model to downstream tasks, following standard SFT practices (Naveed et al., ; Zhao et al., ).
In the second stage, we jointly fine-tune the MatryoshkaKV projections with the Matryoshka training strategy and the LoRA parameters.
Further discussion on the superiority of this recipe is detailed in Appendix C.

**Table 2: Accuracy of our Matryoshka method after SFT based on LLaMA2-7B-base on four downstream tasks: PIQA, GSM8K, HellaSwag, and OpenbookQA, at seven KV cache budgets. Degradation from baseline is shown in brackets.**
| Model | Budget | PIQA | GSM8K | HLSG | OBQA | Avg. |
| --- | --- | --- | --- | --- | --- | --- |
| LLaMA2<br>7B-base | 100.0 % | 84.22 | 34.95 | 93.94 | 83.2 | 74.08 (-0.00%) |
| 87.5 % | 83.84 | 35.25 | 93.17 | 81.80 | 73.52 (-0.76%) |  |
| 75.0 % | 83.30 | 32.90 | 91.47 | 81.40 | 72.27 (-2.44%) |  |
| 62.5 % | 82.75 | 31.46 | 89.86 | 79.60 | 70.92 (-4.27%) |  |
| 50.0 % | 79.33 | 31.77 | 86.29 | 76.60 | 68.50 (-7.53%) |  |
| 37.5 % | 75.35 | 26.91 | 76.10 | 70.80 | 62.29 (-15.9%) |  |
| 25.0 % | 69.04 | 16.38 | 56.10 | 61.40 | 50.73 (-31.5%) |  |

Figure: Figure 3: Evaluation loss of four budgets vs. the number of training samples during 1 epoch of SFT on GSM8K (Left). Evaluation loss of models with and without PCA initialization, using a 50% cache budget, vs. the number of training samples during 4 epochs of SFT on GSM8K (Right).
Refer to caption: x3.png

Results.
We report accuracy on four zero-shot benchmarks at seven KV cache budgets in Table 2.
As shown, our method demonstrates notable performance in the SFT scenario.
It achieves 92.47% of the baseline’s average accuracy while utilizing only 50% of the KV cache budget.
On simple tasks like PIQA, our method retains 89.47% of the full-cache performance with a 37.5% cache budget.
However, for more complex tasks such as GSM8K, a 50% cache budget is necessary to achieve comparable results.
Furthermore, we report the evaluation loss at four budgets: 100%, 62.5%, 50%, and 37.5% during the second stage of SFT on GSM8K in Figure 3 (Left).
It shows our method simultaneously optimizes models under various KV cache budgets and maintains the hierarchical structures present in the orthogonal matrices.
These findings highlight the robustness of our approach, delivering consistent performance across both CPT and SFT scenarios.

### 5.3 Ablation Studies

We conduct ablation studies on various components of our method to verify their effectiveness.

Figure: Figure 4: Comparison between PCA and distilled MatryoshkaKV Projections after CPT with and without greedy search for adaptive compression levels. We report average accuracy on datasets mentioned in the experimental setup of Section 5.1 (Left). Comparison between with and without Matryoshka training strategy and orthogonal constraint after SFT on GSM8K. We report the relative accuracy compared with the LLaMA2-7B-base model fine-tuned with LoRA on GSM8K, utilizing the full KV cache (Right).
Refer to caption: x5.png

W/o greedy search for adaptive compression levels.
We evaluate our trained models without our greedy search for adaptive compression levels.
Figure 4 (Left) represents the average accuracy on four datasets mentioned in Section 5.2 as the cache budget varies.
For the exact numerical values, please refer to Table 3 in Appendix D.
To ensure that each head in the LLM plays its due role, we set 25% as our minimum cache budget for each head.
At a 37.5% cache budget, the average accuracy improves by 1.92%, indicating the significance of our search algorithm for further KV cache compression.
Furthermore, our MatryoshkaKV demonstrates robustness even when applying a uniform compression rate across all layers and heads, in contrast to PCA projection, which fails to handle this setting effectively.

W/o Matryoshka training strategy.
As discussed in Section 4.2, we point out that the tuning process w/o Matryoshka training strategy can destroy the hierarchical structures present in the orthogonal matrices inherited from the PCA ones.
To validate this, we train MatryoshkaKV projections with a fixed KV cache budget of 50%.
The result is displayed in Figure 4 (Right).
We observe that fixing the compression rate at 50% hinders the potential for further compression.
Moreover, when the budget exceeds 50%, the model’s performance does not improve significantly but even deteriorates, indicating the hierarchical structure of projections is destroyed.

W/o orthogonal constraint.
We investigate the necessity of imposing the orthogonal constraint during training, with experimental results presented by Figure 4 (Right).
After training without orthogonal constraint on GSM8K, we observe that non-orthogonal projections achieve performance comparable to orthogonal projections when the cache budget is less than 50%.
However, when utilizing a full KV cache budget, this model is unable to maintain the performance of the base model.
This is due to the non-orthogonality of the projection matrix, which prevents LLM from replicating the attention mechanism of the base model.
This phenomenon also validates our discussion in previous Section 3.2.

W/o PCA initialization.
To demonstrate the necessity of using PCA results to initialize projections, we train an LLM equipped with randomly initialized orthogonal matrices on GSM8K and impose orthogonal constraints.
In Figure 3 (Right), we report the evaluation loss during the second stage of SFT on GSM8K.
Despite training for four epochs, randomly initialized orthogonal projections consistently fail to converge to an optimal solution, and the text generated by our fine-tuned LLM projection is composed of meaningless symbols.
This highlights the critical importance of PCA initialization.

### 5.4 Heterogeneous Compression Rates Visualization

Figure 1 shows the heterogeneous compression levels across all attention heads inside our MatryoshkaKV LLM distilled from the LLaMA2-7B-base.
We acquire these results by leveraging the greedy search for adaptive compression levels on the ARC-C, ARC-E, and WinoGrande datasets.
We observe that shallower layers require larger KV cache budgets, while in deeper layers, only a minority of specific heads require a relatively high budget.
PyramidKV (Cai. et al., ) also observes that the model aggregates information globally from all available content in lower layers, indicating that KV cache inside lower layers can exert a substantial influence over the final output and should be allocated at a relatively high budget.
Therefore, allocating more cache in lower layers and less in higher ones is superior to maintaining a uniform KV cache size across layers.
Also, as Wu et al. () point out, retrieval heads with high retrieval scores in LLaMA2-7B-base, are much more important than other heads and should be preserved in KV cache compression. These findings are consistent with our observations.

Moreover, we observe that keys can be more compressed than values.
As shown by the heatmaps in Appendix D, the compression of values affects downstream tasks more than keys.
Specifically, according to our greedy search for adaptive compression levels, for a 37.5% KV cache budget, the optimized key cache budget is allocated 32.28%, and the value cache budget is allocated 42.72%.

## 6 Conclusion

In this study, we delve into how to compress the KV cache in LLMs by applying low-rank projection matrices to the feature dimension.
We first investigate data compression using the canonical orthogonal projection method through PCA.
We observe significant performance degradation at a relatively high compression rate, indicating that PCA projection is suboptimal for preserving global outputs due to LLMs’ nonlinearity and compounding effects.
To bridge the gap, we directly optimize orthogonal projection matrices for KV cache compression in LLMs with a distillation objective using an elaborate Matryoshka training strategy.
After training, we show that adaptive compression rates for different layers and heads ensure optimal performance compared to uniform compression rates across all layers and heads in LLMs.
Experimental results demonstrate significant performance gains and flexibility in achieving desired compression rates compared to traditional PCA projection.

## 7 Future work

We directly tune orthogonal projections using the Matryoshka training strategy with a distillation objective.
Experiments demonstrate that our method efficiently compresses the KV cache along the feature dimension.
Having established the compatibility of MatryoshkaKV cache with GQA, we intend to enhance its versatility by integrating our feature compression techniques with contemporary token merging and eviction methods, such as $\text{H}_{2}\text{O}$ (Zhang et al., ) and StreamingLLM (Xiao et al., ).
This integration aims to optimize compression rates, especially for extended sequences, and ensure broader compatibility with existing KV cache methodologies.

## Appendix A Error Analysis

Consider in a decoder layer with key and value has same head dimension: $d=d_{k/v}$, our orthogonal projection $U^{K},U^{V}\in\mathbb{R}^{d\times d}$
, error can be computed at a cache budget $r/d$ as:

$$ $\displaystyle\mathcal{L}$ $\displaystyle=\left\lVert\text{Attention}(Q,K,V)W^{O}-\text{Attention}(\tilde{ Q},\tilde{K},\tilde{V})W^{OV}\right\rVert_{F}^{2}$ $\displaystyle=\left\lVert\text{Softmax}\left(\frac{QK^{\top}}{\sqrt{d}}\right) VW^{O}-\text{Softmax}\left(\frac{\tilde{Q}\tilde{K}^{\top}}{\sqrt{d}}\right) \tilde{V}U^{V\top}_{r}W^{O}\right\rVert_{F}^{2}$ $\displaystyle=\left\lVert\text{Softmax}\left(\frac{QK^{\top}}{\sqrt{d}}\right) VU^{V}U^{V\top}W^{O}-\text{Softmax}\left(\frac{\tilde{Q}\tilde{K}^{\top}}{ \sqrt{d}}\right)VU^{V}_{r}U^{V\top}_{r}W^{O}\right\rVert_{F}^{2}$ $\displaystyle=\left\lVert\text{Softmax}\left(\frac{QK^{\top}}{\sqrt{d}}\right) V\left(U_{:,r:d}^{V}U_{:,r:d}^{V\top}+U^{V}_{r}U^{V\top}_{r}\right)W^{O}-\text {Softmax}\left(\frac{\tilde{Q}\tilde{K}^{\top}}{\sqrt{d}}\right)VU^{V}_{r}U^{V \top}_{r}W^{O}\right\rVert_{F}^{2}$ $\displaystyle=\left\lVert\left(\text{Softmax}\left(\frac{QK^{\top}}{\sqrt{d}} \right)-\text{Softmax}\left(\frac{\tilde{Q}\tilde{K}^{\top}}{\sqrt{d}}\right) VU^{V}_{r}U^{V\top}_{r}+\text{Softmax}\left(\frac{QK^{\top}}{\sqrt{d}}\right) VU_{:,r:d}^{V}U_{:,r:d}^{V\top}\right)\right\rVert_{F}^{2}$ $$

Consider the original parameter $W_{O}$. By donating $\mathcal{L}_{QK}=\text{Softmax}\left(\frac{QK^{\top}}{\sqrt{d}}\right)-\text{
Softmax}\left(\frac{\tilde{Q}\tilde{K}^{\top}}{\sqrt{d}}\right)$ and $A=\text{Softmax}\left(\frac{QK^{\top}}{\sqrt{d}}\right)$ is a constant, we just need to minimize:

$$ $\displaystyle\mathcal{L}$ $\displaystyle=\left\lVert\left(\mathcal{L}_{QK}VU^{V}_{r}U^{V\top}_{r}+AVU_{:, r:d}^{V}U_{:,r:d}^{V\top}\right)\right\rVert_{F}^{2}$ $\displaystyle=\left\lVert\mathcal{L}_{QK}+\left(A-\mathcal{L}_{QK}\right)VU_{: ,r:d}^{V}U_{:,r:d}^{V\top}\right\rVert_{F}^{2}$ $$

While PCA on value states minimizes $\left\lVert VU_{:,r:d}^{V}U_{:,r:d}^{V\top}\right\rVert_{F}^{2}$
and PCA on query and key states minimizes $\mathcal{L}_{QK}=\text{Softmax}\left(\frac{QK^{\top}}{\sqrt{d}}\right)-\text{
Softmax}\left(\frac{\tilde{Q}\tilde{K}^{\top}}{\sqrt{d}}\right)$, these optimizations do not necessarily guarantee the minimization of the global error $\mathcal{L}$, showing the PCA projection is suboptimal and has the room to be optimized to make the global error minimized.

## Appendix B weight merging method

Given that both $W^{Q}$ and $W^{K}$, as well as our orthogonal projection, operate on hidden states, consolidating parameters evidently reduces computational time.
However, many LLMs utilize RoPE Su et al. (), introducing a relative position embedding between $W^{Q}$ and $W^{K}$, which complicates integrating the parameters with our unitary transform.
This issue has been addressed in prior works Saxena et al. (); Yu et al. (). The approach in Saxena et al. () involves maintaining the merged parameters and transforming the compressed dimension cache back to its original dimensions for reapplication of RoPE.
This does not reduce peak memory usage for attention and necessitates RoPE for all past tokens.
Alternatively, (Yu et al., ) compresses the key states post-RoPE, which prohibits the merging of $W^{Q/K}$ and $U^{K}$.
However, as only a single new token requires orthogonal transformation and dimensionality reduction during inference, the time increase is merely slight as shown in (Yu et al., ).
Consequently, our treatment of RoPE in the present study is influenced by (Yu et al., )’s methodology.
The integration of the weight parameters of $W^{O}$ and $U^{V\top}$, given RoPE has no impact on value states, the details of our weight merging methods can be formulated as follows and in Figure 5

$$ $\displaystyle\text{MSA}(X)$ $\displaystyle=\text{Concat}(\text{head}_{1},\text{head}_{2},\ldots,\text{head} _{H})W^{O}$ $\displaystyle=\text{Concat}(A_{1}V_{1},A_{2}V_{2},\cdots,A_{h}V_{H})W^{O}$ $\displaystyle=\text{Concat}(A_{1}V_{1}U^{V}_{1}U^{V\top}_{1},A_{2}V_{2}U^{V}_{ 2}U^{V\top}_{2},\cdots,A_{h}V_{h}U^{V}_{H}U^{V\top}_{H})W^{O}$ $\displaystyle=\text{Concat}(A_{1}V_{1}U^{V}_{1},A_{2}V_{2}U^{V}_{2},\cdots,A_{ H}V_{H}U^{V}_{H})\left(\tilde{U^{V}}W^{O}\right)$ $\displaystyle=\text{Concat}(A_{1}\tilde{V_{1}},A_{2}\tilde{V_{2}},\cdots,A_{H} \tilde{V_{H}})W^{OV}$ $\displaystyle\text{where}\quad A_{i}$ $\displaystyle=\text{Softmax}\left(\frac{Q_{i}K_{i}^{\top}}{\sqrt{d_{k}}}\right )\text{is the attention weights of a given head in each layer}$ $$

Figure: Figure 5: After obtaining an orthogonal matrix through training, we merge the parameters in this way, reducing the number of matrix multiplications required during inference without incurring any inference time overhead. Truncation can be achieved simply by removing the columns corresponding to $W_{OV}$ , thereby reducing peak memory consumption.
Refer to caption: extracted/5927254/figure/weight_merge2.jpg

## Appendix C Two stage SFT

In this section, we provide a detailed discussion on our observations regarding fine-tuning with LoRA and the orthogonal matrix. We elaborate on the issues stemming from calculating covariance on a limited sample subset and performing spectral decomposition, which may lead to suboptimal parameters. We hypothesize that larger gradients during training can arise from task-specific distributions, such as in GSM8K, affecting the alignment of LoRA weights with the base model.

To mitigate these issues, our two-phase training approach involves initially training only the LoRA weights to ensure adequate adaptation to downstream tasks. In the second phase, we introduce simultaneous training of the unitary transformation matrix and the LoRA weights, focusing on maintaining performance while compressing the cache effectively. We also explore the impact of using separate learning rates for the LoRA and orthogonal matrix parameters to further investigate these phenomena. Extensive experimental results are provided to support our findings.

## Appendix D Experiment

We present some experimental results using a uniform compression rate across all heads after CPT and SFT in our MatryoshkaKV LLM.

**Table 3: Accuracy of our distilled MatryoshkaKV Projections after CPT on six benchmarks w/o greedy search for adaptive compression levels.**
| Model | Budget | Method | HLSG | ARC-C | ARC-E | PIQA | WG | CSQA | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| LLaMA2<br>7B-base | 100.0% | baseline | 74.00 | 35.93 | 50.97 | 78.50 | 61.64 | 65.93 | 61.16 |
| PCA | 72.04 | 36.95 | 52.38 | 76.66 | 61.72 | 67.24 | 61.17 |  |  |
| MKV | 72.05 | 37.29 | 52.38 | 76.66 | 61.72 | 67.32 | 61.24 |  |  |
| 87.5% | PCA | 30.28 | 23.73 | 30.34 | 58.05 | 51.30 | 22.60 | 36.05 |  |
| MKV | 72.22 | 35.93 | 52.20 | 76.28 | 62.12 | 65.27 | 60.67 |  |  |
| 75.0% | PCA | 25.47 | 27.80 | 27.51 | 52.67 | 49.72 | 20.56 | 33.96 |  |
| MKV | 70.98 | 34.58 | 55.20 | 76.77 | 61.56 | 63.64 | 60.46 |  |  |
| 62.5% | PCA | 24.22 | 28.81 | 27.51 | 51.58 | 50.28 | 21.29 | 33.95 |  |
| MKV | 69.22 | 37.29 | 55.73 | 75.22 | 59.35 | 64.21 | 60.17 |  |  |
| 50.0% | PCA | 24.04 | 28.47 | 25.22 | 52.29 | 50.67 | 20.72 | 33.57 |  |
| MKV | 66.62 | 34.24 | 52.91 | 75.46 | 58.41 | 62.00 | 58.27 |  |  |
| 37.5% | PCA | 24.08 | 28.47 | 25.40 | 50.76 | 49.49 | 18.35 | 32.76 |  |
| MKV | 62.38 | 32.20 | 50.26 | 73.34 | 56.67 | 55.28 | 55.02 |  |  |
| 25.0% | PCA | 23.98 | 29.49 | 26.28 | 51.20 | 50.36 | 16.22 | 32.92 |  |
| MKV | 51.91 | 27.46 | 44.44 | 69.64 | 54.54 | 44.39 | 48.73 |  |  |
| Mistral-v0.3<br>7B-base | 100.0% | baseline | 75.50 | 42.03 | 63.14 | 80.25 | 65.43 | 70.68 | 66.17 |
| PCA | 75.46 | 42.03 | 62.96 | 80.25 | 65.35 | 70.27 | 66.05 |  |  |
| MKV | 75.44 | 42.03 | 62.96 | 80.25 | 65.51 | 70.27 | 66.08 |  |  |
| 87.5% | PCA | 37.09 | 22.03 | 34.57 | 59.85 | 53.67 | 33.99 | 40.20 |  |
| MKV | 77.01 | 42.37 | 62.43 | 80.09 | 65.51 | 70.52 | 66.32 |  |  |
| 75.0% | PCA | 30.58 | 20.68 | 30.86 | 58.92 | 51.14 | 24.65 | 36.14 |  |
| MKV | 75.55 | 40.34 | 63.49 | 80.47 | 64.48 | 70.60 | 65.82 |  |  |
| 62.5% | PCA | 28.91 | 21.69 | 26.46 | 56.58 | 51.14 | 21.70 | 34.41 |  |
| MKV | 73.95 | 38.98 | 62.61 | 79.22 | 64.40 | 68.39 | 64.59 |  |  |
| 50.0% | PCA | 27.40 | 23.73 | 26.28 | 55.01 | 50.43 | 22.77 | 34.27 |  |
| MKV | 71.65 | 36.95 | 60.85 | 78.40 | 62.19 | 66.91 | 62.83 |  |  |
| 37.5% | PCA | 25.77 | 21.69 | 24.34 | 53.70 | 49.57 | 21.46 | 32.76 |  |
| MKV | 68.63 | 33.56 | 56.26 | 77.48 | 59.83 | 62.16 | 59.65 |  |  |
| 25.0% | PCA | 24.91 | 26.10 | 25.40 | 52.67 | 48.30 | 19.74 | 32.85 |  |
| MKV | 59.21 | 25.42 | 48.68 | 73.83 | 54.30 | 45.13 | 51.10 |  |  |

Figure: Figure 6: Two-phase SFT on PIQA.
Refer to caption: extracted/5927254/figure/piqa_better.jpg

Figure: Figure 8: Two-phase SFT on HellaSwag.
Refer to caption: extracted/5927254/figure/hellaswag.jpg

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
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?
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
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?
soundness: 3
strengths: 1. The paper is easy to follow.
2. Much stronger performance than the PCA baseline when the compression ratio is low.
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.
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.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
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?
soundness: 3
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.
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.
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

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: see weakness
soundness: 3
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
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 3
presentation: 3
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.
soundness: 2
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.
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.
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.

Meta-review: 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

Decision: Accept (Poster)