TransMLA: Migrating GQA Models to MLA with Full DeepSeek Compatibility and Speedup
Authors: Fanxu Meng
Pingzhi Tang
Zengwei Yao
Xing Sun
Muhan Zhang
Venue: NeurIPS 2025 spotlight ()
Keywords: Deepseek
MLA
RoPE
PCA
LLaMA
Qwen
GQA
MHA
Primary area: deep_learning

Abstract: Modern large-language models often face communication bottlenecks on current hardware rather than computational limitations. 
*Multi-head latent attention (MLA)* addresses this by compressing the key-value cache using low-rank matrices, while the Absorb operation prevents the KV cache from reverting to its original size, significantly boosting both training and inference speed.
Despite the success of DeepSeek V2/V3/R1, most model providers have heavily invested in optimizing GQA-based models and, therefore, lack strong incentives to retrain MLA-based models from scratch.
This paper demonstrates that MLA provides superior expressive power compared to GQA with the same KV cache overhead, thereby offering a rationale for transitioning from GQA to MLA.
In addition, we introduce TransMLA, a framework that seamlessly converts any GQA-based pre-trained model (e.g., LLaMA, Qwen, Gemma, Mistral/Mixtral) into an MLA-based model. 
For the first time, our method enables *direct conversion of these models into a format compatible with DeepSeek's codebase*, allowing them to fully leverage the existing, highly-optimized support for the DeepSeek architecture within inference engines like vLLM and SGlang.
By compressing 93\% of the KV cache in LLaMA-2-7B, we achieve a **10x speedup** with an 8K context length while maintaining meaningful output. 
Moreover, the model requires only **6B tokens** for fine-tuning to recover comparable performance across multiple benchmarks.
TransMLA provides a practical path for migrating GQA-based models to the MLA structure, and when combined with DeepSeek’s advanced optimizations—such as FP8 quantization and Multi-Token Prediction—further inference acceleration can be achieved.

TransMLA: Migrating GQA Models to MLA
with Full DeepSeek Compatibility and Speedup
Fanxu Meng', Pingzhi Tang':
Zengwei Zengeel Yao', Xing Sun Muhan Zhang 1.2
Institute for Artificial Intelligence, Peking University
State Key Laboratory of General Artificial Intelligence, BIGAI Tencent Youtu Lab, Shanghai, China 'Xiaomi Corp., Beijing, China
https//gitthbb.com/uuluaaaPP
Abstract
Modern large-language models often face communication bottlenecks on current hardware rather than computational limitations. Multi-head latent attention (MLA) addresses this by compressing the key-value cache using low-rank matrices, while the Absorb operation prevents the KV cache from reverting its original size, significantly boosting both training and inference speed. Despite the success of DeepSeek V2/V3/R1, most model providers have heavily invested optimizing GQA-based models and, therefore, lack strong incentives to retrain MLA-based models from scratch. This paper demonstrates that MLA provides superior ex- pressive power compared to GQA with the same KV cache overhead, thereby offering a rationale for transitioning from GQA to MI.A. In addition, we introduce TransMLA, framework that seamlessly converts any GQA-based pre-trained model (c.g.. LLaMA, Qwen, Gemma, Mistral/Mixtral) into an MLA-based model. For the first time, method enables direct conversion of these models into a format compatible with DeepSeek's 's codebase, allowing them to fully leverage the existing, highly-optimized support for the DeepSeek architecture within inference engines like vLLM and SGlang. By compressing 93% of the KV cache in LLaMA- 2-7B, achieve 10x speedup with 8K context length while maintaining meaningful output. Moreover, the model requires only 6B tokens for fine-tuning to recover comparable performance across multiple benchmarks. TransMLA provides a practical path for migrating GQA-based models to the MIA structure, and when combined with DeepSeek's advanced optimizations as FP8 quantization and Multi-Token ceriitionnfuuthae inference acceleration can be achieved.
Introduction
Advanced Large language models (uhh as GPT-4o (1), Claude 3.7 Sonnet (2]. Gemini-
2.5 (3). LLaMA-4 [], Mistral-3 (5), Qwen-3 (6]. DeepSeek V3/R1 7 8). Gemma-3 [9]. and
Phi-4 10 represent rapidly evolving frontier for both research and applications. At their core, LLMs rely on next-token prediction [I] 12]: tokens are generated sequentially, process that requires self- self-attention be computed over all preceding tokens at each generation step. To avoid the computationally expensive recalculation of these states, implementations store the intermediate key-value (KV) pairs in cache. Yet, model and context sizes grow, the KV cache itself becomes major bottleneck for inference.
"Equal contribution. Corresponding author: muhan@pku. edu.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Cacked
RoRoPE BrogFold
DoPE [8]
Kuye
IRY-PPA
GQA
MiQA
C
(a) Given the same KV cache size, the expressiveness increases in the order of GQA, MLA, and MQA.
(b) TransMLA concentrates positional information into Kmpe and compresses Kagee and V.
Figure I: (a) GQA, MLA, and MQA can be equivalently transformed in one direction, illustrating a gradual increase in expressive power. (b) RoRoPE aggregates positional information in the first head, eliminating the need for RoPE in others. FreqFold further enhances this effect. After balancing the magnitudes of Kepp and V. joint low-rank approximation is applied to compress the KV cache.
To mitigate these challenges, Group-Query Attention (GQA) B partitions the query heads into groups, where every head within a group shares single key and value head. GQA thus generalizes both Multi-Query Attention (MQA) [14] and Multi-Head Attention (MHA) [[5]: it degenerates to MQA when single group is used, and reduces to MHA when the number of groups equals the number of query heads. While both GQA and MQA cut the size of the KV cache relative to MHA, they do so the cost of model quality. Post-training K VV-cache compression techniques-such as Duo-Attention [16,, KiVi [17]. KV-Quant (18). H2O [99 and Palu [20 further shrink memory usage, but their non-standard implementations demand specialized optimizations, hindering widespread adoption.
with DeepSeek V2 [2] and further refined in DeepSeek- V3 77/R11 8 offers pre-trained KV-cache compression strategy that strikes excel- lent balance between computational efficiency and model quality. Models equipped with MLA deliver state-of-the-an results while driving training and inference costs new lows. Moreover, the DeepSeek team's 's ongoing commitment to open-source releases provides highly optimized implementations and deployment recipes, making these advances readily accessible the community.
Multi-Head Latent Attention
In this paper, we first prove that MLA consistently offers higher expressive power than GQA under the same KV cache overhead, which theoretically explains the advantage of MLA. However, practical hurdle preventing model vendors from switching to MLA is the substantial prior investment on GQA-based models. This motivates to ask:
seamlessly convert a GQA-baaed model, such as LLaMA 44 and Qwen 6 to MLA so that we can inherit the model weights and pretraining effort, rather than training MLA from scratch?
based pretrained
key obstacle to converting a GQs-based model to MLA is that the Absorb operation 20]. which DeepSeek uses to switch between the compute and memory-efficient modes, is blocked. This occurs because every query/ key head has its own Rotary Positional Embedding (RoPE) [22]. Borrowing from DeepSeek's 's Decoupled RoPE scheme, concentrate the positional signal in K into small subset dimensions, Knope- The remaining dimensions, Kuope, contain negligible positional content. This decoupling allows us to drop the RoPE from Knope and then absorb its up-projection into the query projection-mirroiinn mirroring the DeepSeek conversion.
enabling scamless MLA
To efficiently concentrate positional information into fewer dimensions, we introduce RoRoPE novel technique that performs principal component analysis (PCA) on the key output, applies rotation across the two ends of RoPE, and consolidates the principal components of all attention heads into the dimensions of the first attention head. We theoretically prove that the product remains invariant after rotating the query and key using an orthogonal matrix U, as long as U satisfies two conditions: (1) rotation occurs only within the same dimensions across all attention heads and (2) the real and imaginary components of RoPE are rotated in the same manner. Additionally, by exploiting the frequency similarity between adjacent RoPE dimensions, we propose FreqFold, technique that improves the concentration efficiency.
Finally, we find that the mn-norm of Knope is much larger than that of V. If we run PCA on the concatenated matrix [Kmpei V] without adjustment, the principal components are dominated by Knopes leading to a severe loss of information from the value subspace and sharp decrease in accuracy. We therefore introduce Balanced Key-Value (BKV) procedure: we first rescale K Kaoe and

---

V so that their norms match, and only then perform joint PCA. This simple normalization restores the balance between the two subspaces and delivers marked improvement in compression quality.
The above innovations collectively form our TransMIAA method. Using TransML.A, we compressed the KV cache of LLLMM-22 by 68.75%, with only a 1.65% performance drop across 6 benchmarks for training free. In contrast, a concurrent method, MHA2MLA [23], experienced 21.85% performance decline. At a compression rate of 93%, the model still maintained meaningful responses, and after training with 6B tokens, its performance was mostly restored. We evaluated both the original (GQA/MHA) models and their ConssLA--coverttee
counterparts on different hardware setups using vLLM. The TransMLA models achieved up to a 10x speedup over the originals, demonstrating the significant potential of our approach. Moreover, the TransMLA models are fully compatible with DeepSeek's s code, enjoying DeepSeek's ecosystem to accelerate inference and seamlessly integrate with various hardware and frameworks.
2 Related Work
Autoregressive decoding in large language models necessitates storing past activations- key-value (KV) pairs-in a cache to avoid recomputation. Because the size of this cache grows linearly with sequence length, its memory footprint quickly becomes the limiting factor for very long contexts. Consequently, shrinking the KV cache without compromising accuracy has become a pivotal research focus, motivating spectrum of architectural innovations and compression strategies.
Multi-Query Attention (MQA) 14] and Stoup-Querry Attention (GQA) [3] shrink the KV cache by letting every query in group share a single key and value head. Although both schemes save memory relative to Multi-Head Attention (MHA) [5], they usually give up some accuracy. Multi-Head Latent Attention (MLA)-intee introduced with DeepSeek V2 (21] and refined in later releases DeepSeek V3/R1 -offers more favorable trade-off, delivering near-state-of-the-arr quality while cutting training and inference costs. Grouped Latent Attention (GLA) [25] provides a parallel-friendly implementation of latent attention that further accelerates MLA inference. By contrast, Tensor Product Attention (TPA) [26) tackles the memory bottleneck by dynamically factorizing activations, slashing the runtime KV cache by an order of magnitude, but it necessitates training the model from scratch. TransMLA fills this gap: rather than proposing yet another attention variant, it converts an existing GQA model into an MLA model with only light fine-tuning, restoring accuracy while inheriting MLA's memory and speed advantages.
21 24
Another approach is to optimize the KV cache of existing pre- praited models. For example, dynamic token pruning is employed by LazyLLM (27), A2SF 28). and SnapKV [29]. These methods selectively prune less important tokens from the KV cache. Sharing KV representations across layers, as in YONO 30). MiniCache 31.. and MLKV [2),, reduces memory by reusing the same KV cache across multiple layers. This can drastically lower memory usage and speed up inference. Although effective, both families of methods usually require custom kernels or runtime tweaks, complicating deployment and limiting adoption. TransMLA. by contrast. plugs directly into the mature DeepSeck dosystem--covvetee checkpoints load out-of-the-bon delivering MLA-level speed-ups across every DeepSeek-supporte platform.
There are two works most related to TransMLA. One is Palu 20 which reduces KV cache size by applying low-rank decomposition on both the keys and values, enabling speedup through tailored optimizations. However, Palu does not specifically handle RoPE, which prevents it from using the Absorb operation during inference. Therefore, Palu needs project the compressed representations back to their original size. This projection incurs significant computational overhead during inference, limiting the overall acceleration. Another concurrent work, MHA2MLA (23). also claims to convert MHA to MLA and decouple RoPE from the main computational path. It is important to clarify that MLA is not simply GQA extension of MHA2MLA both TransMLA and MHA2MLA support MHA and GQA architectures. However, MHA2MLA determines which RoPE dimensions to remove solely based on the norms of the query and key vectors, which tends to cause larger information loss when pruning the same proportion of positions. Also, the distribution of important dimensions in MHA2MLA is uneven, requiring sparse indexing that complicates optimization and acceleration. Their work reports compression ratios of the KV cache but does not demonstrate actual inference speedup. Furthermore, MHA2MLA directly applies joint singular value decomposition to KV, resulting in higher loss compared to our balanced key-value PCA method.
3

---

3 Preliminary
Rotary Position Embedding and Group Query Attention Details are provided in Appendix [A]
Multi-Head Latent Attention MLA operates an input sequence, where x, IRD is the embedding of the 1-th token and D is the model's hidden dimension. To save KV cache, first computes low-rank latent features <V using down-projection matrix W E g Then, up-projection matrices WUK E Radxrav and Wey R used to derive the key (k) and value (v) representations for all heads, where d is the per-head dimension. Similarly, MLA decomposes the query projection (using WDQ R''' and W Ridxro where rg is the query rank) to reduce activation memory. For positional embedding, MLA uses a decoupled RoPE strategy. This employs additional multi-head queries E. E Rd4 and a shared key k, E Rdd to carry the RoPE signal, which are generated from W E idld"xr, and WKR. E Rd"xD (d is the per-head RoPE dimension). The computation of the key and query representations is formulated follows:
KV = W DDKV Xt:
ci DD2 Xt,
[K, 2 [ = k. = WUK ot
KV.
[q 9:
9. = q. = =wuec ci
kf" = (W
[ q:2i
qthb) = qi ROPE, (WOR c.).
kt,- = [KC.;: kF"
(1)
qr = [q g..l.
(2)
MLA features a dual-mode capability, tailored for different operational stages. For the compute- intensive training phase, it adopts configuration resembling standard MHA (Equation (3), offering slightly lower computational overhead than conventional MHA. Conversely, for communication intensive inference, it can seamlessly switch to an MQA-like setting (Equation 4 In this latter mode, the latent features function as a shared large KV head, interacting with all query and output heads to efficiently produce the final output. This mechanism, known the Absorb operation, is crucial for accelerating inference.
wi,ii v vi WUY KV
qe, [WURT C L q k =c
K.t Vd+ dul S
softmax,
6,,1 softmaxy(
+ di 1c;
KV
y1 (o,,1; 01,2; o.A))
= W
[W UV o.. W o..1),
(4)
where w (UKKVV) denotes slices of the projection matrices corresponding to the i-th attention head. One of the main contributions of this paper the seamless support for the Absorb operation. significantly enhancing inference speed.
4 TransMLA
In this section we formally present TransMLA, motivated by two observations:
1. For fixed KV-cache budget, MLA is strictly more expressive than GQA. As proven in Appendix B and illustrated in Figure Ia any GQA layer can be rewritten as an MLA layer by introducing a single additional projection matrix. The reverse transformation is not always possible, implying that MLA subsumes GQA. This equivalence holds even when RoPE are present, although it requires the equivalent MLA layer to be expressed in the absorbed form.
2. Inference acceleration occurs when MLA uses a smaller KV cache. Although one can build an MlA-equivalent representation of a GQA model, speedups arise only if the number of stored KV vectors is actually reduced. TLa achieves this by converting the GQA-based network into DeepSeek-like MIA architecture while compressing its KV cache. This allows the transformed model to run directly on DeepSeek's 's inference stack and realize the full memory-latency benefits.
4.1 Merging Grouped Heads to a Latent Head
The core insight behind establishing the strictly greater expressive power of MIAA over GQA, and the first step in the TransMLA transformation, lies in unifying GQA's g grouped key-value (KV)

---

heads into single MLA latent representation. This transformation replicates GQA's repeat_kv mapping. where each query i maps to a specific KV group j Ti/(h/g)] To achieve this, we initialize the MIA up-projection W W as sparse block matrices: for each i-th query head, only the block corresponding to its designated j-th KV group set to an identity matrix, while all others are zero. This ensures that single key representation can be shared across ! query heads while maintaining interactions solely with its designated partner. W is initialized analogously to map to the j-th value head. Consequently, the g identical RoPE rotations must also be consolidated into single operator, denoted RoPE, which applies the original rotation pattern (of per-head dimension d) repeatedly every d dimensions to the unified key. In this way, the computation of GQA attention is transformed into the following form:
q 9:2: 9c,h = q W
Ic c11 = cVV WDDKV Xt:
(5)
c..
RoPE
WHIKT qr.i)
k = RoPE (c V. cY
(6)
q.i K. Vd
softmax,
Vj. y: = w o W o.,n],
(7)
j=1
where:
WDKV = > WV R2ddxD
if (j 1)d + K;
w Wee E Rdxgd [k, w ( 0 otherwise.
This GQA-equivalent formulation (Equation 66 is crucial because it makes the model structurally compatible with the absorb operation, enabling seamless switching between execution modes. However, this equivalence alone provides no acceleration. It is evident that the total KV cache size remains unchanged (c E ), which is identical to the original GQA model. Furthermore, as shown in Equation (6. the effective attention dimension increases from to gd (since q and k." are both in ) leading to higher computational costs. To achieve actual acceleration, compressing the KV cache is therefore essential. This unified representation is advantageous for compression, as merging multiple KV heads allows for better identification of shared principal components. Moreover, this merged-key structure is a prerequisite for efficiently decoupling RoPE in the subsequent steps.
Overview of the RoRoPE decoupling pipeline. Blue lines and orange lines denote the real and imaginary components, respectively. We first gather dimensions with the same rotational frequency (0)) across all heads. A joint Principal Component Analysis (PCA) is then applied to these dimensional real and imaginary components. This process concentrates the dominant positional signal into single principal component for each frequency, which is then represented by a standard RoPE (RoPE RoPE 2), while the remaining components are decoupled (NoPE 3-8).
4.2 Head-wise Rotation for Decoupled RoPE with Minimal Loss
To enable the essential KV cache compression identified in the previous section, we introduce a method, RoRoPE, As illustrated in Figure applying RoRoPE to the merged head removes the bulk of the positional signal from K.
In the GQA-equivalens form (Equation 7 the attention dot product q.. R.t is computed over the merged gd ad-diimesioonll vectors. Recall that RoPE operates by rotating d/2 independent 2D subspaces. Because our RoPE operator applies the same rotation pattern to all g heads, we can regroup this total dot product by summing over these d/2 independent subspaces. We define the real and imaginary components (Appendix A of the 1-th 2D subspace (where E [1, d22) as the (21 - 1) and (21)-th dimensions of any single head, respectively. Let q. (1.real) and L. (l.imag) be the g-dimensional vectors formed by gathering all real (dims [2| - 2, - 1 + d,... ..... and imaginary (dims [21, 21 + d.....) components, respectively, from the (absorbed) query heads. The key vectors k (?.eea) and k] (t.immg) are defined similarly from k. The total adddimensional dot product can then be expressed as the sum of d/2 independent 29 di-deennsional dot products:
d/2
9
R
R
RoPE,,
(lat..
(f.ree))) (1.imne)
a6 ( RoPE,,1 ([R(fren)) kg] ((imae)))
(8)
5

---

Group
Group 2
Group 3
Group 4
Dim
Dim
Dim
Dim
Dim
Dim
Dim
PCA
RoPE
NoPE
NoPE
NoPE
Figure 2: Overview of the RoRoPE decoupling pipeline. Blue lines and orange lines denote the real and imaginary components, respectively. We first gather dimensions with the same rotational frequency (8)) across all heads. A joint Principal Component Analysis (PCA) is then applied these g-dimensional real and imaginary components. This process concentrates the dominant positional signal into single principal component for each frequency, which is then represented by standard RoPE
where
(1.rcca)) RoPEg, ([x!4.caa)) x (1.mmg)) cos t0:
(1.mma)) sin t0 x x
(1.rcal)
|x (limg)))
(lreal)) x
(9)
When the I-th concatenated real and imaginary components of â.,, and k, are multiplied by an orthogonal matrix Uj Rox9, the inner product with RoPE applied remains invariant. Specifically,
d/2
RoPE, ((U,af. (l.ral), U.q(limme) ((llimae)) RoPE,;
I2 q k
(10)
This invariance holds because RoPE a linear rotation independently applied different heads, and any orthogonal transformation applied to the g-dimensional components preserves the inner product. However, this reveals a critical constraint: the same orthogonal matrix U, must be applied both the real and imaginary components within each 2D subspace. For a detailed proof and our proposed solution, please refer Appendix[[ [
In practice, our RoRoPE method uses Principal Component Analysis (PCA) to compute the required orthogonal matrices {U:)tE(L....d/2), from the key activations for each 1-th subspace. To apply these rotations required by Equation (10p. we transform the weight matrices W and WUK directly. This is because rotating the activations (qr.s; k, is equivalent to rotating their respective generation matrices. The rotation effectively concentrates the essential information into the first few components (heads), and in turn, allows us to perform the final decoupling: we remove the RoPE encoding from the non con-principal components, preserving positional information only within the principal components.
However, applying PCA directly to each 2g-dimensional subspace independently is suboptimal, as this space is often too small for effective compression. We address this by exploiting a key property of RoPE: adjacent rotational frequencies are often similar (#; F 0t+1). This similarity allows us to treat adjacent subspaces as effectively equivalent, grouping them to perform PCA over a larger, combined latent space. This technique, which we call FreqFold, achieves superior compression by identifying principal components across a broader set of dimensions. We defer detailed analysis of FreqFold to Appendix

---

4.3 Balancing Koppe and V for Improved Joint Low-Rank Compression
In the previous section, we split the key heads into one carrying positional information and the others without positional information, achieving minimal loss. We then apply Principal Component Analysis (PCA) jointly on the values and the non-positional components of the keys (i.e. Knope), using activations collected from a small calibration dataset, thereby compressing the projection matrices into low-rank latent space, However, we observed that although the principal components of the keys were effectively separated with RoRoPE, the norm of the residual key features remained significantly larger than that of the value features. This imbalance caused the direct decomposition to favor principal component directions dominated by the keys.
To mitigate this, we scale WDK by dividing it by nope
E2 W DK nope X [|2]
(11)
E,
W DV x,
and correspondingly scale WUK by multiplying it by a. Here, W E R(s-1)dxDD represents the last g 1 heads that does not use RoPE in WDK
This transformation is mathematically equivalent and does not affect the overall model outputs, while significantly enhancing the effectiveness of KV cache compression in subsequent steps. More details is provided in the Appendix E and
Experiment
5.1 Main Experiment
In this section, we present our main experimental results. Following the experimental setup of MHA2MLA, we converted two modeLs-smolLM 1.7B and Llama 2 7B7--nn into the MLA architec- ture. We evaluated the models' performance on six benchmarks three stages: before conversion, immediately after conversion without further training, and after conversion followed by training. For the training process, we used a subset of the pretraining corpus used for the smoll model. The fine-tuned results of MHA2MIA are taken directly from the original paper. Our experiments were conducted on an 8-GPU machine, each GPU having 40GB of memory and delivering 312 TFLOPS of FP16 compute power. Detailed hyperparameter settings are provided in the Appendix G
From Table T we observe that TransMLA efficiently facilitates architecture migration across various models and KV cache compression ratios. Notably, the untrained performance of MLA models initialized with TransMLA shows minimal degradation in capability compared the original mod- els sss-sinniiccanlly less than the degradation observed with MHA2MLA under equivalent KV cache compression. In fact, using TransMLA to compress the KV cache of Llama 2 7B to just 7.03% of its original size still results in better performance than MHA2MLA LA's compression to 31 31.25% on the same model. This highlights the effectiveness of our proposed techniques includes RoRoPE, FreqFold and activation- based balanced KV low-rank factorization.
The low-loss transformation achieved by TransMLA enables to recover the original model per- formance with minimal training overhead. As shown in the table, TransMLA achieves stronger performance than MHA2MLA-66 while using significantly fewer training tokens. For instance, when transforming smolLM 1.7B and compressing the KV cache to 31.25% of its original size, we only need 4.9% of the training data used by MHA2MLA and hours training to surpass its performance.
5.2 Key Norm Analysis Reveals the Impact of RoRoPE and FreqFold
In this section, we conduct a detailed analysis of the distribution of key activations in the attention module to demonstrate the effectiveness of our proposed methods.
Figure 3a presents the average L2 norm of each key dimension in the first attention layer of the
LLaMA 8B model, computed on subset of the WikiText-2 (39 dataset. We compare the original model (in blue), the model transformed using our RoRoPE equivalence method (in orange), and the further approximated model using 4D FreqFold (in green). The top and bottom halves of the plot correspond to pairs of dimensions that share the same rotation angle in RoPF, which we refer to as the real (Re-dim) and imaginary (Im-dim) dimensions.
7

---

Table 1: Commonsense reasoning ability of two LLMs with TransMLA compared MHA2MLA. The six benchmarks include MMLU (133. ARC easy and challenge (ARC, (4). PIQA (355). HellaSwag (HS. (361), OpenBookQA (OBQA, 37), and Winogrande (WG, [8). Tokens refers to the number of tokens used for further training after the TransMLA conversion. A value of indicates that the model was evaluated immediately after conversion, without any fine-tuning.
Model
Tokens
KV Mem.
Avg.
MMLU
ARC
PIQA
HS
OBQA
WG
SmoILM-I.7B
IT
55.90
39.27
59.87
75.73
62.93
42.80
54.85
-68.75%
40.97
27.73
41.96
63.00
29.19
34.40
49.53
-81. 21.5%%
37.14
26.57
32.73
55.77
26.90
31.40
49,49
-87.50%
34.01
25.32
27.15
51.36
25.47
26.20
48.54
-MIAAMMLA
-68.75%
54.76
38.11
57.13
76.12
61.35
42.00
53.83
-81. 31.25%
54.65
37.87
56.81
75.84
60.41
42.60
54.38
87.50%
53.61
37.17
55.50
74.86
58.55
41.20
54.38
-68.7%%
51.95
35.70
55.68
73.94
53.04
39.80
53.51
-81.25%
47.73
32.87
47.89
69.75
48.16
36.20
51.46
87.50%
44.12
29.97
41.72
66.87
41.15
34.80
50.28
TransMIA
300M
-68.75%
55.24
38.60
58.95
74.97
61.52
43.00
54.38
700M
-81.25%
54.78
37.79
57.53
75.52
59.88
42.80
55.17
IB
-87.50%
54.01
37.24
56.32
74.81
60.08
42.40
$3.20
LLaMA-2-7B
2T
59.85
41.43
59.24
78.40
73.29
41.80
64.96
-68.75%
37.90
25.74
32.87
59.41
28.68
28.60
52.09
-81255%
34.02
25.50
26.44
53.43
27.19
22.60
49.01
-87. 97.50%
32.70
25.41
25.79
50.60
26.52
19.40
48.46
-MIAAMMAA
-68.75%
59.51
41.36
59.51
77.37
71.72
44.20
62.90
-81. 2.5%
59.61
40.86
59.74
77.75
70.75
45.60
62.98
87.50%
58.96
40.39
59.29
77.75
69.70
43.40
63.22
75%
58,20
39,90
57.66
77.48
70.22
41.00
62.90
-87.500 57.50%
51.19
34.39
45.38
71,27
60.73
37.40
57.93
-92.97%
43.26
28.93
36.32
63.38
45.87
31.60
53.43
TransMLA
500M
-68.75%
59.82
40.87
59.18
77.91
71.82
45.20
63.93
3B
88.500
59.36
40.77
58.84
78.18
71.28
43.60
63.46
-92.97%
59.19
40.41
58.68
77.53
70.39
45.00
63.14
LLaMA-3-8B
1ST
63.84
46.20
65.75
80.47
76.20
45.60
68.82
-71.855%
54.13
36.38
52.84
73.83
64.34
37.00
60.38
TransMLA
30B
-71.875%
63.39
46.18
66.30
80.30
76.33
45.00
66.22
60B
-71.85%%
63.76
47.39
66.96
80.41
77.10
44.80
65.90
We observe that the original model exhibits a highly uneven norm distribution across key dimensions, with numerous outliers. This suggests that naively removing RoPE from certain heads would likely result significant performance degradation. After applying our RoRoPE transformation, as shown by the orange line, the key dimensions with large norms are nearly all concentrated to the first two heads (dimension 0-128). Further applying the 4D FreqFold approximation compresses the tail (i.e., higher-index dimensions) even more, leading to even sharper concentration of high-norm key dimensions. This concentrated structure is highly beneficial for the subsequent RoPE removal step as shown in Figure 3b
In Figure 3b we present the log-perplexity of the LLaMA 8B model on WikiText-2 as RoPE components are progressively removed. We observe that our proposed RoRoPE method significantly outperforms MHA2MLA's per-head dimension selection strategy, especially at high removal ratio. Furthermore, incorporating extreme removal rates. At 90% removal ration, RoRoPE 4D-FreqFold still maintains log- perplexity about 2, while MHA2MLA reaches nearly 6, which no longer generates meaningful outputs. At the same time, we observe that overly aggressive FreqFold (i.e., using too many dimensions) can degrade performance, as the loss introduced by approximation of nearby dimensions outweigh
approximation leads to even better performance under

---

2.00
Original ROROPE
8-
Ly of No
50-
Perplexity
Vanilla PCA
MHA2MLA
RoRoPE
RoRoPE & 2D-FreqFold
Ly-turn ifiddm
4-
RoRoPE & UD-Freqfold
16
RoRoPE & 8D-FreqFold
64
128
192
256
320
384
448
522
40 RoPE Removal Ratio (%)
80
Dimension
(a) Key norms of the first layer.
(b) RoPE removal results, evaluated on WikiText-2.
Figure 3: Visualization of key norms and RoPE removal results on LLaMA 8B model. The top and bottom halves of the left figure correspond to pairs of dimensions that share the same rotation angle RoPE, which we refer to as the real (Re-dim) and imaginary (Im-dim) dimensions.
Origine leoeor
20
Key
W-based
WX--besed
wmw
9-
W-bssed with
Wi--based BKY
10
Loo of Perpl Perplexity
5
128
384
640
Balaneed trnom
2.0
5
2.5
2.0
6.0
128
256
384
512
640
768
896
1024
75.0
125
87.5
Dimension
Key-Value Cache Compression Ratio (36)
(a) Norm magnitude of keys and values before and after KV balancing is applied.
(b) Comparison of different methods of KV cache compression, with and without KV balancing.
Figure 4: Visualization of the norms of keys and values for the first layer of LLaMA 3 8B and the perplexity results after joint low-rank compression of keys and values under WikiText-2. Here, W-based and WX-based refer to PCA applied on the attention weights and the activation outputs, respectively. BKV denotes the application of KV balancing.
the benefit in concentrating the principal components. This figure suggests that for LLaMA 8B, the sweet spot lies in applying RoRoPE combined with 4D FreqFold.
5.3 Key-Value Value Norm Disparity Motivates KV Balancing
visualize the norm magnitudes of the key and value activations in the first layer of LLaMA 8B before and after KV balancing. Note that both the key and value shown here are activations after applying the RoRoPE principal component concentration, and the first head of the sey-reserved for RoPE-is excluded. As result, the value and the remaining key components shown in the figure precisely the elements aim to compress jointly into dower-dimensional space via PCA.
In Figure 4a
It is evident that even after removing the first head with the highest norm, the overall norm of the key remains significantly larger than that of the value. This norm disparity poses a substantial challenge for joint compression into shared latent space. In particular, such imbalance bias the PCA toward directions aligned with the key, rather than the value, leading to suboptimal representation of the value components.
The lower part of Figure 4a shows the norm distribution of keys and values after applying KV balancing. At this point, the norms of keys and values become more aligned, which is beneficial for performing joint PCA. This observation is further supported by the results in Figure where KV balancing consistently reduces the loss incurred by jointly applying low-rank approximation keys and values- whether the PCA based on weights or on activations. Figure 4b also demonstrates that activation-based PCA yields significantly better results than weight-based PCA.

---

5.4 Haroware-Agnostic Inference Speedup with TransMLA
By converting MHA/GQA models into MLA models that are fully compatible with the DeepSeek codebase and compressing the KV cache, TransMLA enables us to leverage all optimizations and tooling available in DeepSeek. Using the vLLM framework, we achieve substantial real-world inference speedups.
In Figure 5 we benchmarked the inference performance of MLA model-with 92.97% reduction in KV cache size-on three consumer-grade AI accelerators with different compute capabilities and memory sizes: 165.2 TFLOPS with 24GB memory, 312 TFLOPS with 40GB memory, and 320 TFLOPS with 64GB memory. The figure shows the inference speedup of the MLA model relative to the original MHA model. Low-rank Q and Full-rank Q indicate whether the query projections were also compressed. Context length represents the total sequence length (i.e., context length plus generated tokens).
Our experiments show that the inference speedup of MLA models increases as the context length grows, which aligns with our expectations. Since the primary performance gain of MLA stems from KV cache compression, longer contexts lead to more substantial savings and thus higher speedups. Remarkably, for the 8K context window on the first hardware platform, the trrasfooodr model achieves an impressive 10x inference acceleration. To the best of our knowledge, the MHA2MLA method has not reported any inference speedup results.
644
10
1K
2k
4k
1k
24
4%
8K
16k
32k
1k
2K
4k
8%
18k
32k
1k
24
48
8K
15k
32K
(a) 165 TFlops 24GB
(b) 312 TFlops 40GB
(c) 320 TFlops 64GB
(d) 1979 TFlops 80GB
Figure 5: Inference speedups with TransMLA comparing to the original LLaMA2 7B model three consumer- grade Al accelerators. Context length represents the total sequence length.
Conclusion, Limitation and Future Work
In this work, we demonstrate that the expressive power of TransMLA is stronger than GQA under the same KV cache. To help existing GQA transition to the MLA structure with minimal cost, we propose the TransMLA method. By using the RoRoPE method, the multi-head KV positional information concentrated into the first head, and FreqFold further enhances this extraction effect. The positional information of the remaining query-key heads is removed, and the Balance KV norm method is used to jointly compress the values and the remaining heads of keys. TransMLA can convert models such as LLaMA and Qwen into MLA-base models, and the converted model incurs very little loss compared the original model, with performance being recoverable through training with only few tokens. Additionally, TLA can easily leverage the DeepSeek ecosystem for accelerated inference, achieving significant throughput improvements across various hardware platforms.
Although TransMLA significantly reduces the loss caused by decoupling RoPE through RoRoPE and FreqFold, and mitigates KV cache compression loss via KV balancing, error accumulation becomes increasingly severe for longer texts, making recovery more challenging. Moreover, potential issues may arise during training aimed at enhancing recovery performance, including subtle performance degradation and hallucinations, which must be carefully considered. To address these challenges, is worth exploring more advanced mathematical approaches that can better reconcile the norm disparity between keys and values, achieving improved performance under the same compression rate and ultimately enabling a truly training-firee conversion process.
All speed comparisons in this paper are conducted on the single GPU. When multi- Multi-GGU parallel inference is required, MLA is less compatible with tensor parallelism (TP) than GQA, since multiple devices must replicate the same latent cache, which typically necessitates the use of data parallelism (DP). How can the speed of TLansMLA be improved under TP? Furthermore, TransMLA should be integrated with pruning, quantization, token selection, and other optimization techniques to fully explore the upper limits of inference acceleration.
10

---

7 Acknowledgements
This work supported by the National Key R&D Program of China (2022D00660300), Center of Excellence, Peking University, and CCF-Tencent Rhino-Bird Open Research Fund.
References
[1] OpenAl. Hello GPT-4o, 2024. URL http://pss//opaaaaaccc/////////a
[2] Anthropic. Claude 3.7 sonnet, 2025.
URL http::/www..attiupiii..cmmeeeeet
[3] Google. Gemini 2.5: Our most intelligent ai model, 2025. URL URt_ttpp::/bbogg.oogll
[4] AI@ Meta. Llama 4 model card. 2025. URL http:://gttat...cm//essaaaaleea
11 ama htta-/mnddds//loommmaaaaaaaaa https//1aaaaa//////
MAA.e///DDDEEEDDDDD.A_ CARD_ md
[5] Mistral. Medium is the new large, 2025.
URL
mistral medium-2
[6] Qwen. Qwen3: Think deeper, act faster, 2025. URL btttps://qeenlm.gitiie qwen3
[7] 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.1945 19437, 2024.
[8] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Iims via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025.
[9] Gemma Team, Aishwarya Kamath, Johan Ferret, Shreya Pathak, Nino Vieillard, Ramona Merhej, Sarah Perrin, Tatiana Matejovicova, Alexandre Ramé, Morgane Rivière, et Gemma 3 technical report. arXiv preprint
2025.
[10] Marah Abdin, Jyoti Aneja, Harkirat Behl, Séhastien Bubeck, Ronen Eldan, Suriya Gunasekar, Michael Harrison, Russell J Hewett, Mojan Javaheripi, Piero Kauffmann, et al. Phi-4 technical report. arXiv preprint arXiv:2412.08900 2024.
[11] Alec Radford. Improving language understanding by generative pre-training. 2018
[12] Tom Brown, Benjamin Mann, Nick Ryder, Mclanic 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, 13:1877-1901, 2020.
[13] Joshua Ainslie, James TherT. 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.
[14] Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint 19 /50, 2019.
[15] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Fukasz. Kaiser, and Illia Polosukhin. Attention is all you need. NeurlPS, 2017.
[16] Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, and Song Han. Duoattention: Efficient long-context 1lm inference with retrieval and streaming heads. arXiv preprint arXiv:2410.10817 2024.
[17] Zhiyang Liu, Dong Zhang, Xinyi Li, and Ji Wu. Kivi: Quantized key-value representation for efficient long- context transformers. arXiv preprint arXiv:2402.0673) 2024.

---

[18] James Hooper, Li Dai, Zhen Zhang, and Seung- Hwan Lee. Kvquant: Quantization for efficient key-value caching in transformer models. arXiv preprint arXiv:2402.12345, 2024.
[19] 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-34777 10, 2023.
[20] Chi-Chih Chang, Wei-Cheng Lin, Yhien-Yu Lin, Chong- Yan Chen, Yu-Fang Hu, Pei-Shuo Wang, Ning-Chi Huang, Luis Ceze, Mohamed S Abdelfattah, and Kai-Chiang Wu. Palu: Compressing kv-cache with low-rank projection. arXiv preprint arXiv:2407.21118 2024.
[21] DeepSeek-AL. Deepseek-v2: A strong, economical, and efficient mixture ef language model. CoRR, (aas22405.04434, 2024. URL fttt:::://oo..og/1...... 04434
2405
[22] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputingg 168:3, 2024.
[23] 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-baset llms. arXiv preprint arXiv:2502. 14837, 2025.
[24] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang. Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025.
[25] Ted Zadouri, Hubert Strauss, and Tri Dao. Hardware-efficient attention for fast decoding. arXiv preprint arXiv:2505.21488) 2025.
[26] Yifan Zhang. Yifeng Liu, Huizhuo Yuan, Zhen Qin, Yang Yuan, Quanquan Gu, and Andrew Chi-Chih Yao. Tensor product attention is all you need. arXiv preprint arXiv:2501.06425, 2025.
[27] Qichen Fu, Minsik Cho, Thomas Merth, Sachin Mehta, Mohammad Rastegari, and Mahyar Najibi. Lazyllm: Dynamic token pruning for efficient long context Ilm inference. arXiv preprint arXiv:2407. 14057, 2024.
[28] Hyun-rae Jo and Dongkun Shin. A2sf: Accumulative attention scoring with forgetting factor for token pruning in transformer decoder. arXiv preprint arXiv:2407.20485 2024.
[29] Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Llm knows what you are looking for before generation. arXiv preprint arXiv:2404.1 14469, 2024.
[30] Tian Sun, Li Zhang, and Shuang Wu. You only need one: Efficient ky sharing across transformer layers. Proceedings of the 42nd International Conference Machine Learning (ICMI.), 2024.
[31] Akide Liu, Jing Liu, Zizheng Pan, Yefei He, Gholamreza Haffari, and Bohan Zhuang. Mini- cache: Kv cache compression in depth dimension for large language models. arXiv preprint arXiv:2405.14366. 2024.
[32] Zayd Muhammad Kawakibi Zuhri, Muhammad Farid Adilazuarda, Ayu Purwarianti, and Alham Fikri Aji. Mlkv: Multi-layer key-value heads for memory efficient transformer decoding. arXiv preprint arXiv:2406.0922 2024.
[33] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Mcasuring massive multitask language understanding. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL. htt.ttttt///peeeeeeeee http:://paaaeeeeeeeeeee/ee/aaaaa
[34] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the A12 reasoning challenge. CoRR, abs/ abs/0580. 2018. htLp:tt///aaxx org/ h/s///////// wwwp///..... 01883.055557
12

---

[35] Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. In The Thirty-Fourdt Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pages 7432-7439, AAAI Press, 2020, doi: 10. 2..169/AAALLV331005.0. aaai. v34105. 6239
URL http:: //oiii org....rgg11110 1609/
[36] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? In Anna Korhonen, David R. Traum, and Lluís Mårquez, editors. Proceedings of the 57th Conference of the Association for Computational Linguistics, ACL 2019, Florence, Italy, July 28- August 2, 2019, Volume I: Long Papers, pages 4791- 4800. Association for Computational Linguistics, 2019. doi: 10. 1/1993/V1//19-1472 URL https://doi.c sttps://doi.org/100 10. 18653/v1/pppp 1472.
[37] Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? A new dataset for open book question answering. In Ellen Riloff, David Chiang, Julia Hockenmaier, and Jun't ichi Tsujii, editors, Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium, October 31 November 4, 2018, pages 2381-2339 Association for Computational Linguisties, 2018. doi: 10.18653/V1/D188 1260. URL httttp://pi.orrrreeeeeet http:///...rr....o.oaa////////// v1855/v1//18-1000
[38] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: an adversarial winograd schema challenge at scale. Commun, ACM, 64(9):99-106, 2021. doi: 10.1145/347438. URL httttttttt/:..... orgooogg///o 10. 01145/37443
[39] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
[40] Loubna Ben Allal, Anton Lozhkov, Guilherme Penedo, Thomas Wolf, and Leandro von Werra. Smollmn-crppus. 2024. URL http:///wgggggaccccccoo/aaaaaaas smollm-corpo
[41] Anton Lozhkov, Loubna Ben Allal, Leandro von Werra, and Thomas Wolf, Fineweb-edur the finest collection of educational content, 2024. URL https::/huugiingtiiiiii cotpbceeet//ttttttts.see
HigggingFaceFF///ineeo.. edu
[42] Xun Wu, Shaohan Huang, and Furu Wei. Mixture of lora experts. arXiv preprint 2rXiv:2404.13626 2024.
[43] Anton Lozhkov, Raymond Li, Loubna Ben Allal, Federico Cassano, Joel Lamy-Poirier Noua- mane Tazi, Ao Tang, Dmytro Pykhtar, Jiawei Liu, Yuxiang Wei, Tianyang Liu, Max Tian, Denis Kocetkov, Arthur Zucker, Younes Belkada, Zijian Wang. Qian Liu, Dmitry Abulkhanov, Indraneil Paul, Zhuang Li, Wen-Ding Li, Megan Risdal, Jia Li, Jian Zhu, Terry Yue Zhuo, Evgenii Zheltonozhskii, Nii Osae Osae Dade, Wenhao Yu, Lucas Krauß, Naman Jain, Yixuan Su, Xuanli He, Manan Dey, Edoardo Abati, Yekun Chai, Niklas Muennighoff, Xiangru Tang, Muhtasham Oblokulov, Christopher Akiki, Mare Marone, Chenghao Mou, Mayank Mishra, Alex Gu, Binyuan Hui, Tri Dao, Armel Zebaze. Olivier Dehaene, Nicolas Patry, Canwen Xu, Julian McAuley, Han Hu, Torsten Scholak, Sebastien Paquet, Jennifer Robinson, Carolyn Jane Anderson, Nicolas Chapados, Mostofa Patwary, Nima Tajbakhsh, Yacine Jernite, Carlos Muñoz Ferrandis, Lingming Zhang, Sean Hughes, Thomas Wolf, Arjun Guha, Leandro von Werra, and Harm de Vries. Starcoder 2 and the stack v2: The next generation, 2024.
[44] Keiran Paster, Marco Dos Santos, Zhangir Azerbayev, and Jimmy Ba. Openwebmath: An open dataset of high-quality mathematical web text, 2023.
[45] Stack Overflow. Stack overflow, 2025. URL https://staccoverriiice com 2025-05-21.
Accessed:
[46] Zechun Liu, Barlas Oguz, Changsheng Zhao, Ernie Chang, Pierre Stock, Yashar Mehdad, Yangyang Shi, Raghuraman Krishnamoorthi, and Vikas Chandra. Llm-qat: Data-free quantiza- tion aware training for large language models. arXiv preprint arXiv:2305. 17888 2023.
13

---

[47] Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Gabriel Martín Blázquez, Guilherme Penedo, Lewis Tunstall, Andrés Marafioti, Hynek Kydliček, Agustín Piqueres Lajarin, Vaibhav Srivastav, et al. Smollm2: When smol goes big data-centric training of a small language model. arXiv preprint arXiv:2500.002337 2025.
[48]
Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Gabriel Martín Blázquez, Guilherme Penedo, Lewis Tunstall, Andrés Marafioti, Hynek Kydliček, Agustín Piqueres Lajarin, Vaibhav Srivastav, Joshua Lochner, Caleb Fahlgren, Xuan-Son Nguyen, Clémentine Fourrier, Ben Burtenshaw, Hugo Larcher, Haojun Zhao, Cyril Zakka, Mathieu Morlon, Colin Raffel, Leandro von Werra, and Thomas Wolf. Smollm2: When smol goes big data-centric training small language model, 2025. URL /wwpvv..vvvv///////////////a/av. org wwwpaabb..220000 02737
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All main claims made in the abstract and introduction (see Section I are well-aligned with the theoretical and experimental results presented in the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: See Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
15

---

Justification:
All theorical results are proved either in the main paper or in supplemental
material.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: All the information to reproduce the experimental results given in the paper, and the corresponding code will be released both in the supplemental material and github repository.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
16

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: All data used in the paper can be publicly accessed, and code will be released upon publication. 'he supplemental material contains detailed instructions on how to download and preprocess the data, as well as how to run the provided scripts to reproduce the main experimental results. Environment details, dependencies, and exact commands are documented to ensure faithful reproduction. For any baseline methods, scripts and data preparation steps are also included or referenced.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
bublic/guides/ httpr//guuiiee/oeSubbbiieeee
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips. cc/pbblic/pbcccc//cc/pcppp guwdes/dddeeee httpp////i.gggaeeeeeoo///aaaaaaa
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: In Section 5 the experiment settings are briefly discussed. Detailed experi- mental configurations are provided in the supplemental material.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Due to limited space in the main text for figures and tables, detailed reporting of error bars and statistical significance information is provided in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
17

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We have provided detailed information about the GPU parameters, including memory size and compute capability.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make it into the paper).
9. Code ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: We have carefully reviewed and ensured that all aspects of our research conform fully to the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: See supplemental materials.
18

---

Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considcrations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning conrectly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We have carefully reviewed and verified the licenses and terms of use for all existing assets used in this paper, including code, datasets, and pretrained models. The original creators are properly credited with citations to the corresponding papers.
Guidelines:
existing assets. The authors should cite the original paper that produced the code package dataset.
The answer NA means that the paper does
The authors should state which version of the asset used and, possible, include URL.
19

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the datasets.paperssii paperswi thcode. com/datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification:
We release new code assets alongside the paper. The documentation includes details on usage, training, license, and limitations, and is provided with the code repository. Assets are anonymized for submission.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
20

---

The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core methods of this research do not involve any large language models. LLMs were not used in any important, original, or non-standard capacity in this work.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to LLM policy http:s///earrppscccc/ooorrrerrre for what should or should not be described.
21

---

Contents
Introduction
2
Related Work
3
Prelininary
TransMLA
Merging Grouped Heads to Latent Head
4.2
Head-wise Rotation for Decoupled RoPE with Minimal Loss
4.3
Balancing Kooe and V for Improved Joint Low-Rank Compression
5
Experiment
5.1
Main Experiment
5.2
Key Norm Analysis Reveals the Impact of RoRoPE and Freqloold
5.3
Key-Value Norm Disparity Motivates KV Balancing
5.4
Hardware-Agmostic Inference Speedup with TransMLA
10
Conclusion, Limitation and Future Work
10
Acknowledgements
11
Additional Preliminaries
24
Rotary Position Embedding (RoPE)
24
A.2
Group Query Attention
24
Enhanced Expressive Power of MLA with Decoupled RoPE
25
B.
Introduction
25
B.2
Proof of Proposition]]
25
B.2.1 GQA < MLAFactoizeel
26
B.2.2
MLAAaagtrizee
MQA
26
B.2.3 Strict Inequalities: GQA < MLAFactorizedd
MQA
27
B.3
Expressiveness of MLA with Decoupled RoPE
27
Proof of RoPE Inner Product Invariance under Orthogonal Transformation
28
D FreqFold: Detailed Mechanism, Example, and PCA Efficiency
30
D.I Detailed Explanation of FreqFold and RoRoPE's PCA
30
D.I.1 Illustrative Example of FregFold
31
D.2 Formalizing the Benefit of FreqFold in PCA
31
D.3 Proof of Proposition
32
D.4 Discussion the Trade-off in Freq Fold
33
22

---

Balancing Vayue Norms and Low-Rank Approximation
33
KV Balancing: Purpose and Formulation
33
Low-Rank Approximation of NoPE Keys and Values using PCA
33
Theoretical Analysis of Balancing Scheme for PCA in Matrix Decomposition
34
Theoretical Analysis
34
FLI1
Problem Setup
35
F.1.2
Proof of Proportionality
35
R11.3
Verify Minimum Conditions
36
Experimental Settings of Fine-tuning
37
H
vLLM Benchmark Details
37
Inference Speed Ablation
38
Evaluation GQA Models
39
K Distillation Experiments
40
Ablation Study: Calibration Dataset
40
M Ablation Study: Core Components
40
Evaluation the LongBench Benchmark
41
Extension to Multi-Group MLA
42
Case Study
43
23

---

A Additional Preliminaries
A.1
Rotary Position Embedding (RoPE)
RoPE 22 is position encoding method that encodes the absolute positions with different rotations and incorporates the explicit relative position dependency the self-attention formulation. applies different rotations to tokens in different positions to encode the position information.
Consider x, RR to be the embedding of the t-th token with the hidden size d. The RoPE operation upon X produces a representation x, that encodes both semantic and positional information:
cos 10 cos 10
sin 10, sin 10,
x
cos t0
sin t02
x = RoPE,
=
X
4)
cos t0
sin t02
(12)
X 1)
Cos 102/2
(d)
sin 10a/2
cos t0 2/2
X
sin 2/2
(d)
1d-1)
where @ denotes the element-wise multiplication two vectors, x E IR denotes the i-th element of X,, and 0, 10000 - is the i-th rotation angle. If we interpret every two elements in the embedding as representation in the complex coordinate system, can divide x; into paired dimensions, where the odd-indexed dimensions x;
(2k-1) represent the real parts and the even-indexed
dimensions x (2k) represent the imaginary parts.
A.2 Group Query Attention
Let the 1-th token of the input sequence be x, E R,D, where D denotes the hidden dimension. To reduce the memory overhead of the KV cache, GQA divides the h query heads uniformly into groups, with all query heads within group sharing the same key and value vectors. Specifically, let WºQ E RidxD W WV e Ro D and W E DDxdd be the projection matrices for the query, key, value and output, where d D/h denotes the dimension per head. GQA first computes the concatenated queries qr, keys kt, and values Vt, and then slices them into heads or groups for attention computation:
9r.13 94.5 9ch) = W9
(13)
[k kz,) kt,g] k; W X2,
(14)
[Vr,1: Vi,21 Vt.aa Vt V Xx,
(15)
where each qr, E Rdd corresponds to the query vector of the i-th head, and ke,,, Vt.j E IRd correspond to the key and value vectors of the j-th group.
Using the notation in Section A. after applying RoPE qu kr, we can obtain the attention output for the 1-th token as follows:
T Rt 3.13/1
O,1
softmax,
1.
(16)
Vd
y: O 1 Or.2..... 0.,a)
(17)
As we can see, in GQA, each key and value head corresponds to h query heads. When = h, GQA becomes MHA, and when GQA becomes Multi-Query Attention (MQA).
24

---

B Enhanced Expressive Power of MLA with Decoupled RoPE
B.1
Introduction
This section provides a theoretical analysis to demonstrate that Multi-Head Latent Attention (MLA) with decoupled Rotary Position Embedding (RoPE), as deseribed in Section of the main paper, possesses greater expressive power than croupeddQuees Attention (GQA) (Section A.2 This analysis assumes comparable KV cache sizes and number of query heads.
Our primary argument focuses on the core projection mechanisms that generate queries, keys, and values, abstracting away from the specifics of RoPE application initially. We first present the following proposition concerning the relative expressiveness of these core mechanisms:
Proposition 1. Given the same KV cache size and number of query heads, the expressiveness of the attention projection mechanisms follows the order: GQA < MLAfactorizeed MQA.
Here, MLAFactorized refers to an attention mechanism employing low-rank factorization for its key and value projections, representing the content-processing aspect of the full MLA. It is im- portant to note that in the proposition, the query projection in MLAPactorizedd does not undergo low-rank factorization; this differs from the full MLA, where the query also factorized. After proving this proposition, an Pactorreee core for its content components and an MQA core for its decoupled RoPE com- ponents, is thereby more expressive than GQA. For this analysis, primarily consider the impact of the architectural structure on representational capacity, setting aside the direct effects of RoPE itself the expressiveness comparison between the fundamental GQA, MLA-Factorized, and MQA structures.
will discuss how the full MIA architecture, which incorporates such
O Cached During Inference
d
gd
gd
V
proj
g
K
proj
Q
GQA2gd
MLA
MQA
Figure 6: Comparison of Solti-Queryy Attention (MQA), Group Query Attention (GQA), and Multi- Head Latent Attention (MLA). In this work, illustrate that given the same KV cache size, the expressiveness increases in the order of GQA, MLA, and MQA. In the figure, h, d, g denote the number of heads, hidden dimension of each head, and the number of groups (K/V heads) in GQA, respectively. In MQA, the head dimension is set to gd to align the KV cache size with GQA and MLAA. As result, the KV cache size per token per layer for all three approaches is 2gd
B.2 Proof of Proposition
1
Let D be the hidden dimension of the input token X; E IR,, h be the number of query heads, and d = D/h be the dimension per head. In GQA, query heads are divided into groups. For fair KV cache comparison, the latent dimension for keys and values in MLA Mattoroorie (rtw) and the head dimension of MQA will be related to gd. Specifically, if the KV cache per token in GQA is 2yd for both keys and values, then in MI.A Factorizied The = 2yd, and in MQA, the head dimension is also 2gd; this ensures the KV cache sizes are aligned.
25

---

B.2.1 GQA S MLA Factorized
In GQA, query head qe attends to key k,, //((//0]] and value /x The GQA key projection w E Rodx produces 9 distinct key vectors :k.,i....;kk Similarly, WV E RgdxX produces value vectors. We define effective per-query-hear -ead projection matrices w E iidd and W Rhdx for GQA:
W
w
where WiK W
(18)
W
W
W
where W.Y W 14/0/9911
(19)
W
Here, W is the k-th x block of W Thus, KJ,; = W/* = ks/i//N//))> and similarly for values. The matrices W and VV have ranks at most gd.
An FaAoooideieed mechanism generates keys via k,,, = (WUK (W-DVV x,));. where WDKV grexXD and WUK E giddxre,, A similar formulation applies for values with WUV Rhddrr.
To demonstrate expressive capability, GQA S MLAPactorizedd we set Tke = 2gd. Let WDKV WK
R22dxD We seek WUK WUV Rdxxx2dd such that W WUKWDKY WV
WV
WUV WDKV, This is achieved by setting Wirk W Rddd2dd (the block for head i) as selector
matrices:
WUK
Oaxd: ....
Oaxds Laxds Odx
Oaxa]
(20)
blocks
2g
WUY
0axd ...
Odxd Laxd; Odxds Odxd
(21)
- blocks
g-k blocks
where [i/(h/g]]. Thus, GQA's key/value generation can replicated by an MLA Pactoriadd model with fke = 2gd and specific sparse structures for WUK and WUV The KV cache size 2gd (sequence length) preserved since we will be caching e WDKV, x, R22d On that account, the theoretical expressive power of GQA is less than or equal to that of MLAPactorized given the same KV cache size.
B.2.2 MLA Factorized & MQA
Consider an MLA-Factorized model where queries are qı,i W.ºx, x; (assuming W. E IR is the i-th block of W) and keys are k,,, (W. UK (W DKV x,)). The attention score for head involves kj,
= T( (WV DWV x;)).
(22)
This can be rewritten as:
J. = (WWUK)TTC xz) (WWDKV x,).
(23)
W
26

---

Let qr,i = W//x, X E 2 and vg = WDRKY, X, E R2add The computation of attention output
becomes:
110 Σ softmax
(24)
Vd
W [o :0,2i:: 0t,
Y:
WWV
of softmax,
JokU
WUV
Wo
(25)
Aliokv Vi )okV
WUV
softmax,
WP
This is an MQA formulation where each modified query fr, (now of dimension 2gd) attends to shared key and value c This indicates that the computations within MLA-Factorized can be structured use shared intermediate key and value representations akin to MQA's core. Thus, any MLA-Factorized model can be represented as an MQA model with a shared key/value of dimension 2gd.
B.2.3 Strict Inequalities: GQA < MLA Factorized < MQA
The relationships are strict:
GQA Facorriereeed When GQA is represented MLAA model, the up-projection trices WUK and WUV must adopt specific sparse, block-selector structures. general ML.Apattoriod model imposes such constraints; W and WUV are typically dense and fully learnable. This allows general MLA Pactorzze to create distinct key (and value) vectors by combining features from the -i-dimmesioonl latent space in complex ways. GQA restricted to g unique key (and value) vectors that are merely replicated h/g times. If h > g. MLLA can generate a richer set of interaction patterns. Thus, MLAPactorized has strictly greater expressive power.
FlAFoctrieeedd < MQA Consider the bilinear form x, Mx, in the attention score. In MLAPactorizzdd for head i, MALLA, (W:9) (WW determined by the smallest one among the ranks of w. e RdxD WUK gdxx2dd and WDKV R22dxD, which most d.
Factorizeds The maximum rank of the transformation is
However, in the MQA form derived from MLAfactorizedd, the rank the interaction matrix here, (W"R)TWe WDD is determined by the smallest one among the ranks of w E R2ddx and WDRVV D22ddD, which at most 2gd.
Since 2gd 2 d, MQA allows for potentially higher-rank interaction between the (modified) query and the shared key representations compared the per-head effective rank in MLAApactoizedd original formulation. This indicates that MQA has a greater representational capacity for the scoring mechanism.
B.3 Expressiveness of MIL with Decoupled RoPE
The full MLA architecture, as defined in Section 3 (main paper), employs a decoupled RoPE strategy. The query qı and key k, for head (in the MHA-like training paradigm, Equation 3 are:
qr. qq L
(26)
k,
k k.J
(27)
where k, is shared RoPE key component across all heads for token t. The bilinear attention score (numerator of the softmax argument) for head i between query at and key at is:
(q kG. + (q..) TT.t kf
(28)
Let's analyze the two components of this score:
27

---

1. Content Component Interaction:
( j. The content keys ky,, are derived from WUK (WDKV xz). This key generation mechanism for k9 is precisely that of the MLA ML.AFattriaa Factorized model discussed in Section B. As established, MLA-Factorized is strictly more expressive than GQA for the non-positional part of the representation.
2. Positional Component Interaction: (kf.) This interaction, where h distinct query- side RoPE components Ef, attend single, shared key-side RoPE component k. is an MQA structure specifically for the positional information. As shown in Section B.2..3 B.2.3 MQA is strictly more expressive than MLAFactorized, and by extension, GQA.
In summary, we have demonstrated that the expressive power of MLA with decoupled RoPE is stronger than that of the traditional GQA. However, it is worth noting that in the previously proven proposition, the Mattoocorie does not have a low-rank decomposition on the query; this differs from DeepSeek MIA. In the full MLA architecture, the query is also decomposed.
C Proof of RoPE Inner Product Invariance under Orthogonal Transformation
In this subsection, provide rigorous proof of Equation
T0 namely:
d/2
Σ RoOPE ((Urq))ea)) (aaa) U.qq(..ma)) )) RoPE5 ([V UnK((inme))
9/. K.
(L.eal) (and similarly for other Here, d is the dimension each original attention head. The notation q terms) refers to an g-dimensional vector collecting the (21 - 1)-th components from each of the original attention heads. The matrix U; is an g orthogonal matrix.
(L.cal) 4t
(Limag)
Proof. For the sake of convenience, we omit all i,j, k and let q
and qy.
Similarly, let k2,1 kg (..ral)) and k (f.iag)
The RoPE transformation, as defined by Equation 9 in the main text, applies as follows for query vector at position / and key vector at position j within the 1-th subspace:
(qz,1)* R 4r, (t0())) 9y. sin (t0;)
(qv,t) R dr, sin (t0:) + qy, cos (t0:)
1k2,,) R k cos (j0i) - ky, (j(t)
(ky,:) R = k sin() + ky, cos()
The superscript R denotes the application of RoPE. We use the shorthand = cos((1)), sin((O:). = cos(jor), and 8, sin()
The right-hand side (RHS) of Equation
10 is given by the definition of the RoPE inner product:
d/2
91,
T
kp
Σ
(q
R (qq,)) R [(kr.) (k),))R)
1=1
4/2
Σ (R) (kr,i) +((g))))((()))"
(ky.)) ((k,)))))((())))))
1=1
Let St be the (-th term in this sum:
Si (c.qa. - - (cyke, - sykn,) ) + (siqa, + oqq.1)) (syke, + c.ky,)
= kocj - 4 48, - s: s ke, 818, 41s,
1185 Ka, + s,c ky, + assa kn. a.ggggkyy
(ary + (;:s) kn,) + (s:c) c,s)) (k,ik
cos((t - -1) kn.) + sin((t j)0:) (10)) ka.))
28

---

Now, let's analyze the left- beft-hand side (LHS) of Equation (T)) Let 4 U.q. and L = U.q Similarly, let 22, = 2, and p,1 = g The 1-th term of the LHS sum, denoted S%, is:
S (((g6,) (Kg,) ((g)))) (T (K)
This has the same structure as St, just with primed variables:
S = cos((t - (( k2, + (q(,) K2,) + sin((t j)o. ky, - (dy.)) (K,)))
We need to show that the dot product terms involving primed variables are equal to their unprimed counterparts. Consider the first coefficient term:
(de,)) T ( (U,g.)) ( (U,kk,)) (U
= =92, U1kn, U,UKK, + 4n Ujkya
The last equation holds because Uj is an orthogonal matrix. This matches the corresponding term in Sp.
The same applies to the second coefficient term. In this way, we have proven that S = S for each I E (l.....d/2) This implies that the LHS of Equation (IO) is equal to its RHS:
4/2
Σ ( 21-1s1,
(.al).))) ([U,k]a-Li)) U,k ( [)
9:. R T
This completes the proof, demonstrating that the orthogonal transformation Uz applied to the g- dimensional vectors representing the 1-th 2D subspace components across heads preserves the RoPE-based inner product structure.
In practice, we leverage this rotational invariance property find a set of optimal orthogonal matrices {U,} that concentrate the principal components of the key vectors into the first few attention heads. The preceding proof reveals a critical constraint: for the inner product's 's value to remain unchanged after transformation, the same orthogonal matrix U, must be applied to both the real (21 1) and imaginary (21) components of the key vectors within each 2D subspace. This requirement precludes performing separate PCA on the real and imaginary parts. We must therefore find a single rotation that is jointly optimal for both.
Specifically, we formulate this as a joint optimization problem. First. we process calibration dataset (e.g., Wikitext-2) to collect the key activations at each layer. For each RoPE subspace I E {l.....d/ d/2}. we obtain two collections of n x g-dimensional matrices (where denotes the number of samples): the "real" parts (K and the "imaginary" parts {K,,) To find a single transformation U/ that simultaneously compresses the information from both sets into the first few heads, we proceed as follows.
Let 0.,1 = KT Ky. and 0 K KJ, Ky. be the x covariance matrices of the real and imaginary key components, respectively. Our objective is to find an orthogonal matrix U, that maximizes the variance-or energy-concenee concentrated in the first heads after rotation. This corresponds to maximizing the trace of the top-left m x m submatrix of the summed covariance of the rotated vectors. The problem is formally stated
max U/
Tr [(U) (an., + (a., Ui).) s.t. U
(29)
Here, U, is the g orthogonal optimization variable, and (.)m.. denotes the top-left m m submatrix. The solution to this trace maximization problem is obtained by performing an eigende- composition the summed covariance matrix æI.1 - og. The resulting matrix U, whose columns are the eigenvectors sorted in descending order of their corresponding eigenvalues, the optimal orthogonal transformation U..
By applying this rotation, we ensure that the principal components from both the real and imaginary dimensions of the keys are aligned and concentrated within the first few heads. Consequently, we can discard the RoPE components from the remaining heads both queries and keys while preserving the most significant positional information, thereby minimizing the performance degradation.
29

---

Group
Group 2
Group 3
Group
dim=[1,3]
dim=[2,4]
dim=[5,7]
dim=[6,8]
PCA
RoPE
NoPE
NoPE
NoPE
Figure 7: Pipeline of RoRoPE with FreqFold. RoRoPE encodes the entire frequency spectrum of all attention heads in single latent dimension, which limits its expressive power. FreqFold remedies this by clustering adjacent-frequee fre.uencyguyyy dimensions and extracting their principal components jointly, allocating dinenssonl subspace to similar features. This richer representation enables K Krope to retain far more positional information.
D FreqFold: Detailed Mechanism, Example, and PCA Efficiency
This appendix provides detailed explanation the FreqFold technique, illustrates its operation with a concrete example, and formally connects its benefits to a general principle of Principal Component Analysis (PCA) concerning structured data. This justification clarifies FreqFold's role in minimizing transformation loss towards decoupled RoPE within the RoRoPE framework (Section 4.2
D.1 Detailed Explanation of FreqFold and RoRoPE's PCA
In the RoRoPE framework, Rotary Position Embedding (RoPE) applied. RoPE encodes positional information by rotating pairs of feature dimensions. For each RoPE frequency index E {1..... d/2], the corresponding pair dimensions ([2/ - :: d], [2| d)) from key vectors are rotated. When multiple original attention heads used (say, heads), and their keylquery projection outputs concatenated, the RoPE operation for specific frequency index applies to 2g-dimensional vector segment (formed by concatenating the 1-th 2D RoPE subspace from each of the g heads). RoRoPE
to these 2g-dimensional segments, independently for each then applics PCA via matrices { frequency index l.
The core idea of FreqFold is to approximate numerically similar RoPE base frequencies as being effectively identical. For instance, if RoPE uses original base frequencies 6I, f......, 8i,, that are close in value, MDD-FreFoold might treat them all single, representative frequency 6*.
This approximation has a significant implication for how PCA applied in RoRoPE:
Without FreqFold (Standard RoRoPE PCA): For each distinct RoPE frequency index 7, a separate PCA transformation U, is learned and applied the corresponding 2gddimensional key/query segments.
With FreqFold:
original RoPE frequency indices (say lı,... .,....,.. ,ar...,,))) are grouped together by FreqFold due to their frequency similarity, the M corresponding 2y-dimensional segments are effectively concatenated. Instead of M separate PCAs on 2g-dimensional vectors, single PCA is performed the resulting M 22-diimensional vectors.
30

---

D.1.1 Illustrative Example of FreqFold
Let's consider scenario with 2 key heads, and each head has dhead 8 dimensions. Thus, there are d/2 8/2 distinet RoPE frequency indices per head, which we denote as Φ1 92, Ø3, Φ4- The total number of dimensions is 2 x 8 = 16. The RoPE angles for these 16 dimensions could be conceptualized as follows (repeating for each pair, and across heads):
Head 1 (dims 1-8): (0), Φ1), (02, Φ2), (03, Φ2), (04, ②4)
Head 2 (dims 9-16): (Ø1, @1), (Ø2, o2), (Ø3: 0s). (01, 64)
Case 1: RoRoPE without FreqFold For each frequency index Φi, RoRoPE groups the corresponding dimensions from all g 2 heads. Each such group forms 2g 2 4-dimensional vectors (across N samples).
Group for Φ1: Dimensions {1, 2} from Head 1 and {9, 10] from Head 2. PCA is applied to
these N samples of 4D vectors.
Group for Φ2: Dimensions f3, 4] from Head and (11, 12) from Head 2. PCA is applied
to these N samples of 4D vectors.
Group for Φ3: Dimensions {5, 6} from Head and (13, 14] from Head PCA is applied to these N samples of 4D vectors.
Group for Φ4: Dimensions [7.8) from Head and {15, 16) from Head PCA is applied to these N samples of 4D vectors.
Here, RoRoPE performs separate PCA operations.
Case 2: RoRoPE with 2D-FreqFold 2D-FreqFold implies we are pairing up original frequencies. Suppose FreqFoold approximates Φ1 xs d2 (calling this effective frequency 4 Φ1) and Φ3 R d4 (calling this & n = dg).
Effective Group for A This group now includes all dimensions originally associated with
Φ1 OR 22
Original d1-dimensions: {1, 2} from Head 1; {9, 10) from Head 2. (Forms 4D segment S,,)
Original ø2-dimensions: (3, 4) from Head 1; (11, 12) from Head (Forms a 4D segment Sag))
With FreqFold, these segments Se, and Sø,, are concatenated. PCA is now applied to the N samples of (4 + 4) 8-dimensional vectors formed by [S Sz,.. Effectively, dimensions
{1,2,3, 4] from Head are combined with {9, 10, 11, 12} from Head 2.
Effective Group for Ø Similarly, this group includes dimensions originally for Φ3 OR Φ1.
Original dimeiemssss (5, from Head I; [13, 14) from Head 2. (Forms Saa)
Original da-ddimensions: {) from Head 1; (15, 16) from Head (Forms S.a)) PCA is applied the N samples of 8-dimensional vectors formed by [ S..).
Here, RoRoPE with FreqFold performs PCA operations, but each operates on larger, 8-dimensional vectors which are concatenations of what were previously separate PCA targets.
D.2 Formalizing the Benefit of FreqFoid in PCA
The example above illustrates that FreqFold causes re re-grouping and concatenation of data segments prior to PCA. The benefit of this concatenation is explained by the following proposition. It states that performing PCA jointly on these concatenated segments (as FreqFold enables) is more effective at preserving variance (and thus minimizing loss) than the alternative of performing separate PCAs on the original, smaller segments and then notionally combining their outcomes.
Consider one such FreqFold merge: suppose M original RoPE frequency indices 4...., I' are deemed equivalent by FreqFold. Without FreqFold, each I, would correspond to a dataset X, (e.g., N samples of di-dimensional key segments). With FreqFold, these M datasets are concatenated into single larger dataset X Xmerged = X and PCA is applied to Xergged
31

---

Let M distinct groups of key segments X1, X2.,2...., be identified. Each Xp NNx (where p (I....,M. .....,,))))) consists of N samples of d' d'-dimensional vectors. Assume data in each Xp is mean-centered. Let Sp = X X2 Rdd xd' be its covariance matrix. FreqFold causes these M groups be merged for a single PCA operation.
Proposition 2.
Define V1 =E== Ap,1, where Xp,1 is the largest eigenvalue of Sp. This V1 represents the of variances if each of the M original groups Xp were individually reduced its single most dominant dimension.
(MM-') be the dataset formed by concatenating the features (columns) of these M groups. Lei Sconcat N-T - Z E IR (M (Mt-d') be its covariance matrix.
Let Z = X2...,
Define 1/2 = Eド== Ej=1Ht Hj, where μr 2 P2 ... μM are the M largest eigenvalues of Sconcat- This V2 represents the variance captured if the concatenated data Z is reduced to M dimensions using PCA.
Then, the variance captured by the joint PCA on the TreqFold-merged data (V2) is greater than or equal the sum of variances from optimally reducing each original group one dimension (V1):
V1
This proposition explains that FregFold's strategy of enabling PCA over larger, concatenated segments (formed by merging data from RoPE frequencies deemed similar) is mathematically favored for variance preservation compared to separate, more fragmented PCAs.
D.3 Proof of Proposition]
The objective is to prove that 1/2 2 V1, using the notation from Proposition 2 The proof strategy is to construct specific M -idimensional subspace for the concatenated data Z. We show that the variance captured by projecting Z onto this particular subspace equals Vi. Since the PCA procedure yielding V22 finds the optimal M dimeensioon subspace maximizing captured variance, V/2 must be least V1.
Let Ap,1 be the largest eigenvalue of Sp (covariance of Xp), and Wp,1 E RRd be its corresponding eigenvector. So, Wp 2p.1 and wT. Wp,1 1. The variance Xp.1 wh, Sp
V1 = Epen Ap.1.
For the concatenated data Z, V/2 Σ5= M... By Ky Fan's theorem for matrix eigenvalues:
1/2
max UER UTU-T
Tr(UT Sconcat
where U's columns form an orthonormal basis for an M- --dimensional subspace of I. M'd" Construct U
[ui. /ui..... ..... uM (MM-d))M For (1....,MM, M), define uz T M.d.
Oax1
u = Wp
(as the p-th block of size d')
Orx1
The set (ui,..., uhr} } orthonormal. The variance retained by projecting onto the subspace of U is:
TT Scaneat U*) Σ (un)) Soomcat
Let Sgr be the r- block of Sconcots where Sqr N-1 X,. Note Spp Sp- Each term (ut)² 2 Spm = S. = Ap. So, (T((U*)"S =
En Ap.1 Vj. Since V2 is the maximum possible variance:
2/ " = V/2
Sconeat
Thus. V/2 > V1. This proves Proposition 2
32

---

D.4 Discussion on the Trade-off in FreqFold
While Proposition demonstrates a clear benefit of FreqFold in terms of PCA Sepriieccy-seiiccalllly that merging M original frequency groups allows for greater variance preservation when reducing to M dimensions- it is crucial to acknowledge an inherent trade-off. The foundational assumption of FreqFold is the approximation of numerically similar RoPE base frequencies as effeetivvlyy identical. This approximation, by its very nature, introduces a degree of deviation from the original, precise RoPE formulation.
The extent of this deviation, and thus the potential loss in the fidelity of positional encoding, typically correlates with how aggressively frequencies are grouped. A larger M or looser criterion for similarity when grouping frequencies can amplify this approximation error. Consequently, while increasing the dimensionality of vectors undergoing PCA is beneficial from the perspective of PCA variance capture as shown by the proposition, may simultancously increase the lossiness of the RoPE approximation itself. Therefore, the practical application of FreqFold requires careful balancing act. The parameter M (representing the number of original RoPE frequencies treated as effective frequency for PCA purposes) the specific grouping strategy for frequencies must be chosen to optimize this trade-off.
E
Balancing Key- Value Norms and Low-Rank Approximation
This appendix elaborates on the Key-Value (KV) balancing technique and the subsequent joint low- rank approximation applied to the NoPE (No Positional Encoding) components of the keys and the values, as mentioned in Section 4.3 of the main paper. After the RoRoPE procedure (Section (.22) the key projection matrix W K is effectively split into two components: rope E Rd corresponding to the single head that retains RoPE, and W R(o-1)dxDD corresponding to the remaining g - nope
DK
head components that do not use RoPE. The value projection matrix is denoted as WDV [RgdxD,
E.1 KV Balancing: Purpose and Formulation
The primary goal of KV balancing is to ensure that the principal component analysis (PCA), when applied jointly the NoPE key and value activations, is not disproportionately in- fluenced by components with larger norms. We observed that the activations derived from W WDK nope (i.e., KNoPE, WDD xx) often have significantly larger average norm than those from DDV (i.e., nope
Purpose
Vi DV x²). Without balancing, PCA would predominantly capture the variance within the NoPE key components, potentially neglecting important variations in the value components.
To address this imbalance, we introduce scaling factor This factor is computed as the ratio the expected 1.2 norms of the NoPE key activations to the value activations, based on calibration dataset:
Formulation
E, l: nope
DK
x
(30)
WDV x
where x, E R the t-th input token.
While the main paper states scaling WDKK nope by 1/a and WUK by a for mathematical equivalence in the model's output, for the purpose of deriving the PCA projection, we effectively use scaled NoPE key activations. That is, the activations used to compute the PCA basis are KNoPE,: = 1/a WDDK x
and DVV xt. This ensures that the PCA process considers features from keys and values more equitable footing with respect to their magnitudes. The subsequent low-rank decomposition will then be applied to W WDK and DDV, using the PCA basis derived from these balanced activations.
E.2 Joint Low- Approximation of NoPE Keys and Values using PCA
After determining the scaling factor a, proceed to compress the projection matrices associated with the NoPE keys (WDDK) and all values (WDV) jointly.
The process is as follows:
33

---

1. Collect Calibrated Activations: A small calibration dataset (WikiText-2) is used. For each input
X2 from this dataset, we compute the scaled NoPE key activations KKOPE, and the value activations Ve- These are concatenated to form combined activation vectors:
(KNOPE.. V/:
CNoPE,
(2g-11)
(31)
2. Perform PCA: PCA is performed on the set of collected combined activation vectors (CNOPF.)) This involves computing the covariance matrix of these vectors and finding its principal components. The eigenvectors (corresponding to the largest eigenvalues) are selected to form the columns of a projection matrix RKv E (2((2g1)dd)xr) where Tik is the reduced rank. This matrix Rx captures the directions of highest variance in the (balanced) combined NoPE key and value activation space.
W DK none W
3. Low-Rank Decomposition of Projection Matrices: Let WDKV
be the initial projection matrix that transforms the input X into intermediate NoPE Key and Value
representation CNoPE, WDKV X2 Further, let WUKV = - (WUK nope 0
I 2hdx
0
wev
represent the subsequent collective projection matrix that takes CNOPE. and processes it to produce the actual keys and values required by the attention mechanism for the NoPE components, where W.C E gidxx and W Ro 1)d are two parts of WUK hat participate in and do not participate in the RoPE computation, respectively. The original sequence of operations for these components can be expressed as WUKV WWWDVV' X₂ E R2dd in which the first hd elements correspond to the keys and the following hd elements correspond the values.
To introduce a low-rank bottleneck, we modify both WDV and WUKV using the PCA projection matrix Rgv.
The initial projection matrix WDKV is transformed into WDKV' E ERR WDKV's WDV
(32)
This matrix WDRV' takes the original input x and projects it into compressed The Tre-dimensional latent space, which is the actual content stored in the KV cache for the NoPE components.
The subsequent projection matrix WUKV is transformed into WUKV' E 2hdxTin:
WUKV' = WUKV RKv
(33)
This new matrix WUKV now takes the compressed latent representation as input and produces the final representations for the NoPE components that are used in the attention calculation. As we can see, WUKV' is actually the concatenated form of WUK and WUV in MLA:
(WUK (WUv)
WUKV'
(34)
This joint decomposition allows for a more holistic compression by identifying shared latent structures between NoPE keys and values, guided by the balanced PCA.
F Theoretical Analysis of Balancing Scheme for PCA in Matrix Decomposition
F1
Theoretical Analysis
In this section, we focus on situation where two matrices. W₁ and W., are concatenated together for dimensionality reduction. A potential issue arises when the matrices have different magnitudes, which leads to an imbalance the relative errors during reconstruction. The relative error for one matrix may be small while for the other it may be large. To address this, we propose scaling both matrices by their Root Mean Square (RMS) before concatenating. This scaling ensures that the relative errors for both matrices are approximately equal. The goal of this section is to theoretically prove the validity of this approach by formulating a loss function that reflects the relative errors and demonstrating that the scaling values mentioned will minimize this loss.
34

---

F111 Problem Setup
We are given two matrices W5 and W. of dimensions m x and m x respectively. The matrices are concatenated as follows:
Wx
W
W
(35)
a
where a and b are scaling factors that need to be optimized. These scaling factors are introduced to balance the relative errors during reconstruction. We perform PCA on the concatenated matrix W. The SVD of the matrix W is given by:
rank(W) Σ ==1
W = UVVV
(36)
where U, E, and V are singular vectors, singular values, and right singular vectors, respectively. The r-rank approximation matrix W is given by:
W U,E
(37)
W
W b
W
(38)
where W and W are the low-rank approximations of Wi and W,, When the scaling factors a and b = 1, the matrices W. and W, are not scaled and directly undergo low-rank approximation. The matrices W and W represent the low-rank approximations of We and W, without the effect of any relative scaling based on the Frobenius norms.
To address the imbalance in the relative errors during reconstruction, we define the loss function to equally quantify the relative errors between the original and the approximated matrices. The loss is the sum of the relative errors for W₁ and W u. given by:
WK W' ll +1
WV 1 RMS(Wv)
Wv
C(a. b) RMS(WK)
(39)
Next, we will derive the condition when the loss function reaches its minimum, specifically when the scaling factors and b satisfy the ratio:
a RMS(WA)
(40)
5 RMS(W.)
F.1.2 Proof of Proportionality
Before proceeding with the derivation. we first need to verify the proportionality between and b. which implies that the absolute values and do not affect the conclusion. As long as the ratio a: b is fixed, the final loss will be the same, regardless of the absolute values of and b.
Consider scaling the factors and b by constant 0: a' ca. D' cb.
The concatenated matrix with new scaling factors is:
Wx W.
Wx
W. cb
- W = -UVV
W new
(41)
Thus, the SVD of W new has the same left and right singular vectors U and V as W, but its singular values are scaled by The low-rank approximation of Wnew is:
W new Ur {
V,T = = (U.,,V)) - t
(42)
Comparing this with the block definition of W new. we identify:
W Wkoeow = W/ W v.new W
(43)
The loss function L(a, b) is invariant under scaling + ca and b + ob for any > 0. Thus, it depends only on the ratio g. Therefore, in the following derivations, we will always assume that b = 1.
35

---

E.1.3 Verify Minimum Conditions
We now verify the minimum condition when a
R RMS(W.)
For convenience, we define
We.
A, Wx
B C, W D.
W
1(a, ||A- -lE
||A mink + | DI?
|
mnk
(44)
dC(a,
d
||A CIl.E man + 1 - 1DD mym
(45)
da
da
||A
BB
= -IIA - - da ||||
d
mnkk + mnrk
(IIAA - C1I12))
(46)
|
4 (1IB D1l:))
m. 2o
(47)
BB
when b 1,
RMS(WW)
KMS(W.))
d.(a. 1)
2/1A1 C mng
mny d. (IIA - + || - DIP:)
da
(48)
da
a||A|/?
|B
2||AA C
mn da
d (1IW - - |B17
(49)
a||-||
d mny da
Ei-rt] of
A||AA - C
a||A||/
mIn
|B|
(50)
(51)
W -VV
(52)
dW = UdEV" + +UEAV* + dUEVV U UWdWV- = dE
(53)
+UT
(54)
Since d(UTT) aUT UJ du 0 and d(vTV)- TV VT 0, U dt and VT are skew-symmetric matrices, and the diagonal elements are zero therefore,
do: T dW da u
A 12 of Opu
(55)
da
d.(a. 1) da
A - CI mnk + Exertl
brank(W) 2mn, do. de |BB|
(56)
al|A11/
1 10---1 CI
wank(() o [A 0.:
all All
mnk
B
(57)
rank(W)
A = W
1 T nx]
(58)
C= W ni]
a of I: nd]
(59)
i=1
rank(W)
A-C A-C=w"l,
Σ
nd)
(60)
1=41
therefore,
rank(W)
-
Σ
ou i nelo: mz]
(61)
36

---

Table 2: Composition of the training dataset.
Dataset
Sampling Weight
fineweb-edu-dent
0.70
cosmopedia-v2
0.15
python-edu
0.06
open-veb-mate
0.08
stackoverflow
0.01
0, [A, = o. [on]o: nx]oil: ns]
(62)
when b = 1, a-
mis( W.
we have:
rus(W.) (
dL
(63)
Thus far, we have demonstrated that when the loss function
WK W 2 RMM(SW*)
Wv W/ 2 RMS(WV)
L(a, b)
(64)
the ratio
a RMS(W))
(65)
b RMS(W.)
achieves the best compression performance.
G Experimental Settings of Fine-tuning
Following the experimental setups of MHA2MLA, we fine-tune our models using the prtraining corpus from SmolLM [40]. The dataset comprises Fine FibeWeb-Ead Web-Edu-Dedup [11] Cosmopedia- v2 ---t a synthetic dataset generated by Mixtral
Datasets
42). Python-Edu from StarCoder 43). Open- Web-
445 To ensure fair comparison with the MHA2MLA baseline, we constructed our training dataset using the same data composition strategy. Specifically,
Math 44 and data from StackOverflow
replicate the dataset mixing ratios used in the MHA2MLA setup to maintain experimental consistency, which shown in Table 2
Hy Hyyeepaaraeeees The fine-tuning hyperparameters for models of all sizes are listed in Table 3 the table, entries with a slash (/) indicate a two-step training process.
In
Table 3: Training details across different models.
SmolLM 1B7
LLaMA2 7B
-68.75% %
-87.50%
-68.75%
-87.50%
-92.97%
Batch size
64
64
64
64 64
256 64
Learning rate
le-4
le-4
2e-5
2e-5 2e-5
le-4 / 2e-5
Tokens
300M
IB
500M
2B 1B
5B / 1B
Warmup ratio Ir scheduler
0.03
0.08
0
0 0.03
0 0.03
constant 2048
constant 2048
constant 4096
constant / cosine 4096
constant / cosine 4096
Sequence length
H vLLM Benchmark Details
In Section 5.4, demonstrated the speedup achieved by TranMMA-wwiee compresses 92.97% of the KV cache- compared to the original LLaMA-2-75 model. This section provides the detailed methodology and hardware configurations for this benchmark.
To account for the effects of both the prefilling and decoding stages, we adopt setting where the input and output lengths are equal. For instance, with total context length of 1k, we set the input
37

---

length to 512 tokens and the output length to 512 tokens. Most experiments are conducted using 100 requests to compute the average throughput. However, for shorter context lengths such 1k, inference is extremely fast, leading to some timing fluctuations. To mitigate this, we increase the number of requests to 1000 for more stable measurements.
While the original LLaMA-2-7B model supports a maximum context length of 4096 tokens, we extend this limit to 32k tokens in our evaluation. Detailed throughput results are presented in Table
On a GPU with 165.2 TFLOPS of compute and 24GB of memory, the LLaMA-2-7B model runs out of memory when the context length reaches 16k tokens. In contrast, TransMLA sustains throughput of 414.41 tokens per second under the same conditions. On a more powerful GPU with 320 TFLOPS and 64GB of memory, the throughput of TransMLA will improve further with the release of future optimized versions the framework tailored for this hardware.
employ development version of the vLLM framework. We anticipate that
Table 4: Throughput comparison between LLaMA-2-7b and TransMLA at varying input lengths and number requests.
Throughput output tokens/s)
Context Length
Requests
Model
165.2 TH224GB
312 TF40GGB
320 TF664GB
LLaMA-2-7b
653.81
1579.26
1249.13
IK
1000
TransMLA
3043.65
4062.43
1798.17
LLaMA-2-7b
352.85
850.14
789.31
2K
100
TransMLA
2241.87
2577.01
1080.73
LLaMA-2-7b
173.09
441.37
442.63
4K
100
TransMLA
1318.78
1926.15
1021.03
LLaMA-2-7b
85.80
218.51
216.66
8K
100
TransMLA
832.69
1118.18
870.15
LLaMA-2-7b
OOM
110.58
112.13
16K
100
TransMLA
414.41
601.36
483.22
LLaMA-2-7b
OOM
38.32
55.69
32K
100
TransMLA
OOM
243.81
278.09
Inference Speed Ablation
The strength of the MLA design extends beyond its of low-rank compression- concept that has been explored in prior works. Its true innovation lies in the Absorb operation, which plays a pivotal role enhancing efficiency. During the training phase (which compute-intensin intensiven this operation enables attention computations to be performed using compact 192-dimensional representations per head, resulting in high computational efficiency. In the inference phase (which memory-bounde utilizes shared 576-dimensional KV cache across heads by absorbing the key projections into the query projection matrices. This mechanism significantly reduces memory access overhead, both on high-bandwidth memory (HBM) and on-chip SRAM.
Table 5: Comparison of the impact of num ju__v__ heads, head_ dim, and parameters on inference speed for GQA and MLA. To climinate the influence of MLP layers on the results, the speed test was conducted using models composed of 32 attention layers only.
Method
hgv
d
KV Cache
Param (M)
Throughput (tokens/s)
MLA (with Absorb)
1#1
576
576
48.5
4151.17
MLA (without Absorb)
32*2
192
12288
48.5
248.85
4*2
128
1024
37.7
2714.80
GQA
8*2
128
2048
41.9
1436.92
16*2
64
2048
41.9
1118.83
As shown Table (rows 2-3), disabling the Absorb operation leads 16.7x degradation in MLA's inference speed at a 16k context length, underscoring its crucial role in achieving efficiency.
38

---

This highlights the necessity of supporting the Absorb operation, as it is indispensable for realizing the practical benefits of MLA real-world deployment. The importance of the Absorb operation further corroborated by findings from Palu [20]. In their Figure 5, they demonstrate that removing RoPE allows the Absorb operation deliver 6.17x speedup at a 64k context length. However, when RoPE is retained, Absorb must be disabled, reducing the speedup to 2.91 Unfortunately, since Palu does not support RoPE decoupling, it cannot apply the Absorb operation to widely used architectures such as LLaMA, Qwen, and Mistral.
In Figure 5 we only present the inference speed of the LLaMA-2 model based on MHA. Since GQA already requires fewer KV cache entries than MHA, it is natural to ask whether TransMLA can further improve the inference speed of GQA-based models. To investigate this, we conducted experiments using LLaMA-3-8B, whose key and value projections each contain 8 heads with 128 dimensions per head, resulting in total KV cache size of 2048 dimensions. We applied TransML to reduce the KV cache to 576 dimensions and measured the inference speed across different GPUs. As shown in Figure TransMLA still achieves approximately a 3x speedup, demonstrating its effectiveness even when applied to GQA architectures.
2.4%
3x
1.5
(a) 165 TFlops 24GB
(b) 148 TFlops 196GB
(c) 1979 TFlops 80GB
Figure 8: Inference speedups with TransMLA compared to the original LLaMA3-8B model three Al accelerators. Context length represents the total sequence length.
J Evaluation on GQA Models
The goal Table is to provide a fair and consistent comparison between TransMLA and prior relevant work. Therefore, selected the recently published MHA2MLA (2] our bascline, also focuses on RoPE decoupling and KV cache compression. Following their experimental setup, and LLaMA-2-77 In fact, TransMLA is fully
used two MHA-based models- MatIsssSmooMMMMMM.. compatible with GQA. As stated in related work, both MHA and MQA can be regarded as special cases of GQA. Specifically, when the number of query heads equals the number of key-value heads in GQA, it effectively reduces to an MHA. In this section, we further extend our evaluation to several mainstream GQA models, and the corresponding results without training are presented in Table d
Table 6: Evaluated with Wikitext-2 Perplexity (PPL) of Converting GQA Models to MLA (lower better). Original PPL refers the perplexity of the original GQA model without conversion. Decoupled RoPE PPL represents the result after applying RoRoPE and FreqFold to reduce decoupled RoPE loss. Compressed KV PPL. indicates the perplexity when the KV cache compressed 576 with the help of Balance-KV.
Model
KV Cache
Original ppl
Decouple RoPE ppl Compress KV ppl
Llama-3-88 3-BB
-71.875%
6.1371
8.3997
18.35
Qwen2.5-7B
.43.75%
6.848
7.3059
7,9812
Qwen2.5-72B
-71.875%
4.2687
4,6931
7.7172
Mistral-7B
-71.875%
5.3178
5,5915
7.0251
Mixtral-8x7B
-71.875%
3,8422
4.1407
5.8374
Gemma-2-9B-IT
-85.94%
10.1612
11.421
21.626
GPT-OSS-20B -43.75%
10.2563
9.2175
9,4072
From Table 6 we can observe that for GQA models-ranging from 7B to 70B, whether dense
models or MoE models-TrranMMI... consistently ineurs only minimal performance loss. It's worth noting that MLL is specifically optimized for latent dimension of 576, which adopt a unified
39

---

compression target across all models. Since GQA models typically maintain smaller KV caches than MHA, converting them MLA results less information loss, making the transition both effective and efficient.
K Distillation Experiments
In Table 1, our work follows the experimental setup of MHA2MLA and uses the smollm-corpus for pretraining. However, this dataset differs from the original training data of LLaMA-2-7B. To investigate whether the observed improvements stem from additional learning or genuine recovery of model capabilities, we adopted the data distillation approach proposed in the LLM-QAT [46]. Specifically, we generated a new dataset from the original model to guide the distillation of the converted model. For each model, generated distillation datasets containing 14,000 samples, with each sample consisting of 2,048 tokens. These datasets were used to distill the converted model by minimizing the error between the output of each converted MLA attention module and the corresponding GQA outputs of the original model, as well as aligning the final token predictions. Each atfention module trained for 10 epochs and hyperparameter tuning conducted 10 evenly spaced learning rates within the range [1e-3, 2e-2]. The results are presented in Table 7
Table 7: Commonsense reasoning accuracy for distilling converted MLA-based models from original
GQA models
Model
KV Cache
AVG
ARC
HS
MMLU
OBQA
PIQA
WG
LLaMA-2-76
59.81
59.31
73.13
41.35
41.60
78.40
65.04
TransMLA
-68.75%
59.89
59.82
72.43
41.67
42.60
78.18
64.64
Qwen2.5-7b
63.38
64.45
77.49
47.26
45.40
79.76
65.90
TransMLA
-71.875%
63.12
68.13
74.90
46.33
44.20
79.11
66.06
Mistral-7b-v0.1
64.83
66.71
79.48
45.67
45.60
82.86
68.67
TransMLA
-50%
64.48
66.22
79.00
45.53
45.60
82.70
67.80
According to Table 7 aligning the converted models using only the distilled data from the original models allows them to closely match the performance of the original models. This confirms that the observed performance improvements are not due to additional learning but reflect genuine restoration. We plan to distill more data and complete experiments additional models and compression rates in the future.
L Ablation Study: Calibration Dataset
To investigate the impact of the calibration dataset on the effectiveness of TransMLA, conducted an ablation study focusing on different calibration data sources and dataset sizes. As shown in Table the number of calibration samples noticeably affects the conversion quality, with 64 samples achieving the best overall performance. Using different datasets for calibration also introduces slight variations in results; however, all calibration datasets outperform the baseline that uses randomly generated tokens.
M Ablation Study: Core Components
In Figure 3b and Figure 4b we presented ablation results on perplexity (PPL), focusing on RoRoPE, FreqFold, and BalanceKV. V. In this section, we further evaluate commonsense reasoning accuracy, with the results summarized Table
Removing RoRoPE leads to accuracy drop of approximately 20% compared to the best-performing configuration.
Disabling FreqFold (i.e., setting FreqFold results in about 6.5% reduction in accuracy.
Excluding Balance causes accuracy decline of around 10%.
These findings underscore the significance of each component in sustaining strong downstream reasoning performance.
40

---

Table Ablation study using different calibration data sourees and dataset sizes for converting LLaMA-3-8B to MLA.
Dataset
n sammlls
Decouple RoPE ppl
Compress KV ppl
8.674
155.65
8.217
38.134
4
8,192
22.859
8.050
18.688
wikitext2
16
8.164
17,499
8.217
15.966
64
8.368
14.898
128
8,404
15.080
ptb
128
8.209
19.394
c4
128
7.341
16.018
alpaca
128
7.389
17.398
random 128
15.693
44.956
Table 9: Ablation study of removing each method and its impact on commonsense reasoning performance of the converted model without training.
Method
Avg.
MMLU
ARC
PIQA
HS
OBQA
WG
LLaMA-3-8B
63.92
46.13
65.81
80.79
76.26
45.40
69.14
TransMLA
55.12
37.47
54.31
73.78
65.72
38.80
60.62
RoRoPE
35.59
25.64
28.91
55.77
28.06
24.00
51.14
FreqFold
48.63
31.19
45.24
69.37
56.18
34.60
55.17
BKV
44.91
29.19
37.81
65.56
49.89
32.20
54.78
In the main paper, the default scaling factor BalanceKV is computed as the average L2 norm of KROPE divided by the average L2 norm of V. Table 10 presents an ablation study the impact of the BalanceKV scaling factor on the performance of the converted model.
Table 10: Ablation study on the impact of the BalanceKV scaling factor on the performance of the converted model. Lower perplexity (PPL) indicates better performance.
BalanceK Scaling Factor
Lora PPL (WikiText2)
None
28.8341
(kiroee_me
Grrope_menn___a mean norm
4) Umean pormm porm
18.8790
(kiroge mean, norm X 2) Umean porm
15.3617
Krope mean, _ooeinnomm/aamm Umean porm
15.0796
Arrope_mcan sace__ma___noom (Umean_ norm x 2)
17.8757
Kroppe mean norm (Umean_none norm
4)
27.9021
Optimized per-layer
13.7945
On LLaMA-3-8B, completely removing BalanceKY results in a 3.8-point increase in perplexity, highlighting its critical role in maintaining model quality. We also observed that model performance sensitive to the specific scaling coefficient used in BalanceKV. The default value reported in our paper already yields strong overall performance. To explore this further, we performed a per-layer grid search over the range [0.3, 2.0] with 20 intervals to minimize the output difference from the original model. This tuning produced 1.2-point reduction in perplexity compared the default configuration, suggesting that layer-wise coefficient optimization is both feasible and beneficial for further performance improvement.
N
Evaluation on the LongBench Benchmark
TransMLA achieves KV cache compression by reducing the dimensionality of token representations, making it orthogonal methods such as token pruning and quantization. However, as sequence length increases, the compression error for each token accumulate. To evaluate the effectiveness
41

---

of TransMLA on long-come long-context tasks, we conducted experiments using Qwen-2.5-7B, which supports maximum positional embedding length of 131,072. The model was converted into MLA-based architecture, and the original model was then used to distill the converted one. We report the performance of both the original and converted (distilled) models across multiple long-context tasks from the LongBench benchmark in Table
Table 11: Performance multiple long-text tasks and inference speed (tested 16k-length context) of Qwen after converting GQA to MIA and applying further distillation. Abbreviations: SD-QA Single-Document Question Answering, MD-QA- Multi-Docummen Question Answering, Sum.
Summarization, FSL = Few-Shot Learning, Code = Code Completion, Avg. = Average across all
tasks.
Model
Thpt.
SD-QA
MD-QA
Sum.
FSL
Code
Avg.
Qwen2.5-7B 1233.04
20.29
25.73
24.46
62.11
63.10
39.14
TransMLA
1925.19
16.50
19.02
18.51
44.20
38.39
27.32
TransMLA-dise
1925.19
18.74
28.26
22.39
61.00
59.21
37.92
While TransMLA provides substantial speedup (approximately 1.56x throughput improvement), its zero-shot performance declines significantly when additional training is applied. Although distillation can partially recover accuracy, performance restoration becomes increasingly difficult on long and complex reasoning tasks. Future work will focus reducing the compression-induced degradation TransMLA, aiming to achieve stronger zero-shot generalization on challenging long- context benchmarks.
Extension to Multi-Group MLA
FlashMLA supports only fixed configuration of 64 RoPE dimensions and 512 non-RoPE dimensions for its KV cache, which limits compression flexibility. To address this, we divide the query heads into multiple groups, converting each group into an independent MLA module.
experimentally verified this Table 12 under identical KV cache conditions, ungrouped approach outperforms the two-group version. However, the grouped method enables the model to retain more KV cache, so performs better than the original TransMI.A with single 576-dimensional KV cache.
Importantly, this multi-group strategy remains compatible with FlashMLA inference. We tested the inference speeds of MHA, GQA, two-group MLA, and single-group MLA under 16k context length. Since FlashMLA does not support inference with KV cache dimensions other than 576, we excluded the single-gue goolu 64+1022 MLA configuration. As shown in [2] the wwo-group MLA achieves 6.2x speedup over MHA and 1.58x speedup over GQA, while the single-group group MLA remains approximately 1.9x faster than the two-group version.
Table 12: Performance comparison of multi-group MLA under different RoPE/NoPE dimensions and group settings. Abbreviations: rope num_grp number of MLA groups, ppl perplexity, Thpt. throughput in tokens per second (tested at 16k context length).
rope_dim- dim RoPE dimension, nope dim = NoPE dimension,
Model
rope dim nope dim num grp
ppl
Thpt. (tok/s)
5.4732
11776
64
512
41.6135
141246
LLaMA-2-7B
64
1024
24.0931
64
512
31.2177
73098
6.1371
46247
64
512
25.8047
141246
LLaMA-3-8B
64
1024
8.9768
64
512
2
12,9472
73098
42

---

P
Case Study
To provide an intuitive understanding of TransMML'ss impact on model performance, this section presents several examples from vLLM's docs. We compare the outputs of three model variants: (1) model with 92.97% of its KV cache compressed without any fine-tuning: (2) a model pretrained on 6B tokens, as detailed in Table and (3) model fine-tuned for one epoch on the SmolTalk dataset, following the setup described
The results are summarized in Table
As shown in Table 13 even without any additional training, the compressed model is still able to produce coherent and meaningful responses. This demonstrates the effectiveness of techniques such as RoRoPE, FreqFold, and BKV-PCA in significantly mitigating performance degradation. Moreover, with a modest amount of pretraining or supervised fine-tuning (SFT), the model's performance improves substantially. These findings highlight TransMLA's potential as general framework for converting various GQA models into MILA models, with promising prospects for aligning with the performance of advanced systems like DeepSeek R1.
Table 13: Examples from different model configurations. Red indicates input; black indicates output. "w/o Training" denotes the TransMLA- A-compes-e model (92.97% KV cache) without further training. "Pre-Training" and "Fine-Tuning"" show outputs after pretraining 6B-token corpus and SFT on SmolTalk [88), respectively.
Model
Prompt & Generated Text
w/o Training
Hello, my name is Katiu, my father's dog, the pet of the 3600 year-old tribe, Kint. The Kangs were part of race of reptiles. A small handful Hello, my name is Sasha and I am in third grade at Meadows. You may be wondering what this article about. Well, have been doing lot of research on the water cycle and decided to write about it.
Pre-Training
Fine-Tuning
Hello, my name Emily, and I'm 20-year-old college student. My hobbies include painting, writing, and photography. also enjoy playing the guitar.
w/o Training
The president of the United States is elected by the legislature. The legislature controls the national armed forces, but only provides the funds to establishing national guard.
Pre-Training
The president of the United States is elected to a four-year term by the people of each state in general election held every four years the Tuesday following the first Monday in November.
Fine-Tuning
The president of the United States is not position to be taken lightly. This person is the chief executive of the United States of America, and has immense power and influence.
w/o Training
The capital of France Paris. Its geographical position in the Iberian Plain of France, Spain, Spain, and Morocco are the four largest cities. This region is located in Asia, Spain and Morocco.
Pre-Traine Pre-Training
The capital of France is a major business city and it is favorite destination for businesses from all over the world. It has a strategic location in the heart the European Union, which makes it one of the most popular cities in Europe.
Fine-Tuning
The capital of France is Paris, and it is one of the most popular tourist destinations in the world. It is city that offers something for everyone, from art and history to food and fashion.
w/o Training
The future of AI is in serious risk to create major breakthrough in this emerging phenomenon
the history of artificial intelligence.
Pre-Traie neeTrrinnig
The future of AI is looking bright. With advancements in technology and the increasing availability of data, AI is expected to become more intelligent and capable of performing even more complex tasks.
Fine- Fine-Tuning
The future of AI The future of AI is more nuanced and complex than we might think. Here are some potential developments that could shape the future of Al.
43


Reviews:
[{"review_id": "LJqtsHQq0r", "reviewer": "Reviewer_qHce", "summary": "The authors propose a comprehensive method for converting models based on MHA, as well as simple variants such as GQA, into MLA-compatible architectures, aiming to achieve enhanced expressiveness and overall efficiency. The proposed approach first performs Key fusion, followed by tailored handling of RoPE-related Query and Key components to better adapt them to the specialized RoPE mechanisms in MLA, namely RoRoPE and FreqFold. Finally, a joint PCA-based dimensionality reduction is applied to the Key and Value matrices, guided by the distributional characteristics of their weights. The result is a model that can be effectively adapted to MLA. The authors validate their approach by conducting conversion experiments on models of two different scales, followed by fine-tuning on a small amount of tokens, demonstrating the effectiveness of their method in restoring model capabilities. Furthermore, ablation studies are provided to underscore the necessity of the proposed RoPE handling and weight-based dimensionality reduction strategies.", "questions": "Could the authors provide updated results using more recent models, and align the recovery training data with the original pretraining data to enhance the credibility and relevance of the experimental section?", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n1. The authors design an effective training pipeline that enables the transformation of non-native MLA models into their MLA-compatible counterparts. The empirical results validate the method's effectiveness. In particular, the RoPE adaptation and KV compression techniques offer valuable insights for researchers working on MLA conversion and optimization.\n\n2. The manuscript presents results across different compression scales. Notably, at lower compression ratios, TransMLA is able to retain strong performance even without additional training.\n\n3. The authors employ well-designed visualizations and empirical studies to convincingly demonstrate the necessity and effectiveness of their algorithmic components.\n\n**Weaknesses**\n\n1. A major limitation of the manuscript lies in the choice of baseline models. Both SmolLM-1.7B and LLaMA-2-7B are based on MHA, which, according to prior work and empirical observations, tends to be particularly amenable to KV compression. In contrast, more recent and mainstream models often adopt GQA or similar attention variants, which are known to present greater challenges for compression, especially in the KV cache. As such, the experimental results presented may have limited relevance to current model architectures.\n\n2. The capability restoration experiment for LLaMA-2-7B uses the same training data as that used for SmolLM-1.7B. However, LLaMA-2-7B was trained in the early 2023 using a different dataset. The data used for post-conversion fine-tuning may thus be of higher quality than the original pretraining corpus, implying that the observed gains could result not from genuine restoration, but from additional learning. This undermines the interpretability of the experimental results as a strict recovery study.", "was_revised": true, "final_justification": "The authors have addressed my concern regarding the limited scope of model testing. However, the issue of insufficient adaptability to GQA models remains. Furthermore, the work is oriented more toward theoretical exploration rather than the proposal of an engineering-driven method. Nevertheless, I believe this manuscript carries sufficient significance and value, and its level of completeness, particularly after the addition of the supplementary experiments, is satisfactory.", "rebuttal": "Thank you for the insightful questions.  Below we provide point-by-point responses to your comments. All suggested improvements will be incorporated into the camera-ready version.\n\n**Q1. The experiments focus on MHA models, but most current models use GQA, which is harder to compress.** \n\n**A1.** As stated in lines 29–36 of the original paper, both MHA and MQA can be seen as special cases of GQA. Specifically, when the number of query heads equals the number of key-value heads in GQA, it effectively becomes an MHA. Our goal is to compare TransMLA with prior relevant work under consistent and fair experimental settings. To this end, we selected the recently published MHA2MLA [1] as our baseline, as it also targets RoPE decoupling and KV cache compression. Following their experimental setup, we used two MHA-based models—SmolLM-1.7B and LLaMA-2-7B. Moreover, **TransMLA is fully compatible with GQA**, all of our ablation studies were performed on LLaMA-3-8B, a GQA-based model. \n\nNevertheless, as you suggested, we extended our evaluation to several mainstream GQA models. The results are reported in Table 1.\n\n**Table 1.** Evaluated with Wikitext-2 Perplexity (PPL) of Converting GQA Models to MLA (lower is better). Original PPL refers to the perplexity of the original GQA model without conversion. Decoupled RoPE PPL represents the result after applying RoRoPE and FreqFold to reduce decoupled RoPE loss. Compressed KV PPL indicates the perplexity when the KV cache is compressed to 576 with the help of Balance-KV.\n\n| Model | Original ppl | Decouple RoPE ppl | Compress KV  ppl | KV Cache |\n| --- | --- | --- | --- | --- |\n| Llama-3-8B | 6.1371 | 8.3997 | 18.35 | -71.875% |\n| Qwen2.5-7B | 6.848 | 7.3059 | 7.9812 | -43.75% |\n| Qwen2.5-72B | 4.2687 | 4.6931 | 7.7172 | -71.875% |\n| Mistral-7B | 5.3178 | 5.5915 | 7.0251 | -71.875% |\n| Mixtral-8x7B | 3.8422 | 4.1407 | 5.8374 | -71.875% |\n\nFrom Table 1, we can observe that for GQA models—ranging from 7B to 70B, whether dense models or MoE models—TransMLA consistently incurs only minimal performance loss. It’s worth noting that MLA is specifically optimized for a latent dimension of 576, which we adopt as a unified compression target across all models. Since GQA models typically maintain smaller KV caches than MHA, converting them to MLA results in less information loss, making the transition both effective and efficient. Due to time constraints, we only present the results of the conversion without additional training. We plan to include the full alignment performance for all models in the camera-ready version.\n\n**Q2. The restoration experiment for LLaMA-2-7B uses different data from its original training. Could the improvements come from additional learning rather than true recovery?**\n\n**A2.** As mentioned in A1, our work follows the experimental setup of MHA2MLA and uses the smollm-corpus for pretraining. We noted that the pretraining corpus for LLaMA-2-7B has not been publicly released. To address this concern, we adopted the data distillation approach proposed in the LLM-QAT [2] paper, generating a new dataset from the original model to guide the distillation of the converted model. For each model, we distilled datasets containing 14,000 samples, with each sample consisting of 2,048 tokens. These datasets were used to distill the converted model by minimizing the error between the output of each converted MLA attention module and the original model's corresponding GQA outputs, as well as aligning with the final token prediction. Each attention module was trained for 10 epochs, and we perform hyperparameter tuning over 10 evenly spaced learning rates in the range [1e-3, 2e-2]. The results are presented in Table 2.\n\n**Table 2.** Commonsense reasoning accuracy for distilling converted MLA-based models from original GQA models\n\n| Model | KV cache | avg | arc | hellaswag | mmlu | openbook_qa | piqa | winogrande |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| LLaMA-2-7b | — | 0.5981 | 0.5931 | 0.7313 | 0.4135 | 0.4160 | 0.7840 | 0.6504 |\n| TransMLA |  -68.75% | 0.5989 | 0.5982 | 0.7243 | 0.4167 | 0.4260 | 0.7818 | 0.6464 |\n| Qwen2.5-7b | — | 0.6338 | 0.6445 | 0.7749 | 0.4726 | 0.4540 | 0.7976 | 0.6590 |\n| TransMLA | -71.875% | 0.6312 | 0.6813 | 0.7490 | 0.4633 | 0.4420 | 0.7911 | 0.6606 |\n| Mistral-7b-v0.1 | — | 0.6483 | 0.6671 | 0.7948 | 0.4567 | 0.4560 | 0.8286 | 0.6867 |\n| TransMLA | -50% | 0.6448 | 0.6622 | 0.7900 | 0.4553 | 0.4560 | 0.8270 | 0.6780 |\n\nAccording to Table 2, aligning the converted models using only the distilled data from the original models allows them to closely match the performance of the original models. This confirms that the observed performance improvements are not due to additional learning but reflect genuine restoration. We plan to distill more data and complete experiments on additional models and compression rates in the future.\n\n[1] Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs (ACL 2025)\n\n[2] LLM-QAT: Data-Free Quantization Aware Training for Large Language Models (ACL Findings 2024)"}, {"review_id": "LfC532fTgg", "reviewer": "Reviewer_i2PF", "summary": "The TransMLA framework aims to seamlessly convert existing GQA (Group-Query Attention) pre-trained models into MLA (Multi-head Latent Attention) models. TransMLA uses RoRoPE (Rotary Position Embedding Rotation) and FreqFold: to concentrate the rotational position encoding (RoPE) information into a small number of dimensions (especially the first attention head), thereby removing RoPE from the remaining dimensions to support DeepSeek's unique Absorb operation. It is observed that the norm of the key after removing RoPE is significantly larger than the value. The authors balance the norms of the two by scaling to improve the compression performance of the joint low-rank decomposition. TransMLA achieves significant KV cache compression and inference acceleration.", "questions": "Is MLA really better than GQA? The paper only proves the phenomenon of equal KV Cache, but does it have an advantage under equal parameter values ​​or headdim? Or is GQA good enough for us? There are many ways to solve KVCache, not necessarily MLA. \n\nIf MLA is indeed theoretically superior, and DeepSeek actively open-sources its optimization solution, then why does the GQA model still dominate and require such a complex conversion?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\nIt provides an immediate and feasible path for existing models to utilize DeepSeek's optimization ecosystem, which has strong practical value.\n\nImpressive results were achieved in KV cache compression and inference acceleration. For LLaMA-2-7B, it achieved 93% KV cache compression and 10.6x inference acceleration. \n\nWeakness:\nThe paper does not make significant improvements to MLA, and the innovations proposed are not the source of major performance improvements (such as kvache compression, speed improvement, parameter pruning, etc.) For example, the paper admits that its \"Balance KV technology itself is relatively trivial\", but includes it as an \"innovation\" point as the main contribution.\n\nThe method is overly dependent on a specific ecosystem. All the \"innovations\" in the paper - RoRoPE, FreqFold and even support for Absorb operations - are to achieve compatibility with DeepSeek's unique architecture, but do not discuss the performance differences, advantages and disadvantages of MHA and MLA, and whether there are some improvement solutions. This high degree of coupling has led to doubts about the universality of its methodology, making it impossible to become a general attention mechanism conversion framework with broad scientific value.", "was_revised": true, "final_justification": "Thank you for your detailed experiments and response. Your experiments and theoretical analysis have addressed many of my concerns, but regarding MLA’s effectiveness, I believe it must be evaluated from two perspectives. First, MLA only shows an advantage over GQA when training memory is limited. If memory is not constrained—under the same FLOPs or in an overtrained setting—MLA may not offer clear improvement over GQA due to low-rank issues. This is why I emphasize comparing MLA with the same parameter count rather than cache methods, since compute continues to advance and caching may not be the primary scaling direction. Furthermore, once converted, will subsequent training of the model remain unaffected? Is this risk worth taking? In fact, even after MLA’s introduction, most models still train from scratch using GQA. On the other hand, I support the innovation of converting GQA to MLA, but as a KVCache optimization it is too costly and complex, and many alternative methods exist. Moreover, the authors’ design and contributions are only meaningful in the conversion to MLA. However, it remains unclear whether the MLA conversion method will become the mainstream approach for optimizing cache, and therefore its contribution seems limited. Therefore, I have increased my scores for Quality, Clarity, Significance, and Originality.", "rebuttal": "**Q1. About the innovations proposed in this paper.**\n\n**A1.** This work makes solid theoretical and practical contributions:\n\n**First**, we theoretically prove that MLA has a stronger expressive capacity than GQA under the same KV cache size (see Appendix A), addressing a key gap left by DeepSeek, which only empirically demonstrated MLA’s capabilities. This theoretical framework enables comparative analysis of the expressive power of major attention mechanisms, including MLA, MHA, GQA, and MQA and provides strong motivation for converting GQA to MLA.\n\n**Second**, we are **the first** to prove sufficient conditions for performing invariant query/key transformations across RoPE (see Appendix B), namely, preserving the inner product between RoPE-applied query and key before and after transformation. This breakthrough is crucial for enabling MLA to benefit from inference acceleration **while preserving position information**, which was a significant challenge for prior works. For example, Palu [1] performs low-rank compression before applying RoPE, which not support absorption, requires recomputing the KV representations from the KV cache before computing attention. As a result, the acceleration benefit is limited. MHA2MLA [2] and X-EcoMLA [3] directly truncate RoPE in a non-equivalent manner, leading to significant loss of positional information and degraded performance.\n\nOur theoretical result led to the development of **RoRoPE**, a constrained low-rank compression method, and **FreqFold**, a technique that enhances the expressive power of RoRoPE. Together, these innovations reduce the loss from decoupling RoPE to a minimal level, making GQA-to-MLA conversion with partial RoPE feasible. \n\nLastly, **Balance-KV** is a simple yet highly effective method for improving joint KV compression. To the best of our knowledge, this technique has not been adopted in prior related works [1,2,3], and we list it as a contribution. Please refer to **Reviewer E7fv A1**, where we provide the theoretical justification behind this method.\n\n**Q2.1. Discuss the performance differences, advantages and disadvantages of MHA and MLA**\n\n**A2.1.** In the *Introduction*, *Related Work*, and the section on *Expressive Power Analysis*, we have outlined the strengths and weaknesses of MHA, GQA, and MLA: 1) **MHA**: High expressive power, but slow inference and maximum KV cache. 2) **GQA**: Lower expressive power, but moderate inference speed. 3) **MLA**: A balance between expressive power and inference efficiency, with higher expressiveness than GQA under the same KV cache. MLA has already garnered significant attention in both model architecture design and AI infrastructure communities. Numerous improved variants such as GLA [4], MTLA [5], and GTA [6] have been proposed, and dedicated frameworks and hardware optimizations for MLA inference are actively under development. Given this momentum, our unified framework for analyzing the expressive capacity of multiple attention mechanisms becomes especially valuable and broadly applicable.\n\n**Q2.2.** The innovations are overly dependent on Absorb operations\n\n**A2.2.** The brilliance of MLA’s design, beyond its use of low-rank compression—a concept already explored in prior works—lies in the **Absorb operation**. This operation is central to MLA’s efficiency: during the **training** phase (which is compute-intensive), it allows attention computations to be performed using compact 192-dimensional representations per head, achieving high computational efficiency. During **inference** (which is memory-bound), it employs a shared 576-dimensional KV cache across heads by absorbing the key projections into the query projection matrices, substantially reducing memory access overhead on both high-bandwidth memory (HBM) and on-chip SRAM.\n\nAs demonstrated in Table 3 (rows 2–3), disabling the absorb operation causes MLA’s inference speed at a 16k context length to degrade by **16.7×**, clearly showing its critical role. Consequently, this paper places strong emphasis on supporting the **Absorb operation**, as it is essential for fully realizing MLA’s potential in practical deployment.\n\nThe importance of the absorb operation is further illustrated in Palu [1]. In their Figure 5, they show that removing RoPE allows the absorb operation to achieve a **6.17× speedup** at 64K context. However, when RoPE is retained, absorb must be disabled, and the speedup drops to **2.91×**. Unfortunately, Palu does not support RoPE decoupling and therefore cannot apply the absorb operation to popular models such as LLaMA, Qwen, and Mistral.\n\n**Table 3.** Comparison of the impact of num_kv_heads, head_dim, and parameters on inference speed for GQA and MLA. To eliminate the influence of MLP layers on the results, we conducted the speed test using models composed of 32 attention layers only.\n\n| method | num_kv_heads | head_dim | kv_cache per layer | parameters(M) | Throughput(tokens/s) |\n| --- | --- | --- | --- | --- | --- |\n| MLA (with Absorb)  | 1*1 | 576 | 576 | 48.5 | 4151.17 |\n| MLA (without Absorb) | 32*2 | 192 | 12288 | 48.5 | 248.85 |\n| GQA | 4*2 | 128 | 1024 | 37.7 | 2714.80 |\n|  | 8*2 | 128 | 2048 | 41.9 | 1436.92 |\n|  | 16*2 | 64 | 2048 | 41.9 | 1118.83 |\n\n**Q3.1. Why compare capability under equal KV Cache rather than equal parameter count or head dimension?**\n\n**A3.1.** The inference speed of LLMs is primarily affected by two factors: computational intensity and memory access [7]. As the KV cache length increases during inference, modern hardware primarily hits a **memory bottleneck** rather than compute bottleneck. For example, in Table 3, GQA models with the same head dimension where one doubles its KV cache size show drastic change in inference speed (also about **2×**). In contrast, using the same KV cache size but doubling the head dimensions leads to marginal speed changes. For MLA models, they have the same parameter count before and after absorption, but the difference in KV cache size results in a **16.7× gap** in inference speed. Although MLA has more parameters than GQA, its **reduced KV cache size** leads to **faster inference**. Therefore, comparing models under **equal KV cache size**—the factor most directly affecting inference speed—is both practical and meaningful for analyzing expressive capacity.\n\n**Q3.2. There are many ways to handle the KV cache—why focus on MLA?**\n\n**A3.2.** As discussed in *Related Work*, existing KV cache optimization techniques are orthogonal to ours: 1) **KV cache pruning** reduces the number of tokens stored; 2) **KV cache quantization** reduces the bit-width per value. In contrast, TransMLA focuses on **per-token dimensionality reduction**, and thus complements other techniques rather than replacing them.\n\n**Q3.3. Is MLA really better than GQA?**\n\n**A3.3.** Theoretically, we prove that MLA always has greater expressive power than GQA under the same KV cache size. Empirically, several SOTA open-source models using **576-dimensional** MLA—such as **DeepSeek-V3** and **Kimi-K2**—achieve performance comparable to or even surpassing that of GQA models using **2048-dimensional** KV cache, such as LLaMA-4-Maverick-17B-128E, Qwen3-Coder-480B-A35B, Mixtral-8x22B, and Gemma-3-27B. This provides strong evidence that **MLA can effectively replace GQA** as a core component for high-performance, inference-efficient foundation models.\n\n**Q4. Why do GQA models still dominate, and is such a complex conversion to MLA really necessary?**\n\n**A4.** Pretraining a foundation model is a massive system engineering task that requires **months of continuous experimentation and iteration**. Leading model developers have already invested heavily in GQA-based architectures and accumulated significant optimization experience for both training and inference. Switching to a new architecture would demand substantial system reengineering and delays due to tasks like value alignment before public release, not to mention that DeepSeek R1’s release was only 6 months ago. Despite these challenges, we are already seeing a shift—**Kimi-K2**, a state-of-the-art agentic model, has adopted **MLA**. Our theoretical work provides a compelling reason for major model developers to consider this transition. Moreover, the **GQA-to-MLA conversion method** proposed in this paper allows them to **recycle existing pretrained models**, avoiding the need to train from scratch and thereby **reducing the carbon footprint**.\n\n[1] Palu: Compressing KV-Cache with Low-Rank Projection (ICLR 2025)\n\n[2] Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs (ACL 2025)\n\n[3] X-EcoMLA: Upcycling Pre-Trained Attention into MLA for Efficient and Extreme KV Compression (COLM 2025)\n\n[4] Hardware-Efficient Attention for Fast Decoding\n\n[5] Multi-head Temporal Latent Attention\n\n[6] Grouped-head latenT Attention\n\n[7] Roofline: an insightful visual performance model for multicore architectures."}, {"review_id": "ttZlInkTNB", "reviewer": "Reviewer_KGKd", "summary": "In this paper, the authors propose TransMLA, a framework that converts GQA/MHA modules from existing LLM models (such as Llama) into MLA modules that fit DeepSeek's frameworks. To adopt the RoPE from GQA/MHA to the decoupled RoPE in DeepSeek's MLA, the authors propose RoRoPE and FreqFold to concentrate the RoPE information to the first key/query head. Besides, a balanced key-value method is proposed for model compression.  The proposed method is shown to be a better conversion than existing works (such as MHA2MLA) and could greatly save the memory cost of the KV cache.", "questions": "1. Overall I think Figure 1 is relatively clear to me but they are not well-explained in Section 4. \n - For example, it is not known how $q_{rope}$, $q_{nope}$, $k_{rope}$, $k_{nope}$ are defined in TransMLA.\n - It is not known how the RoPE is removed for $q_{nope}$ and  $k_{nope}$ after applying the PCA rotation matrix. \n - Also, it is not know how the matrices in MLA are exactly initialized (e.g., $W_{dq}$, $W_{qr}$, $W_{uq}$, $W_{dkv}$, $W_{uk}$, $W_{uv}$)\n\n2. The sum sign seems to be missing in Equation (29)\n\n3. It is a little bit not clear to me how this claim \"MLA consistently exhibits greater expressive capacity than GQA given the same KV cache size\" helps in this paper. When setting the same KV cache size, MLA seems clearly stronger than GQA as there are more parameters. It would be great to know if there are any study/theoriem comparing the expressive capacity of GQA and MLA given similar amount of parameters. \n\n4. There seems to be another related work named X-EcoMLA[1]. How is the proposed TransMLA compared to this work?\n\n5. I am wondering how the authors benefit from the DeepSeek optimizations for the inference speed measurement. Is the FlashMLA adopted in this experiment? It would be beneficial to add more details for the inference optimization.\n\n[1] Li, Guihong, et al. \"X-ecomla: Upcycling pre-trained attention into mla for efficient and extreme kv compression.\" arXiv preprint arXiv:2503.11132 (2025).", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The proposed method achieves much better performance than MHA2MLA, especially when there is no additional training. This means the proposed TransMLA is a better approximation to existing MHA/GQA and can be used as a starting point for future research \n\nWeakness:\n- Although the proposed TransMLA achieves good performance, it requires a calibration dataset for PCA and the effect of the calibration dataset is not studied. For example, what is the effect of the size of the calibration dataset? Also, what if we choose another dataset as the calibration dataset?\n- The description of the proposed method is not very clear, and there are some typos (please refer to Questions)\n- There is a lack of ablation studies. For example, what would be the commonsense reasoning performance if we don't enable all the components (e.g., RoRoPE without FreqFold)\n- The authors missed some other related works in the literature (please refer to Questions)\n- The authors tested TransMLA on SmolLM-1.7B and Llama2-7B, but there is no experiment on Qwen/Mistral (which is claimed in the abstract)\n- The proposed method is not compared to other KV cache compression methods such as KVquant and Palu \n- The proposed method is only tested on common-sense reasoning tasks. It would be good to test on some long-context tasks (such as LongBench) since MLA's advantage is mainly on long context lengths", "was_revised": true, "final_justification": "I appreciate the author's responses, which have solved most of my concerns. I have raised my score.", "rebuttal": "We sincerely thank you for your detailed questions and constructive feedback. Below, we address each of your concerns and will incorporate the necessary updates in the camera-ready version.\n\n**Q1. Calibration dataset for PCA: effect of size and dataset choice**\n\n**A1.** As per your suggestion, we conducted an ablation study on the calibration dataset. As shown in **Table 4**, the number of calibration samples does affect the conversion quality, with **64 samples yielding the best performance**. Using different datasets for calibration also impacts results slightly, but all perform better than using randomly generated tokens.\n\n**Table 4.** Ablation study of using different calibration data sources and dataset sizes for converting LLaMA-3-8B to MLA.\n\n|  | n_samples | Decouple RoPE ppl | Compress KV  ppl |\n| --- | --- | --- | --- |\n| wikitext2 | 1 | 8.674 | 155.65 |\n|  | 2 | 8.217 | 38.134 |\n|  | 4 | 8.192 | 22.859 |\n|  | 8 | 8.05 | 18.688 |\n|  | 16 | 8.164 | 17.499 |\n|  | 32 | 8.217 | 15.966 |\n|  | 64 | 8.368 | 14.898 |\n|  | 128 | 8.404 | 15.08 |\n| ptb | 128 | 8.209 | 19.394 |\n| c4 | 128 | 7.341 | 16.018 |\n| alpaca | 128 | 7.389 | 17.398 |\n| random | 128 | 15.693 | 44.956 |\n\n**Q2. Lack of ablation studies (e.g., commonsense** reasoning performance for **RoRoPE without FreqFold)**\n\nA2. We have presented ablations on perplexity (PPL) in **Figure 2 and Figure 3** of our submission for RoRoPE, FreqFold, and BalanceKV. Following your suggestion, we additionally evaluated **commonsense reasoning accuracy**, and the results are shown in **Table 5**:\n\n- Without **RoRoPE**, accuracy drops by ~20% compared to the best configuration.\n- Without **FreqFold** (e.g. FreqFold = 1), accuracy drops by ~6.5%.\n- Without **BalanceKV**, accuracy drops by ~10%.\n\nThese results highlight the importance of each component in maintaining strong downstream performance.\n\n**Table 5.** Ablation study of removing each method and its impact on commonsense reasoning performance of the converted model without training.\n\n|  | Avg. | MMLU | ARC | PIQA | HS | OBQA | WG |\n| - | - | --- | --- | --- | --- | --- | --- |\n| llama-3-8B | 63.92 | 46.13 | 65.81 | 80.79 | 76.26 | 45.40 | 69.14 |\n| TransMLA | **55.12** | **37.47** | **54.31** | **73.78** | **65.72** | **38.80** | **60.62** |\n| - RoRoPE | 35.59 | 25.64 | 28.91 | 55.77 | 28.06 | 24.00 | 51.14 |\n| - FreqFold | 48.63 | 31.19 | 45.24 | 69.37 | 56.18 | 34.60 | 55.17 |\n| - BKV | 44.91 | 29.19 | 37.81 | 65.56 | 49.89 | 32.20 | 54.78 |\n\n**Q3. Experiments on Qwen / Mistral**\n\n**A3.** As mentioned in  **Reviewer qHce A1**, our work follows the experimental setup of MHA2MLA and uses the smollm-corpus for pretraining SmolLM-1.7B and LLaMA-2-7B. To answer your question, we converted Qwen and Mixtral to MLA and conducted distillation experiments to recover the performance. The results are shown in the following table.\n\n| Model | KV cache | avg | arc | hellaswag | mmlu | openbook_qa | piqa | winogrande |\n| - | - | - | - | - | - | - | - | - |\n| LLaMA-2-7b | — | 0.5981 | 0.5931 | 0.7313 | 0.4135 | 0.4160 | 0.7840 | 0.6504 |\n| TransMLA |  -68.75% | 0.5989 | 0.5982 | 0.7243 | 0.4167 | 0.4260 | 0.7818 | 0.6464 |\n| Qwen2.5-7b | — | 0.6338 | 0.6445 | 0.7749 | 0.4726 | 0.4540 | 0.7976 | 0.6590 |\n| TransMLA | -71.875% | 0.6312 | 0.6813 | 0.7490 | 0.4633 | 0.4420 | 0.7911 | 0.6606 |\n| Mistral-7b-v0.1 | — | 0.6483 | 0.6671 | 0.7948 | 0.4567 | 0.4560 | 0.8286 | 0.6867 |\n| TransMLA | -50% | 0.6448 | 0.6622 | 0.7900 | 0.4553 | 0.4560 | 0.8270 | 0.6780 |\n\n\n**Q4. Comparison with KVQuant and Palu**\n\n**A4.** KVQuant is discussed in **Line 34** of our original paper, and Palu is discussed in **Lines 109–114**. As further explained in **Reviewer i2PF’s comments A1 (second paragraph) and A3.2**, TransMLA is orthogonal to KVQuant and can be combined with it. Palu avoids decoupling RoPE by compressing only modules that do not include positional encoding. In their Figure 5, they show that without RoPE, *absorb* enables a **6.17×** speedup comparing to MHA module at a 64K context length. However, when RoPE is present, *absorb* is disabled, and the speedup drops to **2.91×**. This highlights the significance of our contribution, which enables low-rank compression of RoPE—making *absorb* usable and thus enabling more efficient inference.\n\n**Q5. Test on LongBench benchmark**\n\nFollowing your suggestion, we select Qwen-2.5-7B (with max position embeddings of 131072) and convert it into an MLA model. The original model is then used to distill the converted one. We report the performance of both the original and the converted (distilled) models on multiple long-context tasks from the **LongBench** in Table 6.\n\n**Table 6.** Performance on multiple long-text tasks and speed (tested on 16k-length context) of Qwen after converting GQA to MLA and applying further distillation, as described in **Reviewer qHce A2**.\n\n|  | Throughput(tokens/s) | Single-Document QA | Multi-Document QA | Summarization | Few-Shot Learning | Code Completion | Avg. |\n| - | - | - | - | - | - | - | - |\n| Qwen2.5-7B | 1233.04 | 20.29 | 25.73 | 24.46 | 62.11 | 63.10 | 39.14 |\n| TransMLA | 1925.19 | 16.50 | 19.02 | 18.51 | 44.20 | 38.39 | 27.32 |\n| TransMLA-distill | 1925.19 | 18.74 | 28.26 | 22.39 | 61.00 | 59.21 | 37.92 |\n\n**Q6. Figure 1 is not well-explained.**\n\n**A6.** Thank you for pointing this out. We realized that Figure 1 contained too much information, so we have split it into three separate parts: an overall workflow diagram, a frequency-domain transformation illustration for RoRoPE, and a schematic of the FreqFold approximation. Each annotated component in the figures is now explained step-by-step in the Method section. These updates will be included in the camera-ready version.\n\n**Q7. How does proving that MLA has stronger expressive capacity than GQA given the same KV cache size or parameters help support the contributions of this paper?**\n\n**A7:** Please refer to **Reviewer i2PF A1.1 (first paragraph)** and **A3.1**. “MLA consistently exhibits greater expressive capacity than GQA given the same KV cache size” is a **central motivation** of this conversion. \n\n**Q8. MLA seems clearly stronger than GQA as there are more parameters.**\n\n**A8.** More parameters do not necessarily imply stronger expressive capacity across different model architectures. For example, as noted in the original DeepSeek paper, MHA with more parameters demonstrates lower expressive capacity than MLA. One of the key contributions of our paper is a constructive proof showing that, under equal KV cache size, all GQA can be converted into MLA, but the reverse is not true. This provides a formal guarantee that MLA always possesses stronger expressive capacity than GQA.\n\n**Q9. Any other study comparing the expressive capacity of GQA and MLA.** \n\n**A9.** Since the introduction of MLA last year, several studies have proposed hypotheses and conducted experiments to explore the source of MLA’s strong performance. For example, **Jianlin Su** (the author of RoPE) discussed on his blog that increasing the **head dimensions** in MLA may contribute to its improved expressiveness. \n\n**Q10. Comparison with X-EcoMLA**\n\n**A10.** X-EcoMLA removes all positional encodings in the low-rank (nope) part and averages across all heads to extract 64-dimensional RoPE embeddings. This approach lacks theoretical justification for its approximation validity. Both the rope and nope parts of their transformation introduce significant performance drop. As X-EcoMLA is not open-sourced yet, we were unable to reproduce and benchmark their method. Once their code is available, we will provide a direct comparison in the camera-ready version.\n\n**Q11.** I am wondering how the authors benefit from the DeepSeek optimizations for the inference speed measurement. Is the FlashMLA adopted in this experiment? It would be beneficial to add more details for the inference optimization.\n\n**A11.** FlashMLA supports only the Hopper architecture. In our original submission, we did not have access to Hopper GPUs, so we used VLLM with Triton backend for MLA, and FlashAttention backend for MHA. During the rebuttal period, we reran the experiments using Hopper GPUs with FlashMLA backend. As shown in Table 7, compared to the original GQA model on the same GPU, FlashMLA achieves up to an **11.32× speedup at 32K context length**. The detailed information and the code will be publicly available on the camera ready paper.\n\n**Table 7.** Throughput (tokens/s) comparison between LLaMA-2-7B and the converted MLA-based model, evaluated on Hopper architecture.\n\n|  | 1K | 2K | 4K | 8K | 16K | 32K |\n| - | - | - | - | - | - | - |\n| LLaMA-2-7b | 3943.49 | 2122.85 | 1252.37 | 639.32 | 323.58 | 158.37 |\n| TransMLA | 10967.49 | 6461.55 | 5702.00 | 4679.00 | 3436.58 | 1792.79 |\n| Speedup | 2.78x | 3.04x | 4.55x | 7.32x | 10.62x | 11.32x |"}, {"review_id": "v702Zxqm6n", "reviewer": "Reviewer_E7fv", "summary": "This paper presents TransMLA, a framework for converting GQA-based language models (e.g., LLaMA, Qwen) into the MLA (Multi-head Latent Attention) format used by DeepSeek models. TransMLA enables off-the-shelf migration of pre-trained GQA models into the MLA structure, thereby inheriting the benefits of DeepSeek’s inference-time acceleration (e.g., vLLM, SGlang, FP8). The authors propose several innovations to support this conversion: (1) RoRoPE, a PCA-based method to concentrate RoPE positional information into a small number of dimensions, (2) FreqFold, which exploits frequency similarity in RoPE to further reduce redundancy, and (3) Balanced Key-Value PCA, which equalizes key and value norms to enable better joint compression. Extensive experiments on smolLM and LLaMA-2 demonstrate that TransMLA achieves substantial compression (up to 93%) and speedup (up to 10.6×) while preserving model quality with minimal fine-tuning.", "questions": "Can the authors release the converted MLA models (e.g., LLaMA2 7B)? This would greatly benefit the community and validate the plug-and-play nature of TransMLA.\n\nHow sensitive is Balanced KV to the choice of scaling constant? Could this be learned or optimized per-layer?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\nStrong practical relevance: This work addresses a highly relevant engineering challenge to migrate existing GQA models to DeepSeek’s efficient MLA format without retraining from scratch.\n\nWell-motivated and well-engineered: The techniques such as RoRoPE and Balanced KV are practical, intuitive, and empirically effective.\n\nThorough empirical evaluation: The paper includes both performance recovery metrics and inference latency benchmarks across multiple hardware setups, providing a well-rounded view of the benefits.\n\nSubstantial performance gain: At similar compression rates, TransMLA clearly outperforms MHA2MLA both in accuracy and speed, especially under high compression ratios.\n\nCompatibility with ecosystem: Seamless integration with DeepSeek’s toolchain is a strong point, making this work immediately usable for real-world deployment.\n\n**Weaknesses**\n\nBalanced Key-Value PCA is heuristic: While intuitive and empirically validated, the balancing method lack theoretical guarantees or principled formulation (e.g., via optimization objectives).\n\nCompression calibration requires data: The PCA-based compression still relies on a calibration dataset for collecting activations, which may not be trivial in privacy-sensitive or black-box model settings.\n\nClarity in exposition: Some parts, such as RoPE slicing and rotation math (Section 4.2), are quite dense and could benefit from visual aids or step-by-step walkthroughs.", "was_revised": true, "final_justification": "The rebuttal it supplies a theoretical proof for Balanced KV, offers privacy-safe alternatives of distilled data or weight-only SVD, and clarifies RoRoPE/FreqFold with new visual aids and shows a per-layer scaling search that improves perplexity to confirm Balanced KV’s importance. My concerns has been mostly resolved and I maintain my score of 5.", "rebuttal": "Thank you very much for your encouraging review and constructive feedback. Below we provide point-by-point responses to your comments. All suggested improvements will be incorporated into the camera-ready version.\n\n**Q1. While intuitive and empirically validated, the balancing method lacks theoretical guarantees or principled formulation.**\n\n**A1.** Balanced Key-Value PCA (BKV) is indeed a simple yet effective technique. While many prior works (Palu, MHA2MLA, X-EcoMLA) have not adopted this method, we highlight it as a novel contribution. Behind its apparent simplicity, BKV has a **solid theoretical foundation**. We prove that under a joint loss function minimizing the relative approximation errors of both $W_K$ and $W_V$ components:\n$\\mathcal{L} = \\|| \\frac{W_K - W_K^{top-r}}{\\mathrm{RMS}(W_K)} \\||^2 +\\|| \\frac{W_V - W_V^{top-r}}{\\mathrm{RMS}(W_V)} \\||^2$, scaling $W_K$ and $W_V$ to the same magnitude: $W_{\\text{scaled}}^{KV} =\\(W_K / \\mathrm{RMS}(W_K), W_V / \\mathrm{RMS}(W_V)\\)$ achieves the optimal low-rank approximation under this specific loss function. In our proof, we use $a$ and $b$ as scaling factors. We first show that the loss function depends only on the ratio $a:b$, not on the absolute values of $a$ and $b$. By setting $b=1$, the optimization simplifies to finding the optimal value for $a$. Next, we compute the derivative of the loss function with respect to $a$, and show that: $\\frac{d\\mathcal{L}(a, 1)}{da} = 0$ when $a = \\mathrm{RMS}(W_K) / \\mathrm{RMS}(W_V)$, through a series of substitutions. The main substitutions used are:\n$\\||W\\||^2 = \\mathrm{tr}(W^T W) = \\sum_{i=1}^{\\mathrm{rank}(W)} \\sigma_i^2, \\quad \\text{and} \\quad \\frac{d\\sigma_i}{da} = u_i^T \\frac{dW}{da} v_i$,\nwhere $u_i$, $\\sigma_i$, and $v_i$ are the left singular vectors, singular values, and right singular vectors of matrix $W$, respectively.\nThis proof provides a theoretical foundation and a more principled mathematical formulation for our Key-Value balancing strategy. The detailed proof and mathematical derivation will put in the appendix of the camera-ready version. \n\n**Q2. On the need for a calibration dataset and concerns in privacy-sensitive or black-box settings:**\n\n**A2:** Regarding privacy-sensitive datasets, please see our response to **Reviewer qHce A2**: we address this by using distilled data generated from the original model, which avoids direct access to private corpora.\n\nAs for black-box models, TransMLA is primarily intended to accelerate inference, which generally requires access to model internals. However, we will fully open-source TransMLA, including conversion, training, and inference code, as well as converted checkpoints. This allows model providers to apply the method internally without data sharing. Alternatively, they may use our converted foundation models and fine-tune them on in-house data.\n\nIn addition, we also provide alternative implementations of RoRoPE, FreqFold, and BKV that do not require calibration data. Instead of performing PCA on activations, these versions apply SVD directly to the weight matrices. This option will be included in our open-sourced codebase for users who prefer not to rely on calibration datasets.\n\n**Q3. On clarity in exposition, especially for RoPE slicing and rotation math in Section 4.2:**\n\n**A3.** Thank you for this valuable suggestion. Although we made a conscious effort to ensure readability during writing, we recognize that several theoretical components in TransMLA—especially **RoRoPE**—remain abstract and challenging to follow. We have since redesigned the **conversion pipeline diagram** to include each component’s purpose. Additionally, we created intuitive **frequency-domain visualizations** to explain both RoRoPE and FreqFold. These updates significantly improve clarity and allow readers to grasp the motivation even without diving into all the formulas. These diagrams will be included in the camera-ready version.\n\n**Q4. On releasing converted MLA models such as LLaMA2-7B:**\n\n**A4.** Absolutely. We will release all code for conversion, training, and inference, as well as checkpoints for converted models. In the future, we plan to support more models and more tasks, providing the community with faster and stronger **pretrained models.**\n\n**Q5. On the sensitivity of Balanced KV to the scaling constant and potential for optimization:**\n\n**A5.** To investigate this, we conducted an ablation study as shown in Table 8. On LLaMA-3-8B, completely removing Balanced KV results in a 3.8-point increase in perplexity, highlighting its critical role. We also observed that model performance is sensitive to the scaling coefficient used in Balanced KV. The default value in our paper already yields strong overall results. However, following your suggestion, we performed a per-layer grid search over the range [0.3, 2.0] with 20 intervals, aiming to minimize the output difference from the original model. This tuning led to a 1.2-point reduction in perplexity compared to the default, suggesting that layer-wise coefficient optimization is both feasible and beneficial for further improving performance.\n\n**Table 8.** Ablation study on the impact of the Balance KV scaling factor on the performance of the converted model. In the main paper, the default factor is computed as the average L2 norm of $K_{\\text{RoPE}}$ divided by the average L2 norm of V.\n\n| BKV | lora ppl (wiki) |\n| --- | --- |\n| None | 28.8341 |\n| (k_rope_mean_norm*4)/v_mean_norm | 18.8790 |\n| (k_rope_mean_norm*2)/v_mean_norm | 15.3617 |\n| k_rope_mean_norm/v_mean_norm | **15.0796** |\n| k_rope_mean_norm/(v_mean_norm*2) | 17.8757 |\n| k_rope_mean_norm/(v_mean_norm*4) | 27.9021 |\n| optimized per-layer | **13.7945** |"}]

Author rebuttal:

Meta-review: This paper proposes TransMLA, a practical framework for converting MHA/GQA models into MLA-compatible ones, enabling significant KV cache compression and inference acceleration. Key innovations like RoRoPE, FreqFold, and Balanced KV PCA are well-motivated and validated. The paper stands out in terms of empirical rigor and practical relevance. While scope limitations (e.g., lack of GQA baselines, limited theoretical grounding) were noted, the authors addressed most concerns in rebuttal with additional experiments and clarifications. This work is an accept.
Decision: Accept (spotlight)