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
Venue: ICLR 2025 Poster (2025)
Keywords: Inference Optimization
KV Cache Compression
Low-rank Projection
Primary area: foundation or frontier models, including LLMs

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.

Published as a conference paper at ICLR 2025
MATRYOSHKAKY: ADAPTIVE KV COMPRESSION
VIA TRAINABLE ORTHOGONAL PROJECTION
Bokai Lin' Zihao Zeng' Zipeng Xiao Siqi Kou'
Tianqi Hou Xiaofeng Gao' Hao Zhang' Zhijie Deng
'Shanghai Jiao Tong University Huawei University of California, San Diego
(19821172068 zengzihao, xiaozp_25, happy-karry, thiarry,,,iieeeeee zhijied] edu. thou@ccnnect. ust. hk, gao-xi@cs. situ. edu.. en, naczhangeucsd. edu
edu. on
ABSTRACT
KV cache has become 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 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 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 orthogo- nal 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 directly tune the orthogonal projection matrices with distillation objective using an elaborate Matryoshka training strategy. After training, we adaptively search for the optimal compression rates for various layers and heads given vary- ing compression budgets. Compared to previous works, our method can easily embrace pre-trained L.MSs and hold 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
and Sittral-B--O.333-
INTRODUCTION
Large language models (LLMs) like GPT-4 OpenAl et al. 2024) and Claude3 (Enis Hopkins have shown great promise, finding applications in areas such as text generation (Brown et al., Raffel et al. 2023), code completion (Rozière al., 2024), and sentiment analysis (Zhang 2023a). The Key- Value (KV) cache, which is introduced to cache historical information for essential for maintaining context and accelerating the inference of LLMMs. How-
2024)
2020
et al.
self-attention,
as the size of the model and data continues to grow (Fu et al. 2024: Ding et al.) 2024 Chen
ct al. 2024), the KV cache can swiftly lead to system bottleneck in terms of storage and memory transfer (Shi et al. 2024).
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.
2024 Sun al. 2024 Goldstein al. (2024). head number (Ainslie et al.) 2023 Shazeer 2019 Yu
et al. 2024). and sequence length (Wang et al. 2024 Zhang et al. 2023b Li al. 2024 Xiao et al,
2024 Conversely, the exploration feature dimension for KV cache compression significantly lags
behind, partially because of the inherent difficulties of modifying a well-structured feature space.
"Corresponding author.

---

Published as a conference paper at ICLR 2025
ARC-C
ARC-E
WG
Layer Index
Layer Index
Key
25-
Layer
25
10
15
30
10
29
30
25
Hoad Index
Head Index
Head Ind
Value
Layer Index
Layor/nnde Indry
25-
25
25
30
20
155
20
25
30
10
25
25
20
20
25
Head Index
Head Index
Head ex
Figure 1:
Visualization of the feasible compression level for the keys and values in our model distilled from the LLaMA2-7B-bul model. We individually leverage samples in ARC-challenge (ARC-C), ARC-easy (ARC-E) (Clark et al. 2018). and Winogrande (WG) (Sakaguchi et al.) 2019) 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.
This paper aims to tackle this with the help 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 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 L.MSs while retaining reliable generation quality at moderate compression level. Compared to the low-rank architectures Multi-head I Latent Attention (MI.A) (DeepSeek-AL al.) 2024]. the PCA strategy is more approachable due to its training-free nature and also advocated by Saxena et al. (0024). Yet, we note that the PCA projections suffer from quickly degraded perfor- mance when further increasing the compression level. This is because, while the principal compo- nents 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- 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 knowledge distillation objective, enforcing the model output based 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, desire hierarchy over the columns of the projection matrices-an as in PCA-So that we can smoothly trade-off between compression level and performance. To this end, introduce Matryoshka training strategy compute the model output based on the first columns of the matrices, where 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 for different layers, heads, and keys/values during training to disen- tangle the projections in the model. Doing so enables the search for heterogeneous compression rates for different projection matrices during inference and we develop greedy algorithm for this. Heterogeneous compression rates are displayed in Figure
Experiments on both continual pre-training (CPT) and supervised fine-tuning (SFT) exhibit the ef- ficacy of our MatryoshkaKV approach. For the former, we opt to experiment on LLaMA2-7B- base Touvron et al. 2023) with the RedPajama dataset (Computer 2023). To demonstrate com- patibility with Group Query Attention (GQA) (Ainslie et al. 2023), we also apply our approach to the -i (Jiang al. 2023) model. Moreover, demonstrate that MatryoshkaKV
2

---

Published as a conference paper at ICLR 2025
is compatible with other KV cache compression techniques other axes like H2O (Zhag al. 2023b) and KIVI (Liu et al. 2023). We observe that after rarely processing 200 million train ing tokens, MatryoshkaKV achieves 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. 2021) downstream tasks including OBQA (Mihaylov et al. 2021), etc. The results show that our MatryoshkaKY can utilize less than 40% cache while still achieving over 90% accuracy derived from full cache utilization. We also perform ex- tensive ablation studies to chase deep understanding of our approach. The code is available at https://githus com wwtoo/TTeeeeeee
al. 2018), GSM8K (Cobbe
2
RELATED WORK
2024) and PyramidKV CCii. et al, 2024 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 V uses layer-wise approach with recent tokens occupying more weights. CLA (Brandon et al. 2024), YOCO (Sun et al., 2024). and GoldFinch (Goldstein et al.) (024), among others, exploit inter-layer KV cache reuse by sharing KV heads across layers. This signifi- cantly reduces the KV cache size along the head number dimension without compromising model capacity. GQA (Ainslie et al., 2023), MQA (Shazeer) 2019). and HeadK V Yu et al. (2024). espe- cially 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 eviction & merging. KVMerger (Wang et al.
KV cache hidden size compression. DeepSeekv2 (DeepSeek-Al Al et al., 2024 employs MLA tech- niques 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. Eigen-Attention (Saxena et al.] (022)) and HeadKV (Yu et al.] (024)) achieve a 40% re- duction 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 al. 202T) is employed fine-tune model parameters. However, this compression approach on the axis of feature dimension results in sharp decline in model performance when using less than 60% cache budget. Furthermore, fine-tuning the base model with LoRA may lead 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 (-)( with a as the prompt. It is 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 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 subsequent new token y with the attention states softmax((II (TKT /VdVV, where Q E R xd is the query vector for y and K, V R I.xd 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 ((L) to 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. 2024 Sun et al., 2024 Goldstein et al. 2024), merging heads that require caching KV (Yu et al. 2024 Ainslie et al. 2023), evicting or merging
redundant tokens (Xiao et al. 2024 Li al. [202 Cai. al.) 2024 Zhang et al. 2023b). 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

---

Published as a conference paper at ICLR 2025
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" K",V" V RLxr, (r < d). Then, we can only cache K' and V", reducing the storage and memory transfer cost from (d) to (r)) The rank 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 E Rdddd and demand a hierarchy over the columns of U that the optimal r-rank cache can result from the first r columns of U, denoted as U, E Rdx,, 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" and V" A natural choice is using UT, the transposition of the projection matrices, needs to be satisfied. Given that can vary from to d, we identify that U should known that the optimal orthogonal projections for compressing a set of
where U,U,T 2
be orthogonal matrices. 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 long sequence of tokens for each head separately.
Table 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 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 (DeepSekkAAT et al. 2024). We note that PCA projection is also advocated by Saxena et al. 2024); refer to Appendix B for the difference between our attempts and theirs regarding applying projections before or after RoPE (uu et al. 2023 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 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 compres- sion. The whole pipeline can be listed as follows: (1) Obtain the PCA initialization based on a small subset of a general corpus.
(2) Train our model on the corpus. (3) Search for the heterogeneous
compression levels for various heads with a small calibration dataset (5 10 samples) the spe- cific task. (4) Perform inference on that task given the identified compression levels. This section provides the training and inference details of our approach.
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 the original one. This implies a knowledge distillation objective (Hinton al., 2015), which can be instantiated with the KL divergence:
LKD (x (+1x) Ilp' (-|x))
(1)
where we abuse p' to refer to the LLM equipped with low-rank projection matrices. As suggested by the literature (Kou et al. 2024), we also incorporate a language modeling loss to p' 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, employ Cayley parameterization to formulate the orthogonal matrix. Specifically, there is U (T Q) (T -Q)= Q) with Q as skew-symmetric
4

---

Published as a conference paper at ICLR 2025
Fut dimansion ojo states
attention weights
Add & Norm
W
RoPE
Marmul Softmax
Reduced dimensiion stanes
FFN
W
RoPE
flestape
W
w
Add Norm
Attention
Mechanism of tradition attention using full-cache
Mechanism of reseraaeeeeeee w
using reduced-cache
attention weights
wr
RoPE
Matmul Sohmnax
now
wx
RoPE
U
W
Matrrll
W
Uy
cache
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 hier- archy over the columns of projection matrices so as concentrate the principal information on the head dimensions and enable the adjustment of compression level according to resource constraints.
trainable matrix of size x d. Considering that d usually small (c.g.. 64 or 128). the complexity of performing such 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 inher- ited 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 achieve 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. 2022). introducing a Matryoshka strategy for training the projection matrices U. In par- ticular, for each training iteration, we randomly sample r from predefined schedule such as {4, 8, 1i......//// d/4, d/2, d} and use the first columns of U, i..., Ur. to construct the model p' 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 4 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 specific compression budget. Basically, we can first propose compression level for the projection matrix at particular position, assessing the deviation of the model output from the original on predefined calibration dataset (measured by KI. divergence). and determining whether to accept the proposal based predefined tolerance threshold for the devia- tion. Algorithm exhibits a greedy algorithm for accelerating this based on accepting proposals in parallel. Note that this greedy algorithm also applies to the PCA projections.
The recent KV cache compression approach on sequence length aspect (Cai. et al. 2024h also observes that compared to uniformly compressed KV cache using the same rate across all layers (Li et al. 2024), employing a distinct compression rate for each layer results in improved information utilization. Furthermore, observed in (Wu al. 2024), certain retrieval heads within
Discussion.
5

---

Published as a conference paper at ICLR 2025
Algorithm 1: Greedy search for adaptive compression levels in our efficient L.M.
input: An base LLM p (+) and an efficient LLM equipped with MatryoshkaKV projections
p' (+), layer num L, attention head num H, full KV cache feature dimension d, prompt æ, compression rate interval Ar, target cache budget 7/ output: Two tensors R RV E L specifying the heterogeneous key/value compression
rates for each head in each layer.
RK RV +d. d. LLxH
repeat
Ricmps RV + RK RV for Every Layer-l in LLM do
RK
for Every Attention Head-h do
Rie I. Biemp.... & Bih Ar, Ry, Ar
Sth + DKL (p (p(122) 1Il (-12; Rremp: R²
EY, + DKL (p (-(2z) 1ip (-1x; Rk Ryempp)
Riemp,i.h RY Ryemp,L. RK RY,,
Locate the index associated with the minimum value element in the joint error list [e* E Decrement the corresponding compression rate in [RK R²] by Ar.
until Budget (RK .RV) <%
return RKK, RV
LLM consistently attend to crucial information, regardless of contextual variations. The indis- criminate compression rates of these heads can lead to significant performance degradation. These both support the necessity of the proposed heterogeneous KV cache compression approach.
EXPERIMENTS
In this section, we conduct experiments on continual pre-training (CPT) and supervised fine-tuning (SFT) scenarios to demonstrate that our MatryoshkaKV V can not only preserve the foundation knowl- edge of base model but also be compatible with LoRA (Hu al. (02T) for downstream tasks. Fur- thermore, combine our approach with KV cache compression technique targeting the sequence length dimension, referred to as H₂O (Zhang et al. 20236), and additionally implement another experiment by integrating KIVI (Liu al. 2023), KV cache compression strategy focused KV cache quantization. Ablation studies in Section 5.4 validate the efficacy of our proposed method.
5.1 CONTINUAL PRE-TRAINING
TTuvron et al. 2023 and Mistral-vO.3-7B-bane (Jiang al. 2023) as our base models. We conduct continual pre-training (Ke et al.) 2023) using the RedPa- Jama dataset (Computer 2023). To rapidly validate the effectiveness of proposed method, choose a subset of this dataset following ReedPaaaa--aaaa-1aa IT-Sapllee We adopt the Matryoshka training strategy detailed in Section 4.2 and fine-tune MatryoshkaK projections with knowledge distillation loss in Equation [ and language modeling loss, applying 1:3 weighting ratio between the two losses. The projection rank F'k and r. are randomly sampled from predefined schedule set { d] during training and are chosen dynamically with the greedy search for adaptive com- pression levels, as detailed Section 4.3 during inference. During the greedy search for adaptive compression levels, we define the compression rate interval Ar d/8 where the head dimen- sion d for each attention head (LaMAA2-B-basse is 128. We use Opencompass (Contributors) several widely-used zero-shot benchmarks: PIQA (Bisk et al. 2019), ARC- ARC-challenge (ARC-C) (Clark et al. 20188) ARC-easy (ARC-E) (Clark et al. 2018), WinoGrande (WG) (Sakaguchi et al. 2019). HellaSwag (HLSG) (Zellers et al., 2019), and CommonSense
Setup. We select LLaMA2-7B-bass
2023 to test performance
(CSQA) (Talmor et al. 2019). We compare our methods with Eigen-attention (Saxena et al. 2024
(donated PCA) in Table I and ASVD (uuan al. 2024) in Table 7

---

Published as a conference paper at ICLR 2025
Table 1: Comparison between MatryoshkaKV method (donated as MKV in the table) and PCA projection. We use LLaMA2-7B-bar and --ttral--003-7--BBe formance used baseline. Accuracy on HellaSwag, ARC-challenge, ARC-easy, PIQA, Wino- Grande, 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.
our source models, and their per-
Model
Budget
Method
HTSS
ARCC
ARCE
PIQA
WG
CSQA
Avg.
Baseline
74.00
35.93
50.97
78.50
61.64
65.93
61.16
100.0%
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
PCA
71.91
35.93
53.97
76.66
61.40
67.65
61.25
87.5%
MKV
71.58
37.97
53.26
75.95
62.12
69.57
61.74
PCA
70.99
35.59
54.14
76,22
60.06
66.99
60.67
LLaMA2 7B-base
75.0%
MKV
71.58
38.31
55.56
76.01
61.09
66.75
61.55
PCA
67.16
34.24
54.85
74.76
57.77
61.10
58.31
62.5%
MKV
68.03
37.97
56.08
75.12
60.30
65.44
60.49
PCA
42.11
29.83
35.10
58.16
52.57
40.62
43.07
50.0%
MKV
66.78
36.61
55.91
74.32
59.12
61.92
59.11
PCA
24.24
26.44
26.63
51.25
50.36
19.90
33.14
37.5%
MKV
63.97
33.90
51.68
74.97
57.92
59.21
56.94
PCA
23.98
29.49
26.28
51.20
50.36
16.22
32.92
25.0%
MKV
51.91
27.46
44.44
69.64
54.54
44.39
48.73
Baseline
75.50
42.03
63.14
80.25
65.43
70.68
66.17
100.0%
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
PCA
73.46
42.71
63.32
79.54
63.93
70.76
65.92
87.5%
MKV
75.63
42.03
64.37
79.71
65.51
70.35
66.27
PCA
70.75
37.63
61.73
78.18
62.59
68.47
63.23
Mistral-v0.3 7B-base
75.0%
MKV
75.29
43.39
63.14
79.54
64.96
69.12
65.90
PCA
63.48
34.24
55.73
75.90
60.77
62.24
58.73
62.5%
MKV
74.23
40.34
62.96
79.33
64.25
68.63
64.96
PCA
28.12
22.71
28.40
58.16
49.64
22.85
34,98
50.0%
MKV
73.32
38.98
62.08
79.16
61.88
67.08
63.75
PCA
25.04
22.03
28.04
53.86
49.25
21.21
33.24
37.5%
MKV
70.40
35.93
58.91
77.91
60.30
64.29
61.29
PCA
24.91
26,10
25.40
52.67
48.30
19.74
32.85
25.0%
MKV
59.21
25.42
48.68
73.83
54.30
45.13
51.10
Results. We train with a total of 30 GPU hours, processing just under 200 million tokens (20% of the RedPajama sample IT, i.e. 0.02% of the full RedPajama dataset). Table presents the results of our experiments. In zero-shot tasks, our MatryoshkaKV cache substantially reduces the cache foot- print with minimal impact on performance. Specifically, our method retains 93. 10% of LLaMA2- ' average accuracy, while utilizing
7B-base's average accuracy and 92.63% of -0. only 37.5% of the original cache size. For simpler tasks like PIQA, it achieves 88. 38.71% % and 92.00% of the base model's performance with just a 25% cache budget. On more challenging tasks such as ARC-C, 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 bud- get is reduced below 62.5%, achieving just 70.42% accuracy of (LLMMA--B--bar and 52.86% of -istral-vO.3-7BBB
These results underscore the superior efficiency of our approach compared with PCA. We attribute PCA's performance decline to suboptimal projection matrices, whereas method maintains closer alignment with the base model, thereby mitigating this degradation.
Furthermore, evaluate the inference speed of LLM equipped with our MatryoshkaK The results are presented in Table 3 and related discussions are detailed in Appendix

---

Published as a conference paper at ICLR 2025
Table 2: Accuracy of Matryoshka method after SFT based on LLaMA2-7B--bas four down-
stream tasks: PIQA, GSMSK, HellaSwag, and OpenbookQA, at seven KV cache budgets. Degrada- tion from baseline shown in brackets.
Model
Budget
PIQA
GSM8K
HLSG
OBQA
Avg.
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%)
LLaMA2
62.5 %
82.75
31,46
89.86
79.60
70.92 (-4.27%)
7B-base
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%)
0.9
37.5%
WFFCX inttataattr wil PCA Intia Faation
50%
0.8
62.5%
7-
100%
6-
Evall Loss
0.7
Eval Loss
0.6
2-
1,3K
5
6.4k
7.74
2.6k
5.1K
7.7k
10.3K
12
8K
15
2k
17 9K
20 5k
23 OK
25 8K
28 28.2K
Num samples
Num samples
Figure 3:
Evaluation loss of four budgets vs. the number of training samples during 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 epochs of SFT on GSM8K (Right).
5.2 SUPERVISED FINE-TUNING
Setup. We use LLaMA2-7B-base Touvron al. 2023) as our base model and verify the efficacy of our method on PIQA (Bisk et al. 2019). GSM8K (Cobbe et al. 2021). HellaSwag (Zellers et al.) 2019) and Openbook (OBQA) (Mihaylov al. 2018) datasets. We design two-stage training strategy to make Matryoshka training strategy compatible with LoRA (Hu et al. 2021 fine-tuning. Specifically, LoRA is firstly used adapt the base model to downstream tasks, following standard SFT practices (Naved et al., 2024 Zhao et al. 2024). In the second stage, we jointly fine-tune the MatryoshkaKV projections with the Matryoshka training strategy and the LoRA parameters. Further discussion the superiority of this recipe is detailed in Appendix
Results. We report accuracy on four zero-shot benchmarks at seven KV cache budgets in Table As shown, 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 89.47% of the full-cache performance with 37.5% cache budget. However, for more complex tasks such as GSM8K, a 50% cache budget is necessary to achieve comparable results. Furthermore, 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.
report the evaluation loss four budgets: 100%,
5.3 COMPATIBILITY WITH OTHER KV CACHE COMPRESSION TECHNIQUES
To demonstrate the orthogonality and compatibility of method with existing KV cache compres- sion techniques, we conduct extensive experiments utilizing MatryoshkaKV in conjunction with these methods. Based the classification outlined in Section integrate MatryoshkaKV
8

---

Published as a conference paper at ICLR 2025
Table 3: Tokens per second at different KV cache budgets with batch size of 32,
LLaMA2
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.0
60
Relative Accuracy
335
Avg. Accuracy
50-
45
0.4
PCA search sigortthm
MIV W/ search aigonttmm
Orthegonal Constraint ozhia Training sctting
35
we algorithm
MEV aigorthhm
0.0
Cache Utillization
Cache Utilization
Figure 4: Comparison between PCA and distilled MatryoshkaKV Projections after CPT with and without greedy search for adaptive compression levels. We report average accuracy 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-78 ----bbe model fine-tuned with LoRA on GSM8K, utilizing the full KV cache (Right).
with prominent techniques such as KIVI (Hooper et al. (024) for KV quantization, H₂O (Zhang jet al., 2023b), and GQA (Ainslie et al. 2023) for KV cache eviction and merging. We apply Ma- tryoshkaKV V to bility in synergy with GQA (Ainslie et al. 2023).
in Section 5.1 demonstrating its enhanced compression capa-
Combination with H,O. Furthermore, combine methods with H₂O. We first evaluate H₂O and our MatryoshkaKV on datasets mentioned 5.1. To demonstrate improved compression rates in long contexts, we select LongBench (Bai et al. 2024) and calculate perplexity under different cache budget settings of two orthogonal KV compression techniques. For the detailed results, please refer to Table in Appendix E
According to the results, by concurrently using MatryoshkaKV and H2O, the perplexity on long contexts increases by merely 1.02 10% KV cache budget. Additionally, if 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 6 benchmarks, which 91.32% of the baseline.
Combination with KIVI. In addition to integrating with HzO, we also explore the combination of our methods with KIVI (Liu al. (023). a KV cache compression technique based 2-bit cache quantization. Similar to the previous approach, we conduct evaluations on the datasets described Section [5.1] The detailed results of this combination are presented Table [ in Appendix E and analyzed in detail. The results show that our MatryohskaKV can be easily combined with K quantization techniques and achieve higher compression rate.
5.4 ABLATION STUDIES
We conduct ablation studies on various components of our method to verify their effectiveness.
W/o greedy search for adaptive compression levels. We evaluate our trained models without our greedy search for adaptive compression levels. Figure (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 in Appendix D To ensure that each head in the LLM plays its due role, minimum cache budget for each head. At 37.5% cache budget, the average
set 25% accuracy improves by 1.92%, indicating the significance of our search algorithm for further KV cache compression. Furthermore, MatryoshkaKV demonstrates robustness when applying

---

Published as a conference paper at ICLR 2025
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 point out that the tuming pro- cess w/o Matryoshka training strategy can destroy the hierarchical structures present in the orthogo- nal matrices inherited from the PCA ones. To validate this, we train MatryoshkaKV projections with (Right). We observe that fix- ing 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 destroyed.
a fixed KV cache budget of 50%. The result is displayed in Figure 4
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 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 GSMSK and impose orthogonal constraints. In Figure 3 (Right),
report the evaluation loss during the second stage of SFT on GSM8K. Despite training for four epochs, randomly initialized orthogonal projections fail to converge optimal solution, and the text generated by our fine-tuned LLM projection is composed of meaningless symbols. This highlights the importance of PCA initialization.
5.5 HETEROGENEOUS COMPRESSION RATES VISUALIZATION
Figure I shows the heterogeneous compression levels across all attention heads inside our Ma- tryoshkaKV LLM distilled from the LLaMA2-7B-baa 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 minority of specific heads require a relatively high budget. PyramidKV (Cai. et al., 2024) 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 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. (2024) point out, retrieval heads with high retrieval scores in LLaMA2-775 7B-A---BBBe BB-bse, much more important than other heads and should be preserved in KV cache compression. These findings are consistent with 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 37.5% KV cache budget, the optimized key cache budget is allocated 32.28%, and the value cache budget allocated 42.72%.
6 CONCLUSION
In this study, we delve into how to compress the KV cache in LLMs by applying low-rank projec- tion 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 opti- mize 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.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
This work was supported by NSF China (Nos. 92470118, 62306176), Natural Science Founda- tion of Shanghai (No. 23ZR1428700), and CCF-Zhipu Large Model Innovation Fund (No. CCF- Zhipu202412).
REFERENCES
Joshua Ainslie, James Lee-Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points, 2023. URL
org/ abs/ htlp://bs/230055 2305 13245.
Yushi Bai, Xin Lv, Jiajie Zhang. Hongchang Lyu, Jiankai Tang. Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: bilingual, mul- titask benchmark for long context understanding, 2024. URL https: http://axivvv.o org.as///rg//////////
2308. 14508
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piga: Reasoning about physical commonsense in natural language, 2019. URL httts:://axxi..oottt 11641
William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, and Jonathan Ragan
Kelly, Reducing transformer vey-value cache size with cross-layer attention, 2024. URL https:
arxiv. ppxiivorg/abs//200000 12299 12981
Tom B. Brown, Benjamin Mann, Nick Ryder. Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Vors Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL pttps://arxioe / abs 2005. 14165
Zefan Cai., Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, and Wen Xiao. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, 2024. URL
http://wwrr.aabb/2......0 02069
Yukang Chen, Shengju Qian, Haotian Tang. Xin Lai, Zhijian Liu, Song Han, and Jiaya Jia. Longgor Efficient theettuning of long-context large language models, 2024. URL https://arxiv.
org/ 20tg/bb/25330 12307
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try the ai2 reasoning challenge, 2018. URL htttp::////aa
wwtp:///gggrraaaoo//ooo 05457
Karl Cobbe, Vincet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021. URL https://arxiv.
org/abs/ 21 20000.11...1.. 14168
Together Computer.
Redpajama: An open dataset for training large language models, October 2023.
URL htllhttts:///gthubb com http://ttttt...tee..eeeeeeee /RedP a ama Data
OpenCompass Contributors. Opencompass: A universal evaluation platform for foundation models. wwt. wwwpp//eww.oooppp///e/oooeepepee 2023.
DeepSeek-AI, Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi
Dengr, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li. H. Zhang, Hanwei Xu, Hao Yang, Haowei Zhang, Honghui Ding, Huajian Xin, Huazuo Gao. Hui Li, Hui Qu, J. L. Cai, Jian Liang, Jianzhong Guo, Jiaqi Ni, Jiashi Li, Jin Chen, Jingyang Yuan, Junjie Qiu, Junxiao Song. Kai Dong, Kaige Gao, Kang Guan, Lean Wang, Lecong Zhang, Lei Xu, Leyi Xia, Liang Zhao Liyue
11

---

Published as a conference paper at ICLR 2025
Zhang, Meng Li, Miaojun Wang, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Mingming
Li, Ning Tian, Panpan Huang, Peiyi Wang, Peng Zhang, Qihao Zhu, Qinyu Chen, Qiushi Du, R. J.
Chen. R. L. Jin, Ruiqi Ge. Ruizhe Pan, Rumxin Xu, Ruyi Chen, S. S. Li, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shaoqing Wu, Shengfeng Ye, Shirong Ma, Shiyu Wang, Shuang Zhou, Shuiping Yu, Shunfeng Zhou, Size Zheng. T. Wang. Tian Pei, Tian Yuan, Tianyu Sun, W. L.. Xiao, Wangding Zeng, Wei An, Wen Liu, Wenfeng Liang, Wenjun Gao, Wentao Zhang, X. Q. Li, Xiangyue Jin, Xianzu Wang, Xiao Bi, Xiaodong Liu, Xiaohan Wang, Xiaojin Shen, Xiaokang Chen. Xiaosha Chen, Xiaotao Nie, Xiaowen Sun, Xiaoxiang Wang, Xin Liu, Xin Xie, Xingkai Yu, Xinnan Song, Xinyi Zhou, Xinyu Yang. Xuan Lu, Xuecheng Su, Y. Wu, Y. K. Li, Y. X. Wei, Y. X. Zhu, Yanhong Xu, Yanping Huang, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Li, Yaohui Wang, Yi Zheng, Yichao Zhang, Yiliang Xiong, Yilong Zhao, Ying He, Ying Tang. Yishi Piao, Yixin Dong, Yixuan Tan, Yiyuan Liu, Yongji Wang, Yonggiang Guo, Yuchen Zhu, Yuduan Wang, Yuheng Zou, Yukun Zha, Yunxian Ma, Yuting Yan, Yuxiang You, Yuxuan Liu, Z. Z. Ren, Zehui Ren, Zhangli Sha, Zhe Fu, Zhen Huang, Zhen Zhang. Zhenda Xie, Zhewen Hao, Zhihong Shao, Zhiniu Wen, Zhipeng Xu, Zhongyu Zhang, Zhuoshu Li, Zihan Wang, Zihui Gu, Zilin Li, and Ziwei Xie. v2: A strong, economical, and efficient mixture-of-expprts language model, 2024. BtRphtpss: http::///xxxivi.i.e.oo httpr//aab//440.... 04434
Yiran Ding, Li Lyna Zhang, Chengruidong Zhang, Yuanyuan Xu, Ning Shang, Jiahang Xu, Fan Yang, and Mao Yang. Longrope: Extending Ilm context window beyond 2 million tokens, 2024. URL httph:ttps///tc.t.
Maxim Enis and Mark Hopkins. From Ilm to nmt: Advancing low-resource machine translation with claude, 2024. URL
/200..//2444.. 13813
Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannanch Hajishirzi, Yoon Kim, and Hao Peng. Data engineering for scaling language models to 128k context, 2024. URL httpps://///xii
2402./242..10 10171
Daniel Goldstein, Fares Obeid, Eric Alcaide, Guangyu Song, and Eugene Cheah. Goldfinch: High performance rwkv/transformer hybrid with linear pre-fill and extreme kv-cache compression, 2024. http:ttt:://aaxt
abs 220024402222 12077
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network, 2015. URL httpp/ss//ariiii.rrrggg////////p 1503. 02531
Coleman Hooper, Schoon Kim, Hiva Mohammadzadeh, Michael W. Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length IIm inference with kv cache quantization, 2024. (RLLhttps::/// //t::i//aaiviii//a org/ abs/abb//22 22001/241... 18079
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models, 2021. URL https:
//arxiv. wrw.vorrgaaee ab/////a//// 22w. 09685
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chap- lot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William Sayed. Mistral 7b, 2023. URL https:
06825
Zixuan Ke, Yijia Shao, Haowei Lin, Tatsuya Konishi, Gyuhak Kim, and Bing Liu. Continual pre- training of language models, 2023. URL httt:httts///aaa/t org/ abs/ 22022 03241
Siqi Kou, Lanxiang Hu, Zhezhi He, Dengi and Hao Zhang. Clims: Consistency large language models, 2024.
Aditya Kusupati, Gantavya Bhatt, Aniket Rege, Matthew Wallingford, Aditya Sinha, Vivek Ra- manujan, William Howard- -nyder, Kaifeng Chen, Sham Kakade, Prateek Jain, et al. Matryoshka representation learning. Advances in Neural Information Processing Systems, 35:30233-300 -32249, 2022.
12

---

Published as a conference paper at ICLR 2025
Yuhong Li, Yingbing Huang. Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Lim knows what you are looking for before 14469
generation, 2024.
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi: Plug-amd play 2bit kv cache quantization with streaming asymmetric quantization. attps://rguee net lished.
net 1t...............
28167. 37282 2023. Unpub-
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct clectricity? new dataset for open book question answering, 2018. URL atg/abs/18099 02789
Humza Naveed, Asad Ullah Khan, Shi Qiu, Muhammad Saqib, Saced Anwar, Muhammad Usman, Naveed Akhtar, Nick Barnes, and Ajmal Mian. A comprehensive overview of large language models, 2024. htthsts:::/ www.v..iivv.a org/ abs 2207 06435
OpenAI, Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Iige Akkaya, Floren- cia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, Red Avila, Igor Babuschkin, Suchir Balaji. Valerie Balcom, Paul Baltescu, Haiming Bao, Moham- mad Bavarian, Jeff Belgum, Irwan Bello, Jake Berdine, Gabriel Bernadett-Shapinn Christopher Berner, Lenny Bogdonoff, Oleg Boiko, Madelaine Boyd, Anna-Luisa Brakman, Greg Brock- man, Tim Brooks, Miles Brundage, Kevin Button, Trevor Cai, Rosie Campbell, Andrew Cann, Brittany Carey, Chelsea Carlson, Rory Carmichael. Brooke Chan, Che Chang, Fotis Chantzis, Derek Chen, Sully Chen, Ruby Chen, Jason Chen, Mark Chen, Ben Chess, Chester Cho, Casey Chu, Hyung Won Chung, Dave Cummings, Jeremiah Currier, Yunxing Dai, Cory Decareaux, Thomas Degry, Noah Deutsch, Damien Deville, Arka Dhar, David Dohan, Steve Dowling, Sheila Dunning, Adrien Ecoffet, Atty Eleti, Tyna Eloundou, David Farhi, Liam Fedus, Niko Felix, Simón Posada Fishman, Juston Forte, Isabella Fulford, Leo Gao, Elie Georges, Christian Gib- son, Vik Goel, Tarun Gogineni, Gabriel Goh, Rapha Gontijo-Lppes, Jonathan Gordon, Morgan Grafstein, Scott Gray, Ryan Greene, Joshua Gross, Shixiang Shane Gu, Yufei Guo, Chris Hal- lacy, Jesse Han, Jeff Harris, Yuchen He, Mike Heaton, Johannes Heidecke, Chris Hesse, Alan Hickey, Wade Hickey, Peter Hoeschele, Brandon Houghton, Kenny Hsu, Shengli Hu, Xin Hu, Joost Huizinga, Shantanu Jain, Shawn Jain, Joanne Jang, Angela Jiang, Roger Jiang, Haozhun Jin, Denny Jin, Shino Jomoto, Billie Jonn, Heewoo Jun, Tomer Kaftan, tukasz Kaiser, Ali Ka- mali, Ingmar Kanitscheider, Nitish Shirish Keskar, Tabarak Khan, Logan Kilpatrick, Jong Wook Kim, Christina Kim, Yongjik Kim, Jan Hendrik Kirchner, Jamie Kiros, Matt Knight, Daniel Kokotajlo, Lukasz Kondraciuk, Andrew Kondrich, Aris Konstantinidis, Kyle Kosic, Gretchen Krueger, Vishal Kuo, Michael Lampe, Ikai Lan, Teddy Lee, Jan Leike, Jade Leung, Daniel Levy. Chak Ming Li, Rachel Lim, Molly Lin, Stephanie Lin, Mateusz Litwin, Theresa Lopez, Ryan Lowe, Patricia Lue, Anna Makanju, Kim Malfacini, Sam Manning, Todor Markov, Yaniv Markovski, Bianca Martin, Katie Mayer, Andrew Mayne, Bob McGrew, Scott Mayer McKinney, Christine McLeavey, Paul McMillan, Jake McNeil, David Medina, Aalok Mehta, Jacob Menick, Luke Metz, Andrey Mishchenko, Pamela Mishkin, Vinnie Monaco, Evan Morikawa, Daniel Mossing, Tong Mu, Mira Murati, Oleg Murk, David Mély, Ashvin Nair, Reichiro Nakano. Ra- jeev Nayak, Arvind Neelakantan, Richard Ngo, Hyeonwoo Noh, Long Ouyang, Cullen O' Keefe, Jakub Pachocki, Alex Paino, Joe Palermo, Ashley Pantuliano, Giambattista Parascandolo, Joel Parish, Emy Parparita, Alex Passos, Mikhail Pavlov, Andrew Peng, Adam Perelman, Filipe de Avila Belbute Peres, Michael Petrov, Henrique Ponde de Oliveira Pinto, Michael, Pokorny, Michelle Pokrass, Vitchyr H. Pong, Tolly Powell, Alethea Power, Boris Power, Elizabeth Proehl, Raul Puri, Alec Radford, Jack Rae, Aditya Ramesh, Cameron Raymond, Francis Real, Kendra Rimbach, Carl Ross, Bob Rotsted, Henri Roussez, Nick Ryder, Mario Saltarelli, Ted Sanders, Shibani Santurkar, Girish Sastry, Heather Schmidt, David Schnurr, John Schulman, Daniel Sel- sam, Kyla Sheppard, Toki Sherbakov, Jessica Shieh, Sarah Shoker, Pranav Shyam, Szymon Sidor, Eric Sigler, Maddie Simens. Jordan Sitkin, Katarina Slama, Ian Sohl, Benjamin Sokolowsky, Yang Song, Natalie Staudacher, Felipe Petroski Such, Natalie Summers, Ilya Sutskever, Jie Tang, Nikolas Tezak, Madeleine B. Thompson, Phil Tillet, Amin Tootoonchian, Elizabeth Tseng, Pre- ston Tuggle, Nick Turley, Jerry Tworek, Juan Felipe Cerón Uribe, Andrea Vallone, Arun Vi- jayvergiya, Chelsea Voss. Carroll Wainwright, Justin Jay Wang, Alvin Wang. Ben Wang. Jonathan Ward, Jason Wei, CJ Weinmann, Akila Welihinda, Peter Welinder, Jiayi Weng, Lilian Weng.
13

---

Published as a conference paper at ICLR 2025
Matt Wiethoff, Dave Willner, Clemens Winter, Samuel Wolrich, Hannah Wong, Lauren Work- man, Sherwin Wu, Jeff Wu, Michael Wu, Kai Xiao, Tao Xu, Sarah Yoo, Kevin Yu, Qiming Yuan, Wojciech Zaremba, Rowan Zellers, Chong Zhang, Marvin Zhang, Shengjia Zhao, Tianhao Zheng, Juntang Zhuang, William Zhuk, and Barret Zoph. Gpt-4 technical report, 2024. URL ht.tp::/aaxxiv wttp://wwwww.///o//o/ooooo 087774
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2023. URL https: /ttpss//arxxi..... www.a. wttl wwwaaasslb/ 10683
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre, Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Ev- timov, Joanna Bitton, Manish Bhatt, Cristian Canton Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défossez, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usumier, Thomas Scialom, and Gabriel Synnaeve. Code Ilama: Open foundation models for code, 2024. URL http:https///xttto org/ balos//23...22
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande; An adver- sarial winograd schema challenge at scale, 2019. URL
10641
Utkarsh Saxena, Gobinda Saha, Sakshi Choudhary, and Kaushik Roy. Eigen attention: Attention in low-rank space for kv cache compression, 2024. http:/tpss///xii.orr/////a//////
05646
Noam Shazeer. Fast transformer decoding: One write-head is all you need, 2019. URL https:
//arxiv.org/ant www./bs.//a/////a/eoee 19111021510
Luohe Shi, Hongyi Zhang, Yao Yao, Zuchao Li, and Hai Zhao. Keep the cost down: review on methods optimize I1m' s kv-cache consumption, 2024. URL htttt::t//a//iiiii www.wwwwwwwwaaaaeeeae.eeeoeeeeee 2407. 18003
Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, and Yunfeng Liu. Roformer: En hanced transformer with rotary position embedding, 2023. URLhttts:://axt ortpabs/r///////////
2104. 09864
Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, and Furu Wei. You only cache once: Decoder-decoder architectures for language models, 2024. URL whtph:ts:s. //tpxi/.aaiiii.i wttp:///////.....555.52
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Commonsenseqa: A question answering challenge targeting commonsense knowledge, 2019. hRtpshtss://axiiti
abs/1811.00937
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Maric-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023. URL httpt::p///// http//axii......aaa...
arxiv. org. wwtp://..rga/aaaa///eea// 2tt.b///33...2 230880928788888
Zheng Wang, Boxiao Jin, Zhongzhi Yu, and Minjia Zhang. Model tells you where to merge: Adap- tive kv cache merging for llms on long-context tasks, 2024. URL htttps://arxii //tpxii..ooo..///. abs/ tbs/2400 2407. 08454
14

---

Published as a conference paper at ICLR 2025
Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mecha- nistically explains long-context factuality, 2024. URL https: wwwiv.or.oooo...... absp///b///444
15574
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming lan- guage models with attention sinks, 2024. URL BtLLhttp::/ https://axxxx........mm /ttp//aaxvvvvvva wttpv////rgabbb//33333 17453
Hao Yu, Zelan Yang, Shen Li, Yong Li, and Jianxin Wu. Effectively compress kv heads for Ilm, 2024. URL [tttps::://arrtt org/ abs/ 02056
Zhihang Yuan, Yuzhang Shang, Yue Song. Qiang Wu. Yan Yan, and Guangyu Sun. Asvd: Activation-aware singular value decomposition for compressing large language models, 2024.
URL http::t/a//xxi..oorrttte abs/2312. 05821
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence?, 2019. URL httlhtts///// //axpii arxiv. org 1905.0111111 07830
Wenxuan Zhang, Yue Deng, Bing Liu, Sinno Jialin Pan, and Lidong Bing. Sentiment analysis in the era large language models: reality check, 2023a. URL /tttpss://xxii.../// /bs/
23305.1001
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, and Beidi Chen. H2o: Heavy- hitter oracle for efficient generative inference of large language models, 2023b. URL https:
//axxiv.orq/abat..... 220/ 14048
Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolei Wang, Yupeng Hou, Yinggian Min, Beichen Zhang, Junjie Zhang, Zican Dong, Yifan Du Chen Yang. Yushuo Chen, Zhipeng Chen, Jinhao Jiang, Ruiyang Ren, Yifan Li, Xinyu Tang. Zikang Liu, Peiyu Liu, Jian- Nie, and Ji-Rong Wen. A survey of large language models, 2024. URL
www.lvvvrrrgaaae abs/
2303 18223
A ERROR ANALYSIS
Consider in a decoder layer with key and value has same head dimension: d= d our orthogonal projection UK U. E Rdxd we donate the first r columns of orthogonal projection U as U, and the as Unddd the error can be computed cache budget r/d as:
(r)
Attention Q, K, V)WO
Attention Q, R, VV
F
QK
QK Vd
Softmax
VWO
Softmax
VUNT wo
Vd
QK
QQKT Vd
Softmax
VVTOVTWO W Softmax
VUNUN wo
Vd
[QK
QK Vd
Softmax
(UV
(UN.aUN hnd + UUU.TT wo Softmax
VU UU UV WO
Vd
QKT
QKT
QKT Vd
Softmax
Softmax
VUNUKT + Softmax
VUV
UT sind
Vd
Vd
QQ
QRT Vd
Consider the original parameter Wo. By donating Lok = Softmax
Softmax
QQT
and A Softmax
is constant, we just need to minimize:
Vi
15

---

Published as a conference paper at ICLR 2025
(r) [[oo + yra (rrde) ll U LLx + (A- (A-LQk) VU
VUY UVT .rd Srid
While PCA on value states minimizes
and PCA on query and key states minimizes
Lox Softmax QK
QRT vdd
these optimizations do not necessarily guarantee the minimization of the global error L, showing the PCA projection suboptimal and has the room to be optimized to make the global error minimized.
Softmax
vd
The LLM itself has numerous layers, and each layer is nonlinear. Strietly speaking, the error is the output of the last layer of the model after low-rank projection and that of the original model's last layer. Here, we only conduct an intuitive analysis of certain layer. The optimal solution of this optimization problem is complex and difficult to solve mathematically. So, we make these orthogonal matrices trainable to get optimal results.
Theoretically, the optimal solution also changes with the variation of the input data distribution. It difficult for us to model the distribution of all corpora in the world. Therefore, to minimize the error of the model after KV cache compression most tasks as much as possible, consider using data-driven approach for optimization to be a reasonable method.
Attention K, V)WO Attention K, KVWoovl F² use KI- Divergence as a proxy loss to let the distributions of the two models' outputs close to each other. As we have discussed in Section 3.2 to recover the original K and V from the reduced K" and V when using full-rank, the orthogonality of U should be guaranteed. Thus, our optimization objective can be derived as:
To minimize (r)
U arg min
Σ
DKT. (p (-|x)) lIl (-|x;; U.K UN))
rEM
s.t. DKL. (p (-1zz) Ilp' (-|=; UY))= =0
(4)
where we abuse p' to refer to the LLM equipped with low-rank projection matrices, and M is our predefined schedule.
Our orthogonal constraints UUT = on U can guarantee DEL (p (-|w) IP (-Ja;; U. UY) 0.
worth noticing that we only use r < columns to forward for KV cache compression, the rest d columns, i.c. Unrd will not be updated. Thus, although experiments in Appendix G demonstrate that our method is not sensitive to a predefined schedule, we point out that M is must to guarantee all parameters of U to be trained.
B WEIGHT MERGING METHOD
Given that both 8º and W well as our orthogonal projection, operate on hidden states, con- solidating parameters evidently reduces computational time. However, many LLMs utilize RoPE et al. (002)). introducing a relative position embedding between WQ and WK which compli- cates integrating the parameters with our unitary transform. This issue has been addressed in prior works|Saxenn et al. 2024); Yu et al. (2024). The approach in Saxena et al. (2024) involves maintain- ing 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.] 2024) compresses the key states Rost- which prohibits the merging of WQ/K and U 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. 2024). Consequently, our treatment of RoPE in the present study is influenced by (uu et al. 2024 's methodology. The integration of the weight parameters of
16

---

Published as a conference paper at ICLR 2025
Wo and UVT, given RoPE has no impact value states, the details of our weight merging methods can be formulated as follows and in Figure
MSA(X) Concat(headj head2..................
head W
Concat(A1
A&Va)
Concat( (A:taat(AA A2V22U UT .. AA AVVUUUU
Concat Concat(A,/Viy) A2V2U ... AAVAuU)) (iirwo)
Concat( A, 1,, A2 V2, ABV
A Softmax Q
where
is the attention weights of a given head in each layer
dd
U
dk Uy
merged output projector wov
Uyy
output projector wo
block-wise orthogonal
Figure 5: After obtaining 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 correspond- ing to Wov thereby reducing peak memory consumption.
C TWO STAGE SET
In this section, provide detailed discussion observations regarding fine-tuning with LoRA and the orthogonal matrix. We elaborate on the issues stemming from calculating covariance 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 distri- butions, such 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 ensure adequate adaptation downstream tasks. In the second phase, introduce si- multaneous training of the unitary transformation matrix and the LoRA weights, focusing on main- taining 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 support our findings.
17

---

Published as a conference paper at ICLR 2025
32
75.35
77.97
32
22.21
29.19
29.04
Key Head Size
Key Head Size
73. 34
79.33
82.75
83.13
26.
34.27
39.04
39.8
74.65
80.09
83.3
83 84
96
28.2
35,94
38.51
42.76
128
75.68
81.56
84.28
84.22
128
30.25
40.18
40.26
12
32
128
128
Value Head Size
Value Head Size
Figure 6: Two-phase SFT on PIQA.
Figure Two-phase SFT on GSM8K.
32
56.1
76.1
82.23
84.71
32
61.4
70.8
75.0
77.6
Key Head Size
Key Head Size
64
70.01
86.29
89.86
91.85
64
67.0
76.5
79.6
80.4
88.62
93.17
76.2
81.6
128
76.11
09.01
92.26
93.
128
67.8
79.2
01.0
83.2
128
128
Value Head Size
Value Head Size
Figure 8: Two-phase SFT on HellaSwag.
Figure 9: Two-phase SFT OBQA.
D ABLATION STUDY ON GREEDY SEARCH FOR ADAPTIVE COMPRESSION LEVELS
We present some experimental results using a uniform compression rate across all heads after CPT and SFT in our MatryoshkaKV LLM. The results are displayed in 4
18

---

Published as a conference paper at ICLR 2025
Table 4: Accuracy of distilled MatryoshkaKV Projections after CPT six benchmarks w/o greedy search for adaptive compression levels.
Model
Budget
Method
HLLG
ARC-C
ARC-E
PIQA
WG
CSQA
Arg.
baseline
74.00
35.93
50.97
78.50
61.64
63.93
61.16
52.38
76.66
MKV
72.05
37.29
52.38
76.66
61.72
67.32
61.24
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
LLaMA2
75.0%
PCA
23.47
27.80
27.51
52.67
49.72
20.56
33.96
7B buse
55.20
76.77
63.64
24.22
28.81
27.51
51.58
50.28
21.29
62.5%
MKV
69.22
37.29
55.73
75.22
59.35
64.21
60.17
PCA
24.04
28.47
25.22
52.29
30.67
20.72
33.57
50.0%
MKV
66.62
34.24
52.91
73.46
58.41
62.00
58.27
PCA
24.08
28.47
25.40
50.76
49.49
18.35
32.76
37.5%
32.20
50.26
73.34
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
baseline
63.14
100.0%
PCA
75.46
42.03
62.96
80.25
65.35
70.27
66.05
75.44
42.03
62.96
80.25
65.51
70.27
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
Mistral-vO.3 78-base
75.0%
30.58
20.68
30.86
58.92
24.65
MKV
75.55
40.34
63.49
80.47
64.48
70.60
65.82
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
PCA
27.40
23.73
26.28
55.01
50.43
22.77
34.27
50.0%
21.69
37.5%
MKV
68.63
33.56
56.26
77.48
59.83
62. 16
59.65
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
E COMPATIBILITY WITH OTHER KV CACHE COMPRESSION TECHNIQUES
In this appendix, we present detailed results and analysis on the combination of Matryohskak with two orthogonal key-value (KV) cache compression techniques: H₂O (Zhagg et al. (023b) and KIVI (Liu et al. 2023). Both combinations evaluated on the datasets mentioned in Section 5.] and their performance is summarized Table 5 and Table 6 respectively.
19

---

Published as a conference paper at ICLR 2025
Table 5: Results of Combination of Distilled MatryoshkakV Projections and H₂O 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 Matryoshkaky uses 50% cache on the feature dimension axis, the
overall cache utilization is 10%.
H20
MKV
LongBench
HLSG
ARC-C
ARC-E
PIQA
WG
CSQA
Avg.
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
$2.20
76.28
62.12
65.27
60.67
100%
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
4.93
66.62
34.24
52.91
75.46
62.00
37.5%
62.38
32.20
50.26
55.28
25.0%
7.66
51.91
27.46
44.44
69.64
54.54
44.39
48.73
100%
4.18
70.71
36.61
52.38
76.35
60.54
66.50
60.55
87.5%
4.44
71.42
35.25
$3.09
76.33
59.91
641.2
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
$7.30
61.43
57.46
32.88
49.21
73.01
55.12
7.67
44.55
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
$3.79
75.41
54.54
62.00
$8.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%
29.83
49.91
73.61
54.14
7.71
26.44
43.24
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
$2.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
25.0%
43.22
21.02
31.92
63.93
33.09
F COMPARISONS WITH MORE BASELINES
We introduce an additional baseline, ASVD(Yuan et al. 2024), which has been developed to address the low-rank characteristics of L.MM parameters. This approach performs simultaneous compres- sion of both the KV cache and the model parameters, allowing for efficient utilization of memory resources. ASVD provides checkpoints for three specific cache budgets: 85%, 90%, and 95%. In our experiments, we compare our MatryoshkaKV against ASVD under these budgets to evaluate performance and efficiency. The results of these comparisons are detailed in Table 7 where we present the performance metrics for our method alongside those obtained using ASVD.
We evaluate the inference speed of our LLM equipped with MatryoshkaKV and compare it to the -LaMA2-7B-har model. The results are displayed in Table [8, Specifically, these evaluations were conducted during the inference process with a batch size of 32. Our current implementation con- sumes a slightly faster time than the bascline full-KV model. This is because we have not performed system-level optimizations for memory copy and sparse computations involved in our KV mecha- nism.
20

---

Published as a conference paper at ICLR 2025
Table 6: 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 sim- plicity.
Model
Budgel
Method
HLSG
ARC-C
ARC-E
PIQA
WG
CSQA
MKV
36.95
53.26
76.39
67.08
100.0%
MKV+KIVI
69.76
35.93
51.98
76.55
61.48
66. 26
60.49
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
MKY MKV+KIVI
69.30
33.90
54.67
75,90
61.09
63.23
60.08
LLaMA2
75.0%
63.23
MKV
36.27
64.46
62.5%
MKV+KIVI
66.56
35.25
51.68
75.41
59.43
61.59
58.33
MKV
65.08
33,56
52.03
74.81
57.54
60.36
56.98
MKV-
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
$4.09
MKV+KIVI
57.11
28.81
48.85
71.71
50.37
MKY
50.61
25.76
45.33
69.64
54.30
43,90
47.96
25.0%
MKV+
48.12
27.80
42.86
67.85
53.59
40.54
46.78
Table Comparison between our MatryoshkaKV and bascline ASVD. We use uniform compression levels for inference here for simplicity.
Model
Budget
Method
HLSG
ARC-C
ARC-E
PIQA
WG
CSQA
Avg.
100.0%
haseline
74.00
35.93
50.97
78.50
61.64
65.93
61.16
ASVD
71.12
36.95
52,20
76,28
62.35
66.67
60.92
LLaMA2 7B base
MKV
72.59
36.27
53.09
76.44
62.43
66.75
61.25
ASVD
70.45
34.92
52.03
75.63
61.72
64.70
60.06
90%
MKV
36.61
62.05
ASVD
67.23
35.93
62.16
59.29
85%
MKV
72.33
35.93
53.26
76.33
61.80
64.78
61.13
G EXPERIMENTS ON VARIOUS HYPER-PARAMETERS
In Section 5.1 during the training process, we initially predefine the schedule set as (gd) Sub- sequently, we modify the schedule set {1d) while keeping other hyper-parameters unchanged. Then, evaluate the accuracy using the same benchmarks. The results are listed in Table 9
21

---

Published as a conference paper at ICLR 2025
Table 8: Tokens per second at different percentages.
LLaMA2
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
Table 9: Accuracy of our MatryoshkaKV after CPT on six benchmarks. We use uniform com- pression levels for inference here for simplicity. Different hyper-parameters are compared. In the table we donate the schedule { as M2. and the schedule {id) as MI. We use uniform compression levels for inference here for simplicity.
Model
Budget
Method
HLSG
ARC-C
ARC-E
PIQA
WG
CSQA
Avg.
Mi
72.03
36.61
52.56
76.71
61.64
67.16
62.07
100.0%
M2
72.05
37.29
52.38
76.66
61.72
67.32
61.24
M1
72.03
37.29
53.09
76.28
62.75
65.77
62.18
87.5%
M22
60.67
M1
34.92
65.03
LLaMA2
75.0%
M2
70.98
34.58
55.20
76.77
61.56
63.64
60.46
7B-base
M1
69.03
32.88
52.91
74.86
59.19
64.54
59.69
M2
69.22
37.29
55.73
75.22
59.35
64.21
60.17
M1
66.34
32.88
$3.09
74.97
58.25
62.49
$8.59
50.0%
Ma
52.91
62.00
58.27
37.5%
M1
61.55
31.19
49.91
73.83
56.27
52.09
53,78
M2
62.38
32.20
50.26
73.34
56.67
55.28
55.02
M11
50.91
26.10
44.97
68.39
52.72
38.33
46.38
25.0%
M12
51.91
27.46
44.44
69.64
54.54
44.39
48.73
The final results our MatryoshkaKV are not very sensitive to the schedule choice.
22


Reviews:
[{"review_id": "Dl21cA2RKT", "reviewer": "Reviewer_LgQe", "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.", "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.\n\n2. How is $\\Delta𝑟$ determined in the greedy search algorithm? And what values are used in the experiments? \n\n3. In Fig.4, MKV seems to work well with uniform compression rates. Does this always apply to all tasks? \n\n3. Line 96: k→r?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.\n\n2. Greedy search algorithm effectively adapts to differing sparsity in each \n$𝑊_𝑘$ and $𝑊_𝑣$ matrix, showcasing flexibility in compression rates across layers.\n\n3. 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).\n\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JpQn5pxRTj", "reviewer": "Reviewer_XwVW", "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.", "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?\n2. 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is easy to follow.\n2. 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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0vf8hoA4cq", "reviewer": "Reviewer_uec1", "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.\n\nThe 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.\n\nTo 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.", "questions": "Why does the Matryoshka training strategy work better than static compression ratios? What's the theoretical justification?\n\nHow sensitive is the method to the choice of sampling schedule for compression rates during training?\n\nWhy is PCA initialization critical for convergence? Could other initialization strategies work?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 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.\n\nThe 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\n\nWhile 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\n\nLimited evaluation on very long sequence tasks where KV cache compression would be most valuable", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ucGZ5DXXs4", "reviewer": "Reviewer_76BC", "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.", "questions": "see weakness", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 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.\n\n- 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.\n\n- 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.\n\n- 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? \n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "14VFq4tlaP", "reviewer": "Reviewer_9USD", "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.", "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.\n2. Which calibration dataset was used for the results presented in Table 1? \n3. 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.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "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.\n2. 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.\n3. 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].\n4. Results on the more recent family of Llama models (Llama 3/3.1) are missing but crucial to establish the effectiveness of this approach.\n\n[1] Saxena et al., \"Eigen Attention: Attention in Low-Rank Space for KV Cache Compression.\", ArXiv 2024.\n\n[2] Yuan et al., \"ASVD: Activation-aware Singular Value Decomposition for Compressing Large Language Models\", ArXiv 2024.\n\n[3] Zhang et al., \"H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models\", NeurIPS 2023.\n\n[4] Liu et al., \"KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache\", ICML 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)