X-EcoMLA: Upcycling Pre-Trained Attention into MLA for Efficient and Extreme KV Compression
Authors: Guihong Li
Mehdi Rezagholizadeh
Mingyu Yang
Vikram Appia
Emad Barsoum
Venue: COLM 2025 (2025)
Keywords: MLA
Multi-head Attention
LLM
Efficient
TL;DR: Upcycling Pre-Trained Attention into MLA for Efficient and Extreme KV Compression

Abstract: Multi-head latent attention (MLA) is designed to optimize KV cache memory through low-rank key-value joint compression. Rather than caching keys and values separately, MLA stores their compressed latent representations, reducing memory overhead while maintaining the performance. While MLA improves memory efficiency without compromising language model accuracy, its major limitation lies in its integration during the pre-training phase, requiring models to be trained from scratch. This raises a key question:  can we use MLA’s benefits fully or partially in models that have already been pre-trained with different attention mechanisms? In this paper, we propose X-EcoMLA to deploy post training distillation to enable the upcycling of Transformer-based attention into an efficient hybrid MLA variant through lightweight post-training adaptation, bypassing the need for extensive pre-training. We demonstrate that leveraging the dark knowledge of a well-trained model can enhance training accuracy and enable extreme KV cache compression in MLA without compromising model performance. The experimental results show that our proposed method can effectively compress the KV cache while preserving the performance on the benchmarks; specifically, for Llama3.2-1B-Instruct baseline, a 6.4× compression achieves the same average score by using only 3.6B training tokens and 70 GPU hours on AMD MI300, whereas a 10.6× compression have less than 0.1\% average score drop with 7B training tokens and 140 GPU hours. The code for this work is available at \url{https://github.com/AMD-AIG-AIMA/AMD-Hybrid-Models}.

Published as a conference paper at COLM 2025
X-EcoMLA: Upeycling Pre-Trained Attention into MIA for Efficient and Extreme KV Compression
Guihong Li", Mehdi Rezagholizadeh" Mingyu Yang,, Vikram Appia, Emad Barsoum Advanced Micro Devices, Inc. (AMD) {guihong. 1i, mehdi. rew.ghozzzhhooooo zadeh, minddeh,miigyie yanggggggguuyyyyyyyoo.g.go (aang)aamd.. com
Abstract
Multi-head latent attention (MLA) is designed optimize KV cache memory through low-rank key-value joint compression. Rather than caching keys and values separately, MLA stores their compressed latent representations, reducing memory overhead while maintaining the performance. While MLA improves memory efficiency without compromising language model accuracy, its major limitation lies in its integration during the pre-training phase, requiring models to be trained from scratch. This raises a key question: can we use MLA's benefits fully or partially in models that have already been pre-trained with different attention mechanisms? In this paper, we propose X-FcoMI.A to deploy post training distillation to enable the upcycling of Transformer- an efficient hybrid MLA variant through lightweight post-training adaptation, bypassing the need for extensive pre-training. We demonstrate that leveraging the well-trained model can enhance training accuracy and enable
attention into
dark knowledge extreme KV cache compression in MLA without compromising model performance. The experimental results show that our proposed method can effectively compress the KV cache while preserving the performance on the benchmarks; specifically, for Llama3.2-1 B-Instruct baseline, a 6.4x compression achieves the same average score by using only 3.6B training tokens and 70 GPU hours on AMD MI300, whereas 10.6x compression have less than 0. average score drop with 7B training tokens and 140 GPU hours. The code for this work is available at ht aps -M
/ //github. cwm AMD-AMDDAAAAAAA AIMA/ HA
Introduction
Large language models (LLMs) have become ubiquitous, revolutionizing both academic research and industrial applications Brown et al. (2020); OpenAl (2023); Chowdhery et al. (2022); Liu et al. (2024a); Guo et al. (2025). Their success largely stems from pre-training and instruction tuning vast amounts of data, the power of self-attention in Transformer architectures, and the computational capabilities of GPU accelerators.
Despite their widespread adoption, Transformers face two major challenges in their self-attention mechanism: quadratic computational complexity and high memory demands for KV cache storage, particularly when processing long sequences. Significant efforts have been made to mitigate these challenges by: (a) replacing self-attention with new sub-quadratic architectures such as state-space models (SSMs) Gu & Dao; Dao & Gu (2024); Poli et al. (2024); (b) enhancing the efficiency of existing Transformer self-attention mechanisms Arora al. (2024); Ainslie et al. (2023); Zhang et al. (2024b); Yang et al. (2024); Qin et al. (2024); and (c) developing hybrid solutions that combine advantages of both quadratic and sub-quadratic models Lieber et al. (2024); Dong et al. (2024); Wang et al. (2025); Bick et al. (2024). Among these approaches, the second category is particularly appealing, as it requires minimal architectural changes while capitalizing on existing hardware optimized for Transformers. This work focuses on improving the efficiency of existing self-attention mechanisms within this category.
Multi-heat attention (MHA) is fundamental mechanism in Transformer architectures. However, during inference, MHA requires the saving of large amount of key-value (KV) cache, resulting in high
Equal Contribution First Authors, with order determined alphabetically.

---

Published as a conference paper at COLM 2025
Our X-EcoMLA trained on -78 tokens
Our X-EcoMLA trained on -3.68 tokens
55-
16 Teacher Toacher
10 Texcher
30 Taxcher
teacher
BBl leacher
Modell (52.77)
Rase Model (52.77)
Average Accuracy
12.0x
10.6x
6.4k
3.6x KV Cache Compression Ratio (times)
KV Cache Compression Ratio (times)
Figure 1: X-EcoMLA with Different Teacher Sizes: [Right] Our results show that using Llama3.2-1 2-B. 3B, and 8B teacher models enables KV cache compression of Llama3.2- B-BBBBBBBBBBB by 1.9x, 3.6x, and 6.4x using 3.6B tokens respectively, without compromising average accuracy across multiple tasks on the Harness Evaluation benchmark. [Left] With 7B training tokens, we can further compress the KV cache to 10.6x and 12.8x, while maintaining competitive accuracy.
memory consumption. To address this challenge, DeepSeek (Liu et al., 2024a;b) recently proposed multi-head latent attention (MLA), a novel approach that compresses the KV cache while maintaining the performance of LLMs. The mainstream training paradigm for MLA has relied pre-training from scratch using vast amounts of data and computational resources. For example, pre-training the Deepseek-v3 model required 2.664M GPU hours Nvidia H800 clusters. This highlights major challenge: developing models with new attention mechanism demands significant computational resources during pre-training. Given the substantial effort already invested in training Transformer models, natural question arises-can we transfer the rich pre-training knowledge from trained LLMs into more efficient MLA models without training them from scratch?
Existing approaches showcase strong evidence supporting the feasibility of knowledge transfer with architectural adaptations for LLMs, such as MambaInLlama (Wang al., 2025), MOHAWK (Bick et al.. 2024), and HedgeHog (Zhang et al., 2024b). Inspired by these solutions, we introduce X- EcoMLA, a cost-effective knowledge transfer approach designed upcycle pre-trained multi-head attention into MLA. In X-EcoMLA, we initialize MLA from its corresponding pre-trained attention blocks using our static or dynamic SVD approach, followed by distillation from a well-trained teacher model. By leveraging the dark knowledge of a high-quality model, enhance training accuracy and achieve extreme KV cache compression in MLA without sacrificing performance. Our results demonstrate that an 8B teacher model enables 6.4x KV cache compression of the Llama3.22 1B- Inst baseline while preserving 100% of its average score across multiple tasks on the LM Harness Evaluation benchmark. This requires only 3.4B training tokens and about only 70 GPU hours AMD MI300, while pre-training the Llama3.2-18 1BB model requires 370K GPU hours! Furthermore, we achieve a 10.6x compression with 7B training tokens and around 140 GPU hours, with less than 0.1 % average score drop. We summarize our major contributions as follows:
We propose a lightweight post-training approach to upcycle pre-trained attention to MI.,, significantly reducing computational costs by eliminating the need for training from scratch.
We develop static and dynamic SVD-based initialization techniques to improve the conver- gence and accuracy of MLA adaptation.
We demonstrate that leveraging a larger teacher model enables extreme KV cache com- pression while maintaining model performance, achieving up 10.6x compression with minimal accuracy loss. (see Fig. 1).
We validate the effectiveness of our approach through extensive experiments the LM Harness Evaluation benchmark, showcasing its efficiency across various settings and LLMs.
2

---

Published as a conference paper at COLM 2025
2 Related Work
This section gives a brief overview of related work; detailed version is available in Appendix A.
KV Cache Management in Transtormers Transformers store key-value (KV) vectors for each token and each layer during auto-regressive generation, leading to high memory requirements during inference, especially for long sequences. Several methods have been proposed to reduce the KV cache size in Transformers, broadly divided into training- based and post-training approaches Shi al. (2024). Post-training techniques are easier to apply but may degrade performance due to information loss. These include KV eviction (e.g., Heavy Hitter Zhang et al. (2023)), sliding window attention Arora et al. (2024); Beltagy et al. (2020), low-rank projection approach for KV-cache compression Chang et al. (2025), and quantization Kang et al. (2024); Zhang et al. (2024a). Some strategies, like Attention Sink Xiao et al. (2023) and KV merging Wang et al. (2024), aim to mitigate information loss while keeping memory usage low. This paper, however, focuses training- solutions, which tend to offer a better trade-off between efficiency and accuracy.
Training- based KV Cache Management Training- based methods modify attention mechanisms to reduce memory use during inference. Multi-query attention (MQA Shazeer (2019) and grouped- query attention (inA))inslii et al. (2023) reduce KV cache size by sharing keys/values among heads. YOCO Sun al. (2024) reduces redundancy with a shared KV cache across layers. DeepSeek- V2 Liu et al. (2024a) introduces multi-head latent attention (MLA), which compresses hidden states via low-rank projection, reducing cache size while outperforming standard MHA. Inspired by MLA's efficiency, this work explores adapting MLA to traitedined model and try address this question: Can we upcycle pre- pre-trained models to their MLA counterparts without costly retraining?
Upeycling Attention Upcycling refers to upgrading pre- pre-trained models with minimal computa- tion Komatsuzaki et al. (2022). attention upcycling, existing attention blocks are adapted into efficient forms like MLA without full retraining. GQA Ainslie et al. (2023) and Hedgehog Zhang et al. (2024b) achieve this via light fine-tuning or distillation. Hybrid models like MambalnL- lama Wang et al. (2025) and MOHAWK Bick al. (2024) distill knowledge from Transformer attention into Mamba layers. MHA2MLA Ji et al. (2025) introduces a data-efficient fine-tuning approach for converting MHA to MLA via partial RoPE removal and joint SVD-based low-rank approximation. In contrast, X-EcoMLA adopts a unified RoPE design with a shared Key-RoPE vector across all heads (similar to DeepSeek MLA), and employs structured initialization along with mhowledg-disiillaion--ass efficient training to enable effective MLA-based upcycling.
3 Background
In this section, we formalize the mathematical framework of MHA and MLA, following the notation from the original DeepSeek-V V2 technical report Liu al. (2024a) with some slight modifications.
3.1 Multi-Head Attention (MHA)
MHA projects the input hidden state H into three distinct spaces using learned weight matrices:
Q = =HW == HW V= -HWV,
where H E RR is the input sequence representation with / being sequence length and d internal hidden state dimension, and WΩ, WK, w e RRdxnjda are the learnable projection matrices where n1 is the number of attention heads, and d, is the head dimension. The attention scores and final outputs computed as:
QKT
A = Softmax
O = AVWO
where Wo E Rdddd is the output transformation matrix. During inference, MHA requires caching K and V for all past tokens, leading to storage requirement of 2n,d,1.
3.2 Multi-Head Latent Attention (MLA)
MLA introduces a low-rank joint compression strategy for keys and values, reducing the KV cache size. Instead of caching K and V. MLA compresses them into lower-dimensional latent representa-
3

---

Published as a conference paper at COLM 2025
tion CKV:
CKV = HWDKV
(3)
where WDKV E RR is the down-projection matrix, and Pkv < d2 is the compressed dimension for keys and values. The keys and values are then reconstructed from CKV using:
KC C.CWKUUK = CkVW
where WUK R and Wuv Ro are the up-projection matrix for keys and values. Notice that in this paper we consider more flexible setting where the queries and keys could have different dimensionality dgt other than d,. During inference, the learned matrices can be absorbed into the existing projection layers: WUK and WuV can be merged into WQ and WO respectively.
In the meantime, such low-rank compression can be also applied to the queries to reduce the memory usage while training:
c HWPQ, Q cQwue
(5)
where Cº E R represents the compressed latent vector for queries, Vq denotes and query
compression dimension, WDQ E Raxyy is the down-projection matrix and WUQ E RT xnxd
denotes the up-projection matrix.
However, such low-rank KV compression is not compatible with Rotary Position Embedding (RoPH) as it breaks the matrix absorbing mechanism. As result, the authors in Liu al. (2024a) propose decoupled Rotary Position Embeddning where additional multi-head queries Q and a shared key K are applied to carry RoPE, which can be expressed
Q RO K (HWK))
(6)
where WQR Regxn and WKR R represent the matrices to produce the decoupled keys and queries. Then, the RoPE embeddings Q KR and Non-RoPE (NoPE) embeddings Q.,, K are concatenated to perform the attention operation:
Q 10) K repeat((K
where repeat(.) denotes duplicating KR for each head. After concatenation, the same attention operation is applied in Equation 2. During inference, MLA requires caching only CKV and KR. reducing the storage requirement to (Fkv + d.)l which significantly smaller than the standard MHA cache size.
4 Methodology of X-EcoMLA
We begin with pre-trained Transformer model referred to as the "base model".. Our methodology in this paper concerns upcycling the attention modules in the base model into MLA modules to save the KV cache memory while remaining minimum training efforts and performance degradation. To achieve that, we first propose our SVD-based weight initialization approach to better inherit the knowledge from the preitrained model. Additionally, our initialization approach offers both static and dynamic rank selection. After initialization, we adopt the knowledge distillation training process in ManbaInLlama Wang al. (2025), which includes: end-to-end knowledge distillation, and direct preference optimization (DPO).
4.1 SVD-based Weight Initialization
As introduced in Section 3.2, the MLA module comprises several key parameters: WPQ, WuO, WOR WDRV WRR WUK and WUV Correct initialization these matrices is crucial to ensure a smooth transition and preserve much of the original model's knowledge as possible. For other parameters such as the output transformation matrix Wo and the feedforward module, directly copy them from the pre-traine pre-trained base model as initialization.
Although MLA and MHA are fundamentally different, MLA closely approximates low-rank version of MHA if we disregard the positional embedding and intermediate layernorms. Based on this observation, we propose an SVD-based initialization method that initializes the MILA weights using
4

---

Published as a conference paper at COLM 2025
SVD- deD-dcoopoosed low-rank matrices derived from the pre- pre-trained attention weights. Our method is simple and can be illustrated within few lines of pseudocode as shown in Algorithm As demonstrated in our experiments, such straightforward initialization could significantly enhance the knowledge distillation performance compared to random initialization.
For simplicity, we assume the pre-trained attention block employs MHA with weight matrices Wº, WK W E Rdxx,,", However, method can be readily extended to MQA Shazeer (2019) and GQA Ainslie et al. (2023) by duplicating the key and value weight matrices to align with the total number of attention heads. Our approach supports varying query and key dimensions under the constraint d = d, S dz. For the value dimension, we assume that the MI.A and MHA modules share the same dimensionality di.
To initialize wPQ, w and WQR, we begin by performing SVD on the pretrained weight matrix W to obtain:
We = Ug Vg
(8)
where Up E RRax,, Eg E R and Vg E XTo, For the down-projection matrix WPQ, we directly set WDQ = Ug. For the up-projection matrices, we first compute I V/ and reshape it to W R Then we partition it along the last dimension derive the two up-projection matrices: WuQ for the first dgt elements and WOR for the last d,. This can be expressed as:
W = reshape reshape W [:::;: dqil), W@R = reshape W UQR [,, - dy =]),
where reshape(.) denotes the reshape function that merges the last two dimensions of the input tensor.
For the remaining MLA weight matrices associated with keys and values, we first perform a joint SVD on the concatenated W and WV: [WK, Uxo Vir
(10)
where Uho E Rd Eky RT and V/DU E R For the down- pwonn-pojection matrix WDKV, we directly set WDKV Ugp- To derive the up-projection matrix Wuv we set it equal to the last di columns of Ekz V/ For the up-projection matrix for keys WUK we first extract the
first d, columns of Eko V%o and reshape them into W E R Xmg Subsequently, we select the first dz elements along the last dimension of w and reshape back to obtain WUK which can be expressed as:
WUK reshape W UK
[,,:; dg:l).
(11)
Lastly, for the RoPE key embedding matrix WKR we employ different initialization strategy as all attention heads share the same RoPE key embedding in MLA. We first compute the average key projection matrix W augg E dd across all attention heads. Then, we extract the last d, columns from initialize WrR. which can be expressed as: WKR
WK Wang -d, :].
(12)
Clearly, it is crucial to determine the rank values r and Iee for the Q and KV matrices. We propose two methods for the rank selection: (i) Fixed Rank Selection, where constant rank values are used across all transformer blocks. (ii) Dynamic Rank Selection, which determines rank values based on two predefined energy thresholds &g, 6xv E 1). Taking Wo as an example, we use SVD to decompose it and obtain the singular values:
Oj j E [dnn nndz)], where Tj 2 0j+1
(13)
To determine the optimal rank. we define the cumulative energy based the squared singular values. The rank rq for the i-th transformer block is selected as follows:
where E = miníd, (gda))
rq = arg min
E
Σ of
(14)
R
This approach ensures that the selected rank captures at least &, fraction of the total energy E. The same energy-based dynamic rank selection process can be applied the KV weight matrix as well.
5

---

Published as a conference paper at COLM 2025
4.2 Training Process
End-to-End Knowledge Distillation The primary training stage involves an end-to-end distillation using a Supervised Fine-Tuning (SFT) dataset introduced in Wang al. (2025). During this stage, the goal is to minimize the KL divergence loss between the outputs of the student model (i.c.. X-EcoMLA) and the teacher model, which can be expressed as:
x,xxr))l((1lymmil Co (
(15)
where f is the trainable parameters of the student model and 0y is the parameters of the teacher model, which is kept frozen while the distillation process. Such distillation step is crucial for transferring the rich, pre-trained knowledge from the teacher model. In Section A.5.I, we show that this distillation process more effective than plain cross-entropy loss. Note that the teacher model can be any strong, pre- pre-trained model-ever construct the student model.
that is different from the base Transformer model that is used to
Direct Preference Optimization (DPO) In the final training stage, we perform DPO which is a binary cross entropy loss to adjust the preference of the student model. Following Wang et al. (2025), we set the distilled student model itself as the reference model as it makes the training much stabler and produces a sufficient performance gain, which can be observed from the experiments section.
5
Experiments
5.1
Experimental Setup
Model In this paper, we primarily focus on SmolLM-series models 2 SmmoLLM-13MMIIsstrue MoolL.--300M-Instruee
and Llama 3-series models 3 (Llama3.2- I-B-Ittttt 1-B--Instrct) as our base models which span a variety of scales. All the models employ GQA as their attention module. For MLA, we utilize the same number of attention heads and the same head dimension for values. We tried multiple rq. Tku and (d,) while maintaining similar amount of parameters as the base model.
Llama3.2-3B- Instructt,
As mentioned earlier, we employ knowledge distillation, transferring knowledge from teacher models to the base model. In our experiments. we explore the impact of different teacher model sizes, including Llama3.2 B---IIIrucct Llama3.2-31 2--..--BB-nsstueee and Llama3. Instructt Grattafiori et al. (2024). This approach allows us to analyze how the varying sizes of the teacher model impact the performance and efficiency of the distilled student model.
Training procedure Our approach begins with a pre-trained model, so we focus primarily on post-tral postttraaning rather than pre-training the base model again. We two-stage training procedure. Specifically,, in the first stage (knowledge distillation), our X-EcoMLA model is trained with SFT using the KI. loss between the output of the base model and the teacher model. We use AdamW optimizer with B = (0.9, 0.98) and set the training batch size at 96. We use the same dataset as used in the paper Wang al. (2025). It is constructed by aggregating and reformatting data from multiple publicly available sources, including OpenHerrees-2.5 Teknium (2023), GenQA Chen et al. (2024), and Infinity-Insrrue of Artificial Intelligence (BAAI). For the second stage (DPO), based on finetuned models, we use the combination of three datasets for DPO preference tuning: Llama3- ultrafeedback Wang (2024), orca dpo pairs Lian et al. (2023), and ultrafeedback binarized Cui et al. (2023). We set the training batch size at 64 and still use A W as the optimizer. We unfreeze all parameters in our X-EcoMLA model and train them for one epoch for both SFT and DPO.
5.2 Results
Similar to MambaInI Wang et al. (2025), we adopt the LM Harness Eval benchmark Gao et al. (2023) (branch big-refactor) perform zero-shot evaluation on 9 different tasks: ARC-Challenge (HS) Zellers et al. (2019), MMLU (MM) Hendryeks al. (2020), OpenBookQA (OBQA) Mihaylov et al. (2018), PIQA Bisk
(ARC) Clark et al. (2018), ARC-Easy (ARE) Clark et al.
3 http://hhh.....oooooo////o///o//

---

Published as a conference paper at COLM 2025
Model and Seiting
Init. Method
KV-Sies
ARC
ARE
HS
MMLU
OBQA
PIQA
PBMD
RA
WG
Avg.
SmelLMM33M Inst
Base Model_
100%
27.39
43.64
41,91
24_11
33,80
67.30
55.80
32,06
55..2
41.91
EcoMLA
Random (Pes
160)
50%
21.08
30.68
25.27
22.95
25.00
53.54
55.20
22.87
49.57
34.03
1X-EcoMLA DPO
Random (ri=
160)
50%
21.59
30.56
25.
22.95
26.20
53,97
55.20
22.97
49.33
34.22
1X-FeaMIAA
Fixed (ru 1661)
27.05
43.31
40.
24.51
80
51.30
41.31
[X. XX-EcoMLA DPO
Fixed (Ne 160)
50%
28.24
45.33
40.26
26.20
33.60
65.89
55.60
32.63
50.36
42.01
1X-EcoMLA
Dynamic (App = 0.83)
49.5%
26.88
43.52
39.84
24.50
31.80
67.08
55,00
31.87
$1.62
41.35
1X-EcoMIA
Dynamcc (hu
27.99
25.84
33.40
52.73
50.28
41.86
Buse Model
100%
32.76
56.
$2.74
24.37
37,00
70.51
55.40
33,55
54.06
46.28
EX.EaaMiA
Random (ra
288)
50%
22. 10
56.36
26.72
23.22
25,00
55.60
$5.20
25.25
448
53.00
1X-EcoMLA DRO
Random (rks
288)
50%
22.61
35.61
27.12
23.19
56.26
54,80
23.64
49.57
35,49
1X-ExoMLA
Fixed (Yee 288)
50%
33.11
54.55
50,77
23.26
36.60
71,06
55.40
31,96
54. 14
45.65
XX-FeMMIA DPO
Fixed (fg. 288)
50%
34.39
55.77
$1.93
25.68
38.60
69.37
55,20
33.88
53. 53.35
46.46
IX-EcoMLA
Dynamic (6.
0.88)
49.3%
32.85
54.38
50.96
23.47
37.40
31.67
54.62
1X-EcoMLA + DPO
Dynamic (6.,
0.88)
49.3%
34.64
54,97
51.75
25.45
38. 80
69.59
55,20
33,88
54.06
46.48
Llama3.2-11 1-n-
Bice Model
10056
37.97
63,51
60.77
46.09
35.00
74,37
60.20
38.09
59. 67
52.85
1X-EeaMLA
Random (/zs 512)
53.1%
35.32
60.48
54.03
27.77
35.20
71.98
55,80
33.88
55.01
47.72
1X-FoaMIAA DPO
Random (/2- 512)
53.1%
38.99
62.71
56.
23.04
$6.40
36.27
56.
49.44
EX.ELoLLA
Fixed ine 512)
53.1%
36.95
63.89
58.88
43.40
36.00
74,16
58.20
37.32
60.30
52.12
1X-EcoMLA DPO
Fixed the = 512)
53.1%
39.93
63.89
60.73
42.39
37,80
74.92
58.80
10.77
60.54
53.31
1X-EcaMLA
Dynamic (h 0.95)
37.12
43.
413
37.51
60.
[X-EooMLA DPO
Dynamic (A
0.93)
54.7%
41.21
64,86
60.96
42,86
37.60
74.43
$8.60
39.23
58.33
53.12
Llama3,2-3B-BR
Base Model
100%
45.90
67.76
70.36
60.46
36,20
75.57
69.60
40.77
67.17
59.31
1X-FooMLA
Rundom (7
816)
47%
67.34
65.11
46.97
36,40
75.24
39.04
63.69
55.23
1X Xc-EcoMAA+ DPO
Random 0g2 816)
43%
44.80
69.78
67.01
47.37
38.80
76.01
62.60
40.00
64.50
56,80
1X-ExoMLA
Fised (re = 816)
43%
43.09
67.76
69.54
56.96
37.00
75.84
66.00
41.34
67.48
58.33
(X-FooMIA DRO
Fixed (f 816)
43%
48:21
70.45
72.24
57.42
76.55
66.80
46.22
68.
60.54
IX-EcoMLA
Dynamic (6.,
0.95)
43%
42.75
66.41
69.59
57.47
36.80
75,46
67.60
42.30
68.03
58.49
EX-EcoMLA+ + DRO
Dynamic (6. = 0.95)
43%
18:16
69.99
72.26
57.73
39.40
75.79
68.10
46.32
65.90
60.47
Table I: sero-shot evaluation of self diftidistleed X- EcoMLA with different initialization methods (random, SVD with fixed/dynamic rank selection) and base models the LM Harness Eval benchmark across nine tasks: ARC- Challenge (ARC), ARC-Easy (ARE), HellaSwag (HS), MMLU, OpenBookQA (OBQA), PIQA, PubMedQA (PBMD). RACE (RA), and WinoGrande (WG). (f denotes upcycling the "Base Model".)
et al. (2020), PubMedQA (PBMD) Jin et al. (2019), and RACE (RA) Lai et al. (2017), WinoGrande (WG) Sakaguchi et al. (2021).
Self-distillation Evaluation Table 1 shows the benchmark performance of our proposed X- EcoMLA when we the base models themselves as the teacher model, which we refer to as self-distillation. We evaluate three different initialization settings: (i) Fixed rank selection with random initialization, (ii) Fixed rank selection with SVD initialization, and (iii) Dynamic rank se- lection with SVD initialization. For all experiments, we set dzk = d, = 32 and adjust rg. Tko 8q. and 6Gp accordingly to achieve around 50% KV cache compression while ensuring the total number of parameters after the MILA upcycling remain roughly the same. For each scenario, we evaluate training with the full dataset (6.8B tokens for SFT and 0.2B tokens for DPO) and it is observed that DPO significant boosts the distillation performance. For fixed rank selection schemes, it is evident that SVD initialization significantly enhances distillation performance compared to random initialization, yielding 22.8% and 30.91% improvements for SmolLM models and 8.1 % and 6.5% improvements for Llama 3.2 models. Such observation demonstrates the necessity of applying our SVD initialization method to inherit the knowledge from the pre-trained targe models. For dynamic rank selection schemes, we observe similar performance the fix rank selection in most experiments, which demonstrates its effectiveness. Note that our X-EcoMLA + DPO could always achieve better performance than the pre-trained base models with only 43% 54.7% KV cache sizes.
Extreme KV Cache Compression with Larger Teacher Table 2 presents the impact of KV Cache compression on model accuracy across various benchmarks. For more details, please refer to Appendix A.4.4 and Table 11 11. We adopt Ilama3.2 Instruct as the base model and progressively reduce the KV cache size of our X-EcoMLA from 53.1% to 7.81%. With the same base model as teacher, we observe a consistent accuracy drop across most evaluation tasks, which demonstrates the trade-off between reducing memory consumption and maintaining model accuracy.
However, our results reveal that such performance degradation from extreme KV cache compression can be mitigate if we utilize larger teacher models such Llama3.22 3.B SB-Instruct. For instance, when reducing the KV cache to 15.6%, using Llama3.1-8B -BB-Isstt as the teacher recovers 1.56 of the average score (52.94 vs. 51.38) when trained with half of the dataset
and Llama3.1 1-
7

---

Published as a conference paper at COLM 2025
Model and Setting
Teacher
Param
Tokeas
ARC
ARE
HS
MMLU
OBQA
PIQA
PBMD
RA
WG
Arg.
1amm3.2-1B-lam
1,248
37,97
63.51
60.77
46.09
35,00
74.37
60.20
38,09
59.67
52.85
1005: MLAA
Layers U/t
512,
de
32)
53.1%
1X-ReoMLA DPO
Llamm3.2-11--mm
1,238
3.68
39,93
63.51
60.52
41.58
37.20
73.99
59.80
40.48
60.38
53.04
1X RooMLA + DPO
Llama3.2: 2:33 Inst
1.23B
3.6B
42.75
64.81
62.04
43.88
37.40
73.72
59.20
41.44
61.48
54.08
IX-BcoMLA DPO
Llama3.1 1-88118---m
1,238
3.6B
44,03
68.86
63,19
43.81
37,40
73.94
61.40
41.82
61.40
55.13
100% MLA
Layors (hty
256,7,
1184,
4,
KV Size: 28.1%
1X-EcuMLA DPO
Llama3.2-1B IB Inst
1.23B
3.6B
40.02
63.26
58.74
39.79
36.40
72,80
55. 60
40.19
60.38
51.91
1X-BooMLA+ +PO DPO
Llamm3.2 2-3-------88 Iust
1,238
3.6B
40.70
64.35
60.10
41.77
37.20
73.83
57. 30
39.23
61.17
52.91
XX-FoMMI.A DPO
Ilama In8t
1,238
3,68
41.98
41.78
59.00
40.00
60.69
5363
100% MLAA
Layers (7k--
128,
74 1334, d,t.
32)
KV Size: 15.66
IX-EcaMLA DPO
Llama3.2-18 Inst
1.23B
3.6B
39.16
61.83
57.27
37.85
36.20
73.45
56.40
40.19
60.06
51.38
"X-EcoMLA+ DPO
Llama3.2-38- Inst
1.238
3.6B
39.12
62.88
58.41
39.15
37.20
73.39
58.00
39.71
59.75
52.02
1X-KcMMIA DPO
LJama33
1-88-Iost
1,238
3,68
41.30
65.61
64
39.47
(4
24.27
39.53
59.85 83
52.94
X-EcoMLÄ DPO
Llama3.2. 2- list
1.238
7B
40.10
62.68
58.17
39.70
37.80
73.50
56.00
39.33
59.67
51.97
18. A DPO
I 2. Inst
1,238
78
39.33
64.86
58.92
41.86
37.40
73.83
58. 90
39.71
59.59
52,70
XX EcoMLA DPO
Llama3. 3.1 AD Inst
1,23B
7B
42.49
67.13
60.58
42.51
36.60
73.99
59.40
40.38
59.43
53.61
MI
Layors (Pa 64,
1424, 4,p
KV Size: 9.4%
1X-EooMLA DPO
Lluma3.2-18--nnt
1,238
7B
39.16
62.63
56.04
34.90
36.40
22.85
56.40
37.70
58.33
50.49
1X EcoMLA DPO
Llama3.2 Inst
1.23B
7B
37.97
63.55
$6.95
37.54
35.40
72.74
57.00
38.66
59.27
$1.01
1X + DPO
Llama3. 1-8- 1888-Inl
1,238
7B
39.85
67.13
58.45
38.51
37,40
73.83
58.00
39.81
59.27
52.47
100% MLA
Layers (/k,
F, 14400
d,2
32)
7.81
IX-EcoMLA+ DPO
Llama3.2-1B Inst
1,23B
7B
38.48
61.66
55.32
30.62
35,20
72.36
56.60
37.99
59.43
49.74
"X-BooMLA+ + DPO
Alamm3.2-388 388-
1,238
7B
36.18
62.21
55.82
36.41
35.60
72.03
57.00
38.09
60.06
50.38
[X-FoMMLA
Llama Mamma3.1-880 Inst
1,238
7B
36.27
51.35
Table 2: Impact of KV-cache compression and teacher model size on performance. Reducing the KV-ccche size lowers accuracy. but larger teacher models help recover performance. DPO further improves alignment and accuracy. (f denotes upcycling the base model.)
(3.6B tokens). With the larger teacher, our X-EcoMLA achieves even better performance than the pre-trained base model with only 15.6% KV cache size and 3.6B training tokens. As we increase the training tokens to 7B, we could even push the KV cache compression ratio to 9.4% without significant performance degradation compared to the pre-trained base model (52.47 vs. 52.85). These results highlight the effectiveness of leveraging larger teachers and preference tuning to resolve the adverse effects of extreme KV cache compression while maintaining strong accuracy across multiple NLP benchmarrk.
Scalability to Larger Models We have extended X-FcoMLA to two 8B-parameter models (I.lama3- 8B and Llama3. 1-8B). The results in Table 3 demonstrate that even under aggressive KV cache compression, X-EcoMLA maintains performance very close to the full-scale baseline. Notably, at KV-size 10.94%, performance remains nearly on par with the full model.
Model
KV-size
Avg.
ARC
ARE
HS
MM
OBQA
PIQ
PM
RA
WG
Llama3 8B Inst (Base)
100%
65.78
56.66
81.61
75.81
63.82
42.60
78.62
75.00
46.03
71.90
(X-EcoMIA (Ours; Fo. 256)
15.63%
65.16
54.69
75.69
59.20
77.91
74,80
Llama3, 1 8B Inst (Base)
100%
66.63
54.86
79.55
79.23
68.13
43.00
80.90
75.40
44.69
73.88
1X-EcoMLA (Ours; Nop 160)
10.94%
65.85
57.17
77.57
60.13
79.16
76.20
47.85
71.19
Table 3: Performance of X-FcoMLA with 8B models under aggressive KV cache compression
Lystem-Level Inference Metrics We evaluate system-level performance in terms of throughput (sequences/sec) and peak GPU memory usage (GB) for both the baseline model (Llama3. 1-8) and our proposed model. X-EcoMLA-8h, (T 128, 10.67x KV compression), across a range of batch sizes. All experiments are conducted on identical hardware (single AMD MI300 GPU) under consistent settings. As shown in Fig. 2, X- X-EcoMLA-88 achieves approximately 1.7x to 2x higher throughput than the baseline across all batch sizes. In terms of memory efficiency, X-EcOMLA- 8B substantially reduces peak memory consumption. For instance, at batch size 128, Llama3. consumes 143 GB of memory and fails to run larger batches, whereas X-EcoMLA-SDD requires only 28 re-rrpresenting a 5X reduction. These results demonstrate that X-EcoMLA-S0 8B not only maintains strong model accuracy but also delivers significant system- improvements- achieeving higher throughput and dramatically lower memory usage. This makes it well-suited for latency-sensitiv and memory-consee constraieed deployment scenarios.
8

---

Published as a conference paper at COLM 2025
5000
140
120
Paak Memory (GB)
100
Throughput
2000
60
EcooLLA 38 Throughput
1000
X-EcoNLA-BE Hemeeyy
Ueeme3 1 Thro. Thro.ghout
1 89 Memory
20
152
Batch Size
Figure 2: System-level inference performance for Llama3. 1-8B and X-FcoMIAA-BI on 8x AxAD MI300 GPUs. X-FcoMLA enables higher throughput and drastically reduced memory usage across batch sizes. Jlama3.1-80 out-of-memory (OOM) beyond batch size 128, while X X-EcoMLA scales smoothly up to batch size 1024.
Comparison with Existing Post- Tost-Training Low-Rank Methods We conducted experiments to compare solution with two of the most recent SOTA methods -inddss-Mit MHA2MLL Ji et al. (2025) and PALU Chang et al. (2025) and to clarify how X-EcoMLA improves upon these approaches.
Comparison with MHA2MLA Baseline Setup To more comprehensively evaluate the effec- tiveness of X-EcoMLA, we compare it against MHA2MLA Ji et al. (2025) under both continual pretraining and supervised fine-tuning (SFT) settings using the SmolLM 1.7B model. As SmolLM is originally built with MHA-based attention, this evaluation also demonstrates that X-EcoMLA can be seamlessly integrated into existing MHA-based architectures.
In the continual pretraining setting, we evaluate X-EcoMLA on the SmolLM 1.7B base model using 12.5% KV cache size, following the same 6B-token training budget the released MHA2MLA checkpoint'. We compare four configurations: (1) the full-attention baseline (100% KV cache), (2) MHA2MLA continually pretrained, (3) X-EcoMLA pretrained without teacher, and (4) X-FcoMLA with self-distillation. Table 4 (top) summarizes the results where X-EcoMLA without a teacher achieves an average score of 51.94, slightly outperforming MHA2MLA (51.69) by +0.25, When using distilllation the accuracy improves to 52.87, reducing the gap to the full-attention baseline.
In the supervised fine-tuning setting, KV compression ratios: 12.5% and 50%. We use SmolLM -In--nntt as both student and teacher for X-EcoMLA, and follow the joint-SVD Shigher variant for MHA2MLA. As shown in Table (bottom), X-EcoMLA outperforms MHA2MLA across both compression levels. At 12.5% KV size, X-EcoMLA achieves 49.34 vs. 48.19 (+1.15), and at 50%, reaches 50.15 vs. 49.79 (+0.36).
identical SFT data and compare performance under two
We attribute this performance gain to a key architectural difference: the RoPE (Rotary Position Embedding) design. MHA2MLA stores separate Key-RoPE vectors per head. Under a fixed budget 32 RoPE dimensions.
(e.g., 32 dimensions per token), each head receives only Mh EcoMLA adopts the unified RoPE approach used in DeepSeek MI.,, where all heads share a single Key-RoPE vector. This allows each head to fully leverage all 32 dimensions, offering 8x more positional encoding capacity in an 8-head setupacrucial compression.
contrast, X-
crucial for preserving performance under aggressive
Comparison with PALU PALU Chang al. (2025) represents a recent state-of-the-ar approach for low-rank KV-cache compression. It decomposes the linear projection layers into low-rank matrices, caches the compressed intermediate states, and reconstructs the full keys and values on the fly during inference. To evaluate the effectiveness of X- X-EcoMLA, we compare X-EcoMLA against PALU using the Llama3-8B-Instme model.
Following the experimental protocol described in Section 4.2 of the PALU paper, we evaluate zero-shot performance on the used subset of the LM Harness Eval benchmark using PALU's best-
fn1 Smal IM- 1-7
kv_

---

Published as a conference paper at COLM 2025
Method
KV Size
Avg. Acc.
ARC
ARE
HS
MM
OBQA
PIQA
PM
RA
WG
SmolLM 1.7B (Basc)
100%
54.67
46,42
73.48
66.74
27.73
42.00
76.06
62.60
37.03
60,93
MHA2MLA Ji al. (2025)
12.5%
51 199
41.55
69.57
61.43
24.63
39. 39.00
74.70
60.20
35.69
58.41
[X EsoMLA pretrain
480)
12.5%.
51.94
40.19
69.36
62.52
23.79
40.00
75.30
61.40
35.89
59.04
[X-FooMI.A (ri- - 480)
12.5%
52.87
42.41
72.14
62.84
25.55
41.40
75.19
61.40
36.27
58.64
SmolLM 1.7B-Ins (Base)
100%
39.80
56.83
MHA2MLA (SFT)
12.5%
48.19
35.07
60.94
56.18
23.36
36.40
72.36
57.20
34.83
57.38
1X-FooMI.A (SFT, 2he 480)
12.5%
49.34
37.29
62.96
59,05
23.68
38.00
72.52
59.80
34,64
56.12
MHA2MLA (SFT)
50%
49.79
37.63
62.29
59.60
24.09
38. 28
74.05
60.40
34.35
57.46
1X-FooMI.A (SPT, 24 2016)
50%
50.15
37.97
64.06
60.23
24.89
39.00
73.34
60.20
34,93
56.75
Table 4:
Comparison between X X-EcoMLA and MHA2MLA a under both continual pretraining (top) and supervised fine-tuning (bottom) settings on the SmolLM 1.7B model.
performing configurations: G-LRD and J-LRD, both operating at a 50% KV compression ratio. Reported PALU results taken directly from their paper. The results. summarized in Table 5, show that X-EcoMLA achieves an average accuracy of 67.34 at only 15.63% KV cache size, This performance is nearly par with the full model (67.87) and clearly outperforms both PALU-J-LRD (66.19) and PALU-G-LRD (64.45), with respective gains of +1.15 and +2.89 points-deppll using approximately 3x less KV storage.
Model
KV-Size
Avg Acc.
ARC
ARE
HS
OBQA
PIQA
WG
Llama3-8B-Ins (Base)
100%
67.87
56.66
81.61
75.81
42.60
78.62
71.90
PALU-J-LRD
50%
66.19
51.96
79.63
73.20
43.40
76.50
72.45
PALU-G-LRD
50%
64.45
48.99
76.30
70.36
42.60
76.06
72.38
AX-FcoMIAA (Ours)
15.63%
67.34
54.69
81.02
75.69
44.40
77.53
70.72
Table 5: Comparison between X-EcoMLA and PALU baselines on Llama3- 8B. Despite operating at a signifi- cantly lower KV size (15.63%), X-EcoMLA matches or outperforms 50%-KV SOTA methods.
Expanded Discussion and Experimental Details The Appendix provides additional context and supporting results that complement the main paper. In Section I, we present more detailed discussion of related work on KV cache compression, craining-based memory-efficient architectures, and upeycling techniques. Section A.2 outlines the pseudocode for our proposed SVD-based ini- tialization strategy for MLA layers. Hyperparameter for both fixed and dynamic rank selection are reported in Section A.3, helping guide practical configurations. Section A.4 includes supplementary evaluations: long-context benchmarks on LongBench (Section A.4. I); comparisons with H2O as an alternative KV cache compression method (Section A.4.2); hybrid MLA variants that combine atten- tion and MLA layers (Section A.4.3); and detailed results from our extreme KV cache compression experiments (Section A.4.4). In Section A.5, we present ablation studies, analyzing the impact of distillation loss versus direct supervision (Section A.5.1). I). the role of LayerNorm (Section A.5.2), and the trade-off between larger teacher models and longer training (Section A.5.3). Collectively, these sections provide deeper insights into the design choices and empirical robustness of X-EcoMLA across diverse settings and compression budgets.
Conclusion
In this work, we introduced X-EcoMLA, lightweight post-training adaptation approach that enables the upcycling of pre-trained Transformer attention into an efficient MLA or hybrid variant. By leveraging dark knowledge from a well-trained teacher model and employing SVD-based initializa- tion, our method significantly reduces KV cache memory requirements without sacrificing model performance. Our results show that using an 8B teacher model allows us to compress the KV cache size of the Llama3.2. 3.2-a..2-111 IIstructt baseline by 6.4x while preserving 100% of its average score across multiple tasks on the LM Harness Evaluation benchmark. This is achieved with only 3.6B training tokens and about 70 GPU hours on AMD MI300 GPUs. Alternatively, we can compress KV cache by 10.6x using 7B training tokens over approximately 140 GPU hours, while maintaining 99.8% of the average score. These findings highlight the potential of X-EcoMLA as a practical and scalable solution for integrating MLA into existing LLMs, paving the way for more memory-efficiee and deployable models without the need for costly retraining from scratch.
10

---

Published as a conference paper at COLM 2025
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 checkpoints. arXiv preprint arXiv:2305.13245 2023.
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recahl-thooughput tradeoff. arXiv preprini 24Xiv:2402.18666 2024.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.055 05000, 2020.
Aviv Bick, Kevin Y Li, Eric P Xing, Zico Kolter, and Albert Gu. Transformers to ssms: Distilling quadratic knowledge subquadratic models. arXiv preprint arXiv:240..1018, 2024.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 33:1877-1901, 2020.
Chi-Chih Chang, Wei-Cheng Lin, Chien- Yu Lin, Chong- Yan Chen, Yu-Fang Hu, Pei-Shuo Wang, Ning-Chi Huang, Luis Ceze, Mohamed Abdelfattah, and Kai-Chiang Wu. Palu: Kv-cache compression with low-rank projection. In The Thirteenth International Conference on Learning Representations, 2025.
Jiuhai Chen, Rifaa Qadri, Yuxin Wen, Neel Jain, John Kirchenbauer, Tianyi Zhou, and Tom Gold- stein. Genqa: Generating millions of instructions from handful of prompts. arXiv preprint arXiv:2406.10323, 2024.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Lev- skaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstent Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways, 2022.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try are, the ai2 reasoning challenge. arXiv preprint I 1803.05557 2018.
Ganqu Cui, Lifan Yuan, Ning Ding, Guanming Yao. Wei Zhu, Yuan Ni. Guotong Xie, Zhiyuan Liu, and Maosong Sun. Ulrafeedbbck:: Boosting language models with high-quality feedback, 2023.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.2/060 2024.
Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabaleshwarkar, Shih-Yan Yang Liu, Matthijs Van Keirsbilck, Min- Min-Hung Chen, Yoshi Suhara, Yingyan Lin, Jan Kautz, and Pavlo Molchanov. Hymba: A hybrid-hhee head architecture for small language models, 2024. URL https://arrn org/ /211 13676.
11

---

Published as a conference paper at COLM 2025
Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, Jason Phang, Laria Reynolds, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation. 2023.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.22 27..007727773.. 2024.
Albert Gu and Tri Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseck-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.1294: 2025.
Ethan He, Abhinav Khattar, Ryan Prenger, Vijay Korthikanti, Zijie Yan, Tong Liu, Shiqing Fan, Ashwath Aithal, Mohammad Shoeybi, and Bryan Catanzaro. Upcycling large language models into mixture of experts. arXiv preprint arXiv:2410.07524 2024.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Mcasuring massive multitask language understanding. arXiv preprint arXiv:2009.080 03300, 2020.
Tao Ji, Bin Guo, Yuanbin Wu, Qipeng Guo, Lixing Shen, Zhan Chen, Xipeng Qiu, Qi Zhang, and Tao Gui. Towards economical inference: Enabling deepseek's 's multi-head latent attention in any transformer-bet llms. arXiv preprint arXiv:2502. 14837, 2025.
Qiao Jin, Bhuwan Dhingra, Zhengping Liu, William W Cohen, and Xinghua Lu. Pubmedga: A dataset for biomedical research question answering. arXiv preprint 19090099990116 2019.
Hao Kang, Qingru Zhang, Souvik Kundu, Geonhwa Jeong, Zaoxing Liu, Tushar Krishna, and Tuo Zhao. Gear: An efficient kv cache compression recipefor near-lossless generative inference of I1m. arXiv e-prints, pp. arXiv-2403, 2024.
Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upcycling: Training mixture-of- experts from dense checkpoints. arXiv preprint 1rX.v:2212.00005, 2022.
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. Race: Large-scale reading comprehension dataset from examinations. arXiv preprint arXiv: 04683, 2017.
Wing Lian, Bleys Goodson, Eugene Pentland, Austin Cook, Chanvichet Vong, and "Teknium", Openorca: An open dataset gpt augmented flan reasoning traces. https:/ hets:://https///// huggingface ct/ Oww.://eeaaaaaaeeeeeeee www.orr/r//pppo..a..aa 2023.
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, and Yoav Shoham. Jamba: A hybrid transformer-mambe language model, 2024.
Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chengqi Dengr, Chong Ruan, Damai Dai, Daya Guo, et al. Deepseek-v2: strong, economical, and efficient mixture-of- experts language model. arXiv preprint MrXXiv:405.04443. 2024a.
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.194377 2024b.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv: 1809.00000000000 027892 2018.
12

---

Published as a conference paper at COLM 2025
Beijing Academy of Artificial Intelligence (BAAI). Infinity instruct. https://huggingffnt 2024.
co/dat /tt////ttttt///AA/aaiiittttt-yn/
OpenAl. Gpt-4 technical report, 2023.
Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Bjöm Deiseroth, Kristian Kersting, Tajji Suzuki, Brian Hie, Stefano Ermon, Christopher Ré, al. Mechanistic design and scaling of hybrid architectures. arXiv preprint arXiv:2403.17844, 2024.
Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Baohong Lv, Xiao Luo, Yu Qiao, and Yiran Zhong. Transnormerllm: A faster and better large language model with improved transnormer, 2024. URL https: aarxvvar.vvvv//aa///////eaa//aav org/ /bbs/207. 14995.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106. 2021.
Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150, 2019.
Luohe Shi, Hongyi Zhang, Yao Yao, Zuchao Li, and Hai Zhao. Keep the cost down: review on methods to optimize Ilm's ky-cache consumption. arXiv preprint arXiv:2407. 18003, 2024.
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. Advances in Neural Information Processing Systems, 37:7339-7361, 2024.
Teknium. Openhermes 2.5: An open dataset of synthetic data for generalist IIm assistants, 2023.
URL htlpt:/t///// hetttss:/hhuggingffttee co http://aaccccctttttttteeeeeeeeee
Junxiong Wang- Llama3 ultrafeeddacc-armmonn dataset. https://huggingtte -uwtraeeddbaaac---eeo armboaacck-rmrrm.... 2024.
co/
datasets Jtn//uunxiiinno Swuxxiongwaagg/llnne... wwwp.//ggaaangooommameeaeeeeeeee
Junxiong Wang, Daniele Paliotta, Avner May, Alexander M. Rush, and Tri Dao. The mamba in the
llama: Distilling and accelerating hybrid models, 2025. URL https: arxiv. orwoorgggooooo /bss/
2403.
Zheng Wang, Boxiao Jin, Zhongzhi Yu, and Minjia Zhang. Model tells you where merge: Adaptive kv cache merging for lims on long-context tasks. arXiv preprint arXiv:2407.08455 2024.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309. /7453, 2023.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with efffcciieee training. 2024.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine 2019.
really finish your sentence? arXiv preprint
Hailin Zhang. Xiaodong Ji. Yilin Chen, Fangeheng Fu, Xupeng Miao, Xiaonan Nie, Weipeng Chen, and Bin Cui. Pqcache: Product quantization-bart kvcache for long context Il inference. arXiv preprint arXiv:2407.12822 2024a.
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Re. The hedgehog & the
porcupine: Expressive linear attentions with softmax mimicry. In The Twelfth International Conference on Learning Representations, 2024b. URL https: http::///ppeeeeee.ee/m 1d=4g0212N2Nx.
Dt.pa//U...tttps/pppeeeeeeee.... net.eeww.eeeeeeeeeeeeeeeeeeeeeee forum?
Zhenyu Zhang. Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36:34661-34710, 2023.
page
13

---

Published as a conference paper at COLM 2025
A Appendix
A.1 Related Work
KV Cache Management in Transformers Several approaches have been proposed to reduce or compress the KV cache size of Transformers, which can be broadly categorized into training- -based and post-training solutions Shi et al. (2024). Training-based methods involve modifying the model architecture and pre-training it, typically yielding better performance, whereas post-training methods are easier to apply and do not require retraining.
A variety of post-training KV cache management solutions have been explored in the literature. One common strategy is KV cache eviction, such as Heavy Hitter (H2O) Zhang et al. (2023), which defines an eviction policy based on the observation that only few tokens contribute to the highest attention scores. This method retains the most recent and most significant tokens while discarding the others. Another approach is sliding window attention Arora et al. (2024); Beltagy et al. (2020), which restricts attention to a fixed number of recent tokens (or predefined patterns) to maintain bounded KV cache size. Attention Sink Xiao et al. (2023) builds on this by retaining initial tokens in the KV cache to improve performance. Quantization-basee KV cache compression Kang et al. (2024); Zhang et al. (2024a) reduces memory usage by storing KVs in a lower-precision format, while KV cache merging Wang al. (2024) minimizes information loss by merging KV entries instead of discarding them. Although post-training solutions are computationally efficient, they often lead to performance degradation due information loss. In contrast, training-based methods offer better balance between memory efficiency and model accuracy. This paper focuses on training-based solutions, which we review in the following.
Training-based KV Cache Management Training-based solutions modify the attention mecha- nism or replace it with alternative architectures in Transformer models to reduce KV cache memory requirements. For instance, multi-query attention (MQA)Shazee (2019) shares keys and values across all attention heads, reducing the KV cache size by a factor of n compared to multi-head attention (MHA) model with KV heads. However, sharing a single KV across n query heads can be too restrictive. To address this, grouped-query attention heads into groups, allowing each group to share a single set of key and value heads, making balance between memory efficiency and performance. Another notable approach is YOCO Sun et al. (2024), a decode-decoder model that consists of a self-decoder and a cross-decoder module. Instead of storing KV vectors for each layer and token, the self-decoder module provides shared global KV cache to the cross-decoder layers, significantly reducing memory overhead. Multi-head latent attention (MLA), introduced in DeepSeek- V2 Liu et al. (2024a). is another KV-cache efficient variation of MHA. MLA reduces KV cache size by projecting input hidden states into a compressed latent space through low-rank projection, leading demonstrated that MLA can outperform standard MHA while maintaining efficiency.
et al. (2023) divides query
substantial reduction in memory usage. DeepSeek-V
Motivated by MLA's strong performance and efficiency, we focus on adapting MLA for already pre-trained models. However, training-based solutions typically require full pre-training from scratch or extensive continual training. This raises a fundamental question: Can we upcycle pre-trained models to their MLA counterparts without costly retraining? ? In the following section, we review existing solutions for model upcyeling that can be leveraged for this purpose.
Upeyeling Attention In Komatsuzaki et al. (2022), model upcycling is defined as "upgrading an existing model with relatively small additional computational budget." This term has primarily been used to describe the conversion of dense models into mixture-of-experrs (MoE) models in an efficient manner Komatsuzaki et al. (2022); He et al. (2024). In this paper, we focus on the concept of attention upcycling, which involves adapting pre-trained attention blocks Transformer into more efficient forms, such as MLA, without requiring full re-training from scratch. There are several examples of attention upcycling in the literature. For instance, in GQA, Ainslic et al. (2023) propose replacing MHA blocks with GQA and performing light continual pre-training for adaptation. Similarly, HedgehogZhang et al. (2024b) introduces an upcycling method that converts pre-trained attention into linear attention using knowledge distillation.
notable line of work focuses on leveraging the duality between Transformer self-attention and alter- native architectures. MambalnI Jama Wang al. (2025) demonstrates this by replacing some attention
14

---

Published as a conference paper at COLM 2025
Algorithm Python-like pseudocode of the proposed SVD initialization for MLA.
MHA weights: R_Q, W_K, W_V
MLA weights: W_DQ, W_UQ, WLOR, W_DKV, WLKK, WLKR, W_UV
Initialization of W_DQ, w_uQ, and W_QR
J_q, sigma_q, P_q svd (W_Q)
DQ u_q
w_ JQR bar = (sigma q 8 .view d n h.)
W JQ = W. FQU barl:, ", :d qk] vvee n h_ gk)
W_QR bar ;, -d_r:]..iew((e
r_h+d_r)
10
11 Initialization of WLDRV, W_UK, W_KR, UL
12 J_k,, sigma_kv, V_kv = svd storch.cat (W_K, W.V,, -1))
13 W DKV = U_kv
K aavg- K.w..vieww (d, n_h, d_h) (1)
15 W_RR =__Kavgl;, -dr:]
16
17 WUKV = sigma_kv 8 V_kv
_JJ_bar - - W_UKV[;, :d_ham_h)_vieeirrkk n_h, d_h)
w W_ JK -K_u__br[[, :d__gkk view (r_kv, g_)
20 W_ m_UKv[; d_
layers in pre-trained models with Mamba layers, initializing them from their corresponding attention layers, and then fine-tuning using end-to-end knowledge distillation. Similarly, MOHAWK Bick et al. (2024) follows knowledge distillation- based approach for training hybrid attention-Mamba Mamba models. However, MOHAWK differs from MambalnL.lama in some aspects: (a) It does not initialize the student sub-quadratic model from the Transformer attention layers; (b) It incorporates intermediate layer distillation in addition to end-to-end distillation.
A.2 Algorithm
Our simple method for initializing the MLA weights using SVD approach applied to the pre-trained attention weights is summarized in the pseudocode in Algorithm 1.
A.3
Hyper-parameter Selection
In Table and 7, we present the model performance with different hyperparameters for fixed rank selection and dynamic rank selection, respectively. In Table 6, we evaluate three KV ranks (Nkv = 512, 256, 128) and two head dimensions (dgt = 32, 64). We adjust r accordingly to make sure all configurations have approximately the same number of parameters. The results indicate significant performance loss the KV rank The decreases. With the same KV rank, d, 64 generally provides better performance. However, such advantage is more obvious with fk= 128, 256 where rg is relatively large. When Nke = 512, both head dimensions provides similar performance. In Table 7, we explore two thresholds (90% and 95%) for Ty and The and two head dimensions (d,x = 32, 64) for dynamic rank selection. When training with small portion of the dataset (1.6B), we notice that the performance is mainly influenced by the KV rank Fie- Although setting dg = 64 leads to more parameters, it does not necessarily translate to performance improvement, even when trained with the full dataset.
A.4 Supplementary Results
A.4.I Long Context Evaluations
We evaluated our MLA-optimized models on the LongBench benchmark, which covers range of long-context understanding tasks such as LCC, Qasper, QMSum, Multi-News, and SamSum. Table reports the results under various KV-cache size for both Llama3.2-15 2-B and Llama3.2-3B models. Notably, X- X-EcoMLA 3B models achieves a score of 60.03 on LCC, significantly outperforming the full-sized Jama3.2-3B baseline (52.1) 1), despite using only 43% of the KV cache. Across other tasks
15

---

Published as a conference paper at COLM 2025
Configuration
Param
Fg
kk
dgk
KV Size
Tokens
Avg Score
Base
Model:
Llama3.2 2-11
1-1BBBB B-Inst
EX-EcoMLA +DPO
1.23B
864
512
32
53.1%
3.6B
53.04
TX-EcoMLA +DPO
1.23B
480
512
64
56.3%
3.6B
53.14
X-EcoMLA +DPO
1.23B
1184
256
32
28.1%
3,6B
51.91
X EcoMLA +DPO
1.23B
736
256
64
31.3%
3.6B
52.38
X EcoMLA +DPO
1.23B
1344
128
32
15.6%
3.6B
51.38
TX-FcoMLA +DPO
1.23B
864
128
64
18.8%
3.6B
51.60
Table 6: Hyperparameter selection for the internal dimensions of the X X-EcoMLA block under fixed setting with 100% MLA layers, without LayerNorm, and using an identical teacher model as the basc.
Configuration
Param
Tqy
Tko
dyk
KV Size
Tokens
Avg Score
Base
Model:
Llama3.2-11
TX-EcoMLA +DPO
1.22B
90%
90%
32
42.7%
1.6B
51.26
X X--EcoMLA +DPO
1.25B
90%
90%
64
45.9%
1.6B
51.31
1 X-EEcMLA +DPO
1.23B
95%
90%
32
42.7%
1.6B
51.36
1 X-FEcMMAA +DPO
1.27B
95%
90%
64
45.9%
1.6B
51.21
1 X-ECcMLA +DPO
1.23B
90%
95%
32
54.7%
1.6B
52.18
X EcoMLA +DPO
1.26B
90%
95%
64
57.9%
1.6B
51.51
X-EcoMLA +DPO
1.24B
95%
95%
32
54.7%
1.6B
52.40
X- EcoMLA +DPO
1,28B
95%
95%
64
57.9%
1.6B
52.16
X- EcoMLA +DPO
1.23B
90%
95%
32
54.7%
7.0B
53.22
1 X-EcoMLA +DPO
1.26B
90%
95%
64
57.9%
7.0B
53.23
Table 7: Hyperparameter selection for the internal dimensions of X-BcoMLA block under a dynamic setting with 100% MLA layers, without LayerNorm, and using an identical teacher model (Llama-34 1.ama-3,2-11 3,2-1B) as the base.
such as Qasper, Multi-News, and SamSum, our compressed models match or even slightly exceed the performance of their full-cache counterparts. These results indicate that our method scales well to Long-sequence scenarios and is particularly effective in memory contory-nntttainedd environments.
Model
KV-Size
Avg. Acc.
Icc
repobench-pp
qasper
qıasum
multi news
samsun
Llama3 2-11--nn-tt (Base)
100.00%
30.805
1X-EcoMLA (ours)
53.13%
30.77
38.73
40.36
21.13
20.50
25.76
38.11
1X-FooMT (ours)
23%
30.66
38.74
40.54
21.21
20.61
25.62
37.26
Llama3. 2-3B-Inst (Base)
100.00%
40.01
52.11
$4.16
40.42
23.63
26.51
43.21
1X-EcoMLA A (ours)
42.91%
39.29
60.03
56.24
29.94
27.54
40.93
EX-EcoMLA (ours)
25.00%
39.11
59.59
53.94
31.75
20.93
27.19
41.26
Table 8: cont-context evaluation on the LongBench benchmark across varying KV cache sizes. All the X- X-EcoMLA models are trained with Llama3.1-88 1B-IInt as the teacher model
A.4.2 Comparison with other KV Cache Compression Techniques
We compare X-EcoMLA with the widely used H2O method Zhang et (2023), using the same base model (Llama3.22 2--1----IInnunt Instrrutct) and identical KV cache sizes. The evaluation is conducted on the Im-eval-hardness benchmark to assess performance under increasingly aggressive memory constraints. As shown in Table 9, X-EcoMLA consistently outperforms H2O across all compression levels-both in terms of average accuracy and on most individual tasks. Notably, at KV size of 9.4%, X-EcoMLA achieves average accuracy of 50.49%. %., compared to 45.05 % for H2O, with particularly large gains on ARC, ARE, and PIQA. Even at 6.25% KV size, X-FcoMLA maintains strong lead, indicating its robustness under extreme compression. These results demonstrate that X-EcoMLA achieves significantly better accuracy under the same memory budget, making a strong candidate for memory-efficient inference.
16

---

Published as a conference paper at COLM 2025
Model
KV-size
Avg. Acc.
ARC
ARE
HS
MM
OBQA
PIQ
PM
RA
WG
H2O
15.6%
50.30
37.71
57.41
59.91
40.83
31.60
71.11
60.40
37.99
55.80
TX-CcoMLA (ours)
15.6%
51.97
40.10
62.88
58.17
39.70
37.80
73.50
56.60
39.33
59.67
H2O
9.4%
45.05
30.03
43.01
$7.79
33.25
29.60
64.96
58.80
36.08
$1.93
1X-EcoMLA (ours)
9.4%
50.49
39.16
62.63
56.04
34,90
36.40
72.85
40
37.70
58.33
H20
6.25% %
41.30
26.54
34.68
52.75
26.95
28.60
59.03
58. 60
34.26
50.28
1X-EcoMLA (ours)
6.25%
49.74
38.48
61.66
55.32
30.62
35.20
72.36
56.60
37.99
59.43
Table 9: Comparison of X-EcoMLA and H2O Zhang al. (2023) across various KV cache sizes. X-EcoMLA consistently outperforms H2O, especially under aggressive compression.
A.4.3 Hybrid MIA Models
In this section, we include some supplementary results. Table 10 shows the benchmark performance of our X-EcoMLA method on Llama3.2. 2- IB-Inst model when we use the same model as teacher. We evaluate three different initialization settings: (i) Fixed rank selection with random initialization, (ii) Fixed rank selection with SVD initialization, and (iii) Dynamic rank selection with SVD initialization. For the fixed rank selection scenario, we set r = 854, Iku = 512, and dgk d, 32 such that the total number of parameters after the MLA upcyeling remain roughly the same. For the dynamic rank selection case, we apply a threshold of 0.95 for both rq and Tko so that the number of parameters aligns with other setups. We investigate two MLA layer upeycling strategies: upeycling 100% of layers to MLA and upcycling 50% of layers to MLA. For the 100% upcycling strategy, we replace all GQA modules in the base model with MLA. In this scenario, the proposed X-EcoMLA model uses only 53.1% % of the KV cache size for fixed rank selection and 54.7% for dynamic rank selection. For the 50% upcyeling strategy, we replace GQA modules layers 1, 3, 5, 7, 8, 10, 12, and 14. This brings us 78.1% KV cache size for the fixed rank selection and 78% for dynamic rank selection.
For each strategy, we evaluate training with the full dataset (6.8B tokens) and half dataset (3.4B tokens). It is evident that for fixed rank selection schemes, SVD initialization significantly enhances distillation performance compared to random initialization, yielding an 8% improvement for 100% MLA and 3% improvement for 50% MLA.
Model and Setting
Init, Method
KV-Size
Tokens
ARC
ARE
HS
MMLU
OBQA
PIQA
PBMD
RA
WG
Avg.
Llama3.2 1-1-ass
Base
100%
37,97
63.30
60.65
46.05
60.00
38.18
59.67
52.77
100% MLA Layers- Teacher: Meniical to the Base Model
1X- 1X-HooMLA
Random (512)
53.1%
6.88
35.32
60.48
54.03
27.77
35.20
71.98
55 80
33.88
55.01
47.72
1X-EcoMLA
Random (512)
53.1%
7.0B
38.99
62.71
56.20
28.00
36.8
36.27
56.20
1X EceMLLLA
Fixed (512)
53.1%
6.8B
36.95
63.89
58.88
43,40
36,00
74.16
58.20
37.32
60.30
52.12
"X-BeoMLA DPO
Fixed (512)
53.1%
7.08
40.19
63.93
60.67
42,31
37.60
75.03
59 20
40.86
61.01
53.42
1X-EcoMLA
Dynamic (95%)
54.7%
6.8B
37.12
63.64
58.87
43.26
73.72
60.00
37.51
60.22
52.08
1X- EcoML++ DPO
Dynamic (95%)
54.7%
7.0B
40.36
64.31
60.88
42.54
36,80
74.16
61.40
40.77
60.69
53,54
1X EXoBBoMLA
(512)
53.1%
3.1B
37,37
64.35
58.36
42.03
35.00
73.61
57.40
37.03
59.51
51.63
TX- 4coM + DPO
Fixed (512)
53.1%
3.68
39.93
63.51
60.52
41.58
37,20
73.99
59
40.48
60.38
53.04
1X-EcAMLA
Dynamic (95%)
54.7%
3.4B
37.12
63.66
64
$8.44
42.14
34.40
73.61
57.00 57.00
37.22
59.98
$1.50
BooMAA + DPO
Dynamic (95%)
54.7%
3.6B
40,27
62.71
60.55
41.21
36.40
74.16
59.50
39.90
60.14
$2.79
50% MLA Layers, Teacher: Identical to the Buse Model
1X-EcoMLA
Random (512)
78.1%
6.8B
62.79
57.23
38.
56.20
36.75
58.72
50.72
1X-FooMIA 4 DPO
Random (512)
78.1%
7.0B
38.99
63.66
59.00
37.44
37.69
74.59
57 (K)
39.14
661.4
$1.99
1X EX-EcoMLA
Fixed (512)
78.1%
6.8B
37.97
63.01
59.71
44.37
35.50
74.86
60.60
38.37
59.98
52.74
1X-EcoMLA
Fixed (512)
78.1%
7.0B
40.87
63.93
61.95
43.39
37.20
74.15
59.
40.48
60.85
53.66
1X-BcoMLA
Dynamic (95%)
78%
6.8
38.48
63.85
59.78
44.27
35.60
74.54
60 40)
38. 18
60.85
52.88
(X-EcoMLA A DPO
Dynamic (95%)
78%
7.0B
41.64
64.44
61.78
43.58
36.40
74.21
60.00
40.48
60.93
$3.71
1X.BLA
(512)
78.1%
3.4B
37.12
63.55
59.36
44.16
35.20
73.94
57.65
37.70
60.77
52.16
1X-E DPO
(512)
78.1%
3.6B
40.61
64.73
62.06
43.51
37.40
73.78
59.49 59.40
10.29
61.25
53,67
1X-EooMLA
Dynaiic (95%)
78%
3.48
38.05
63.09
59.37
43.60
35.00
74.27
59 80
36.94
60.93
52.33
EX-BAoMLA
Dynamic (95%)
78%
3.6B
39.93
61.76
43.33
36.40
40.57
60.30
53.46
Table 10: shot evaluation of MLA variants with different initialization methods (random, SVD with fixed rank selection. and SVD with dynamic rank selection) on the LM Harness Eval benchmark across nine (ARC), ARC-Easy (ARE), HellaSwag (HS). MMLU, OpenBookQA (OBQA), PIQA, PubMcdQA (PBMD), RACE (RA), and WinoGrande (WG). (f denotes upcycling the base model.)
tasks:
17

---

Published as a conference paper at COLM 2025
4.4.4 More Details Extreme KV Cache Compression Experiments
In this section, we include more details for Table 2 results. For each row, we also show the results of SFT training.
Model and Setring
Teacher
Puram
Tokens
ARC
ARE
HS
MMLLE
OBQA
PIQA
PRMD>
RA
WGF
Avg.
2.
1,248
37,97
63.30
60.65
46.05
34.80
74.32
60:001
38.18
59.67
52.77
100% MLA
Layers (rky
512, 7
864, 6,a
32) KV Size: 53.1% %
1X-EcoMLA
Llama3.2 Inst
1.23B
3.4B
37,37
64.35
58.36
42.03
35,00
73.61
57.40
37.03
59.51
51.63
1X-BcoMLA
Llamm3.2-118 -Ist
1.238
3.6B
73.99
53.01
1N-FooMILA
1. Iost
1,238
3.48
37.71
65. 19
58, 84
43.13
36.20
73:45
59. 201
37.89
99.67
52.25
1X X--EcoMLA a DPO
Llama3.2. Inst
1.23B
3.6B
42.75
64.81
62.04
43.88
37.40
73.72
59.20
41.44
61.48
54.08
1X-EcoMLA
Llama3.2 Inst
1.23B
3.4B
67.38
43.18
73.91
60.10
$3.69
1X- 2X-BooMLA+ DPO
Jamm3.22-88-18
1,238
3,68
44,03
68.86
63,49
43.81
37,40
73.94
61. 40
41 82
61. A4
55.13
100% MI AA
Layers (fu
256,
1184, 4,a.
32) KV Size: 28.1 %
TX-BcoMLA
Llama3.2-18 Just
1,23B
3:4B
37.54
62.84
56.89
41.22
33.6
73. 12
55.4
36.46
59.19
50,70
1X-RooMLA
1,238
3.68
63.26
58.74
22.80
51.91
1X E.ooMMLAA
Llama3.2.33 Inst
1.23B
3.4B
36.35
63.51
57.09
41.30
35.00
73.07
56.30
36.46
60.14
$1.08
1X K-EcoMLA+ DPO
Llama3.2-38 Inst
1.23B
3.6B
40.70
64.35
60.10
41.77
37.20
73.83
57.80
39.23
61.17
52.91
1X-BooMLA
Llama3.2-888
1,238
3.48
41.15
36
73.65
52.12
1X-EoMMLA I DPO
Llama 2. Inst
1,238
3.6B
41,98
66.46
61,33
41.78
37,20
74.27
59:00
40.00
60.69
53.63
100% MLA
Layers (7a
128,
1344, 4,a.
32)
KV Size: 15.6%
1X EX-EcoMLA
Llama3.2 Inst
1.23B
3.4B
36.52
61.41
55,37
38.02
34,60
72.52
56.00
35.60
58.56
49.84
1X-BcoMLA
Llamm3.2-1BB-mm
1,238
3.6B
56.40
51.38
1X-FooMILA
1,238
3.4B
36,26
61. .5.
55. 84
39.28
35.40
71.76
57. 60
35.89
59.27
50.36
(X XX-EcoMLA DPO
Llama3.2.30 Inst
1.23B
3.6B
39.42
62.88
58.41
39.45
37,20
73.39
58.00
39.71
59.75
52.02
1X-EcoMLA
Llama3.2-888 Anst
1.23B
3.4B
36.35
60.10
$1.20
AX- DPO
Laama3.2.2 Inst
1,238
3,68
41,30
65.40 61
59.64
39.47
37. 60
74.25
59. 201
39.5 52
59.85
52.94
1X XX-BoMLA
Llam33.2 2-8-.-11---1 -int
1,238
6,8B
37,54
62.21
56. 36
39.67
35.40
73.
55. 60
35.3 31
55.33
50.41
1X-FooMIA
Llama Insst
1.238
7B
40.10
62.88
39.70
37,80
73.50
39.33
59.67
$1,97
XX EcoMLA
Llama3.2- Inst
1,23B
6.8B
35,58
63.51
56.71
41.38
35.80
72.80
57.20
35.59
58.56
50.83
"X-BooMLA DPO
Llama3.2 2-38-------8 Inst
1,238
7B
39.33
64.56
$8.92
41.86
37.40
73.83
58. 90
39.71
59.59
52.70
AX-FooMIAA
Jaama 2. Inst
1.238
6.88
38.65
66.88
58,46
42.01
73.
60:000
59.12
52,23
X. EcoMAA DPO
Llama3.2-88 Inst
1,23B
7B
42,49
67.13
60.58
42.51
36.60
73.99
59.40 $9.40
40.38
59.43
53.61
100% MLA Layers (Tip
64.
1424, dge
52)
KV Size: 9.4%
1X-EooMLA
Lamm332-118-mnc
1,238
6.8B
37.12
61.32
$4.46
34.89
35.60
72. 36
36.80
35.22
57.30
49.45
1X-EcoMLA
Llama3.2-11B Inst
1.23B
7B
39.16
56.04
72.85
56.40
37.70
58.33
50.49
TX-EcuMLA
Llama3.2-38 Inst
1,23B
6.8B
35.07
61.95
$1.95
38.61
35,20
72.09
57.40
35.98
58.25
19.91
1X-HoMMIA DPO
Llamn3.2 3.-33------
1,238
78
37.97
63.55
56.95
37.54
35.40
72. 74
57.001
38.60 66
59.27
51.01
IX-EcoMLA
Llama3.22
Inst
1.23B
6.8B
36.09
57.01
72.96
58.00
98
59.98
$1.05
XX EcoMLM++ DPO
Llama3.2- Inst
1,23B
7B
40.02
67.17
58.40
38.53
37.80
73.83
58,00
39.43
60.93
52.68
100% MLA
Layers (h = 48, r, 14:00, dia
= 32) KV Sizc: 6.25%
"X-BcoMLA
Llama3. 1-8--ns
1,238
6,88
36.77
60:61
53.51
32.44
33.40
72.
56. 56.60
34.55
58.33
48.71
1X-EcoMIA
Llama Insstt
1.23B
7B
38.48
61.66
55.52
30,62
72.36
56.60
37.99
59.43
49.74
X-EeoMLA
Llama3.2-33 Inst
1,23B
6.8B
33,70
61.32
54.11
35.96
34.60
71.27
56.00
35.22
53.48
48.96
"X- +PO DPO
Llamm3. 2-38--3--888 Inst
1,238
7B
36.18
62.21
55.82
36.41
35.60
72.03
57.00
38.09
60.06
50.38
1X-FooMIA
Ilama Inst
1,238
6.8B
36.35
64.00
55.50
36.65
$4.60
72.31
35.77
58.25
50.21
XX EcoMAA + DPO
Llama3.2-8B. Inst
1,23B
7B
37.71
65.32
57.32
36.27
36.80
72.96
58.00 20
38.76
58.80
51.35
Table 11: Impact of KV-cache compression and teacher model size on performance. Reducing the KV-ccche size lowers accuracy. but larger teacher models help recover performance. DPO further improves alignment and accuracy. (f denotes upcycling the base model.)
A.5 Ablation Studies
A.5.1 Distillation vs. Cross-Entropy
In Table 12, we examine the trade-off between learning from the teacher knowledge (via KL. diver- gence) and direct supervision from the dataset (via cross-entropy loss w.r.t. ground truth) during the SFT distillation stage. We adopt Llama3.2 2-11--Insttuue for both our base model and student model, and we use CE and KL to denote the weights for the cross-entropy loss and KL divergence loss. The results show that relying solely direct supervision (CE KL 0) significantly degrades model performance (48.54 vs. 52.77), underscoring the importance of leveraging teacher knowledge for effective learning.
18

---

Published as a conference paper at COLM 2025
In contrast, incorporating teacher knowledge either exclusively in combination with direct super- vision- yiie-syyelll the best results, indicating the importance of sascher-guide learning in maintaining accuracy. Given these insights, primarily adopt teacher-bre based learning in configurations to minimize hyperparameter tuning efforts unless stated otherwise.
Configuration
CE
KL
Avg Score
Llama3.2 I-IstInst
52.77
1 X-EEcMLA
0
50.84
X-EcoMLAA
0.01
50.93
X T-EcoMLA
1
0.05
50.71
X EcoMLA
1
0.1
50.98
X--EcoLAA
1
0
48.54
Table 12: Comparison of different CE and KL loss weightings in the SFT knowledge distillation phase. The experiment utilizes the same teacher the base model and applies dynamic SVD compression with 8q = 6v 0.95, trained 20% of the dataset.
A.5.2 Impact of LayerNorm
The original MLA module incorporates additional LayerNorm layers between the down- and up- projection operations. However, we observe that it is beneficial to omit those LayerNorm layers in our proposed X-RcoMLA, as evidenced Table 13 and the loss curves Figure the Appendix. By removing the intermediate LayerNorm layers, our proposed X-FcoMI.A demonstrates an improved loss convergence. Besides, across various setups (fixed, dynamic) and different training dataset sizes (3.4B, 6.8B) in Table 13, the removal of LayerNorm layers consistently leads to performance gains.
Loss
Loss with SVD Init, LN
10a
Loss with random init
Loss
10a
0
2500
5000
7500
10000
12500
15000
17500
Steps
Figure 3: Loss curve comparison between random initialization and SVD initialization w/ and w/o LayerNorm layers. All schemes trained with fixed KV rank selection and 100% MIA layers upcyeling.
Figure 3 shows the loss curves of our SVD initialization with and without layer normalization, as well as with random initialization. The results demonstrate that removing layer normalization leads to lower loss values.
A.5.3 Larger Teacher or more Training Data?
Table 14 highlights the impact of increasing training data (tokens) versus using a larger teacher model on both accuracy score and training time. When using the same teacher model (Llama3.22 2--1BIIn IB-ntt)), increasing the number of training tokens (from 3.4B to 6.8B) improves performance but comes at the cost of significantly higher training time (from 4.82 to 9.64 hours).
On the other hand, switching to a larger teacher (c.g.. .lamm3.2-3B. 2-3B-Inst or 1amma3.2-8B-Ima provides notable accuracy improvements with less reliance additional training data. For instance, using the
19

---

Published as a conference paper at COLM 2025
Configuration
Init. Method
LayerNorm
KV Size
Tokens
Avg Score
Base Model:
Llama3.2 1-------1IIs Inst
XX-EcoMLA
Dynamic (Ōgv 0.95)
/
54.7%
6.8B
51.68
1 X E-EcoMLA
Dynamic (6kv 0.95)
X
54.7%
6.8B
52.08
X-EcoMLA
Fixed (Ntv 512)
/
53.1%
6.8B
51.68
1 X-EcoMLA
Fixed (Nhv = 512)
X
53.1%
6.8B
52.12
X- EcoMLA
Dynamic (Ōkr 0.95)
V
54.7%
3,4B
51.18
1 X-EcoMLA
Dynamic (ōkv = 0.95)
X
54.7%
3,4B
51.50
X-EcoMLA
Fixed (Fkv 512)
53.1%
3.4B
50.89
1 X-FcoMLA
Fixed (Fkv = 512)
X
53.1%
3.4B
51.63
Table 13: Comparison MLA with LayerNorm vs. without LayerNorm
Model and Setting
Teacher
Tokens
Training time
Avg Score
Base Model: Llama3.2-1B-Inst
100% MLA Layers (Kkv 512, rg
864)
{X-BcoMLA
Llama3.2-1B-BB
6.8B
9.64 hours
52.12
1 X-EcoMLA + DPO
-lama322-1--hhl
7.0B
10.06 hours
53.42
1 X-BcoMLA
Alama3.2-1BB-Ins.
3.4B
4.82 hours
51.63
1 X-EEOMLA + DPO
Llama3.2-1B-Imm
3.6B
5.24 hours
53.04
XX-EcoMLA
3.4B
6.24 hours
52.25
X-EcoMLA + DPO
3.6B
6,65 hours
54.08
1 X-BcoMLA
2omma3.2-8B-Inst
3.4B
8.54 hours
53.69
X-FEoMLAA + DPO
(lamm3.2-8B--J
3.6B
8.96 hours
55.13
Table 14: Comparison training efficiency and accuracy when increasing training data (tokens) versus using a larger teacher model. Larger teachers yield better accuracy with moderate time increases. The time cost measured on 8 MI300 GPUs.
8B teacher with DPO achieves the highest score, outperforming training with double tokens under smaller IB teacher (55.13 vs. 53.42), even with less training time (8.96 vs. 10.06 hours). However, this comes with a moderate increase in training time.
These results suggest that leveraging stronger teacher model is generally more efficient for improving accuracy than simply increasing training data. While additional tokens help, the benefit diminishes compared to the gains from using larger teacher, making training with larger teacher more effective strategy when computational resources allow.
20


Reviews:
[{"review_id": "7JJda6NKZ6", "reviewer": "Reviewer_YRgt", "summary": "The paper proposes to “up-cycle” LLMs pre-trained with vanilla MHA attention into MLA attention, with the goal of improving memory efficiency at test time without performance degradation.\n\nThey start with instruct-tuned checkpoints and replace a subset (or all) of the MHA blocks with MLA, initializing the MLA weights via SVD on the pre-trained attention weights.  They propose variants of the SVD initialization, as well as random initialization for the MLA up-cycling.  Full and hybrid MLA models are considered. The model is then trained with KL-Distillation from a (usually) larger model using an SFT corpus, and finally aligned with DPO.  They report KV cache compression with minimal loss in performance on LM benchmarks, training with a modest compute budget.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "Though up-cycling vanilla Transformers into more efficient variants isn't novel (MHA -> GQA, softmax attention -> linear in T2R and others), the proposed solution for MLA up-cycling is well-described, and the proposed initialization of the MLA weights by SVD on the pre-trained attention weights is elegant.  The design decisions are well-ablated (compared to random initialization) and show a clear boost, within a modest compute budget.", "reasons_to_reject": "The biggest drawback of this study is that the results are on small models (135M, 360M, 1B, 3B) where the benefits of KV compression are minimal; even for relatively long sequences the memory requirements are fairly modest for these models.  There may be some benefits for embedded devices, but this issue justifies discussion in a dedicated limitations section. \n\nFor larger models where MLA would start to make a bigger difference (7B+), it's quite possible that up-cycling the models in this way would lead to a performance hit, and this remains untested. Such an investigation would require more computation resources which may not be feasible, but the limitation should be noted.\n\nAlthough the authors mention the advantages of MLA in the long-context setting, the evaluations presented are standard LM harness tasks which do not evaluate long-sequence performance; proper long-sequence benchmarks (e.g. SCROLLS [1]) would reveal potential degradation/lack of degradation in this setting.  This may be a limitation of the up-cycling procedure however, as the SFT data used might not be of sufficient context length to keep long-context performance. \n\n[1] Shaham, Uri, et al. \"Scrolls: Standardized comparison over long language sequences.\" EMNLP 2022", "questions_to_authors": "The authors apply X-EcoMLA exclusively to instruct-tuned models with SFT data, leaving open the question of the effectiveness of the approach on base models when continually pre-training on minimal pre-training data (in addition to the potential limitation for long-context performance raised above).  \n\nHave you tried the technique with minimal continued training on pre-training data? Up-cycling works typically used pre-training data, e.g. [2], but [3] analyzed up-cycling with both pre-training and SFT data, and found advantages with each approach. (Some citations of recent up-cycling works are missing, including the original work that introduced up-cycling transformers into RNNs (T2R [1]) and recent extensions to LLMs [2][3].)\n\nCan you comment on the preservation of long-context performance, given the advantage of MLA on long sequences?\n\n[1] Kasai, Jungo, et al. \"Finetuning pretrained transformers into rnns.\" EMNLP 2021\n\n[2] Mercat, Jean, et al. \"Linearizing large language models.\" COLM 2024\n\n[3] Zhang, Michael, et al. \"LoLCATs: On Low-Rank Linearizing of Large Language Models.\" ICLR 2025", "extra_scores": {}}, {"review_id": "OpoqLHxfkx", "reviewer": "Reviewer_VttQ", "summary": "This paper introduces a post-training method X-EcoMLA to transform pre-trained Transformer attention into Multi-head Latent Attention (MLA), enabling extreme KV cache compression without retraining from scratch. By leveraging SVD-based initialization and knowledge distillation from a larger teacher model, X-EcoMLA reduces memory requirements (up to 10.6×) while preserving or even improving model performance on the LM Harness benchmark. The method supports both static and dynamic rank selection and omits LayerNorms for better convergence.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "- The paper introduces a lightweight post-training adaptation mechanism. Unlike prior MLA work that requires expensive pre-training, X-EcoMLA enables efficient KV cache compression on existing pre-trained models with minimal compute (as low as 70 GPU hours), making it highly practical.\n- Across multiple model sizes (e.g., SmolLM, Llama3.2) and benchmarks (LM Harness with 9 diverse tasks), X-EcoMLA achieves up to 10.6× KV compression with negligible performance loss (<0.1%). \n\n- The method democratizes MLA adoption by requiring only post-training on modest datasets (3.4–7B tokens) and affordable compute. This greatly lowers the barrier for academia and industry groups with limited GPU resources to adopt memory-efficient LLMs.", "reasons_to_reject": "- The approach heavily relies on large teacher models (e.g., LLaMA3.1-8B) to recover performance under extreme compression. This raises a concern about practical viability in low-resource settings, where such teachers may be inaccessible. The performance degradation without large teachers (e.g., at 6.25% KV size) is notable and suggests limited standalone utility.\n\n- While the paper demonstrates results on SmolLM and Llama3.2 models, all target models use GQA attention and relatively small scales (≤3B). It remains unclear whether X-EcoMLA generalizes to diverse architectures (e.g., standard MHA, MQA, or non-GQA models) and larger-scale LLMs, which limits the current scope of its applicability.", "questions_to_authors": "- The method relies heavily on large teacher models for extreme compression. Could you quantify the minimum teacher quality/size necessary to recover performance at various KV compression ratios?", "extra_scores": {}}, {"review_id": "O1FR8Qd0Ac", "reviewer": "Reviewer_5zCs", "summary": "This paper presents X-EcoMLA, a post-training distillation framework that enables Multi-head Latent Attention (MLA) in already pre-trained models for KV cache compression. Experiments show that X-EcoMLA can compress KV cache by up to 10.6× with less than 0.1% performance drop, using only a few billion tokens and ~100 GPU hours, thereby offering a more practical path to MLA adoption.", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "- The post-training alternative of MLA avoid the pre-training costs, which makes it more accessible and scalable.  \n- The empirical results look good -- it maintains accuracy at high compression ratios (e.g., ~6.25% of original KV cache).", "reasons_to_reject": "- There are already numerous post-training methods that leverage similar low-rank approximations for MLA or related modules. The paper does not clearly distinguish how its approach differs from or improves upon these existing techniques.  \n- The paper does not include comparisons with strong and relevant post-training KV compression methods.  \n- Despite the paper’s focus on efficient deployment, it lacks system-level metrics such as latency, throughput, or wall-clock inference time. These are essential to substantiate claims about the practical benefits of KV cache compression and to understand trade-offs beyond accuracy.", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces X-EcoMLA, a method for post-training conversion of pre-trained attention mechanisms to Multi-head Latent Attention (MLA) for KV cache compression. In this way, models can be retrofitted to use MLA without the need to train new models from scratch.

Authors engaged with reviewers, addressing their concerns: they demonstrated scalability to 8B models (Llama3/3.1-8B), showed effectiveness on long-context tasks (LongBench), validated applicability to base models and standard MHA architectures (SmolLM), and provided comprehensive comparisons with recent methods (MHA2MLA, PALU, H2O). Experiments show significant throughput gains (x1.7-2) and memory reduction (5× at batch 128). However, reviewer engagement was limited - only YRgt acknowledged the authors' responses.

This paper can have a significant practical impact, and his technically and experimentally strong.
Decision: Accept