The Road to Transformer Upgrades: 20. What Makes MLA Good? (Part 1)

Transformer升级之路：20、MLA好在哪里?（上） 苏剑林 · 2025-05-04 · Machine Learning optimization language model generative model attention

Observations

MLA was proposed in DeepSeek-V2]. This post assumes the reader is already familiar with MLA, or at least with the content covered in the earlier blog post The Endless Tug-of-War Between Cache and Performance: From MHA, MQA, GQA to MLA], so the details of MLA itself will not be elaborated on much. more

The main characteristics of MLA are as follows:

1. During training, MLA is essentially an MHA with qk_head_dims=(128+64) and v_head_dims=128;
2. During decoding, MLA is essentially a KV-Shared MQA with qk_head_dims=(512+64) and v_head_dims=512;
3. The concatenation of [qc, qr] and [kc, kr] in MLA can be understood as a form of Partial RoPE].

Hypotheses

For MHA and GQA, the commonly used head_dims is 128, whereas for MLA, whether viewed from the training side (128+64) or the inference side (512+64), the value is always larger than 128. Combining this with the experience from Breaking the Bottleneck: Building a Stronger Transformer], we arrive at:

Hypothesis 1: Increasing head_dims is one of the key reasons MLA works well.

In addition, the KV-Shared property allows GQA's head_dims or num_groups to be increased under the same KV Cache size, so we also have:

Hypothesis 2: KV-Shared is one of the key reasons MLA works well.

Finally, prior theory and experiments have suggested that Partial RoPE may have a positive effect on performance (see Transformer Upgrade Road: 18, Principles for Choosing the RoPE Base]), so we have:

Hypothesis 3: Partial RoPE is one of the key reasons MLA works well.

Experiments

We now test each of the above hypotheses one by one through experiments.

Setup

The hyperparameters common to all experiments are as follows:

1. A Dense model similar to LLAMA3;
2. hidden_size=2048, num_layers=12, num_heads=16;
3. The optimizer is Muon], with per-head updates in the Attention part;
4. Training length is 4096, total token count is 16B, total training steps is 16k;
5. All experiments only change the Attention module, so parameter counts are not strictly aligned across configurations.

Part I

MLA's KV Cache size is 512+64, which is roughly equal to GQA2-128 (the first number is num_groups, the second is head_dims), so the baselines for comparison are GQA2-128 and GQA1-256. To test Partial RoPE, we add GQA1-256-PR, where the 256 dims of Q and K are split into 192+64, with RoPE applied only to the 64 part and not to the 192 part.

The results are as follows:

$$\begin{array}{c|ccc} \hline & \text{Params} & \text{Loss} & \text{Cache} \\ \hline \text{MLA} & 894M & 2.721 & 576 \\ \text{GQA2-128} & 842M & 2.75 & 512 \\ \text{GQA1-256} & 943M & 2.72 & 512 \\ \text{GQA1-256-PR} & 943M & 2.711 & 512 \\ \hline \end{array}$$

That is:

$$\text{GQA2-128} < \text{MLA} \lesssim \text{GQA1-256} < \text{GQA1-256-PR}$$

This provides preliminary confirmation of the benefits of increasing head_dims and of Partial RoPE. It appears that the seemingly reluctant design choice in MLA of concatenating RoPE and NoPE parts is very likely a key reason behind its outstanding performance! The original paper's claim that MLA even outperforms MHA is most likely also because the MHA being compared against only had head_dims=128.

Part II

To further verify the effect of increasing head_dims, we additionally ran three experiments: MHA, GQA2-192, and MLA-256. MHA is the conventional MHA with head_dims=128; GQA2-192 simply increases GQA2's head_dims to 192; MLA-256 raises MLA's 128+64 to 192+64. The comparison is as follows:

$$\begin{array}{c|ccc} \hline & \text{Params} & \text{Loss} & \text{Cache} \\ \hline \text{MHA} & 931M & 2.721 & 4096 \\ \text{MLA} & 894M & 2.721 & 576 \\ \text{MLA-256} & 989M & 2.705 & 576 \\ \text{GQA2-128} & 842M & 2.75 & 512 \\ \text{GQA2-192} & 899M & 2.729 & 768 \\ \text{GQA1-256} & 943M & 2.72 & 512 \\ \text{GQA1-256-PR} & 943M & 2.711 & 512 \\ \hline \end{array}$$

We can see that although MHA has more total parameters and a KV Cache 7 times larger than MLA, its loss only barely matches MLA's — a result close to the conclusion in the DeepSeek-V2 paper. Furthermore, GQA2-192 outperforms GQA2-128 but still falls short of GQA1-256; and once MLA's head_dims is raised from (128+64) to (192+64), performance improves further compared to (128+64). These phenomena all indicate that increasing head_dims is far more effective than increasing num_groups.

Part III

Next we test KV-Shared, i.e., having K and V share all or most of their dims. Here the main alternative we consider is GQA with head_dims not exceeding 256, while keeping the total KV Cache size close to that of MLA — so with KV-Shared, we can go up to at most GQA2-256.

Since KV-Shared is not fully compatible with RoPE, following MLA's approach, we split the 256 dims into 192+64, where

1. The 192 part carries no RoPE and is shared between K and V;
2. The 64 part carries RoPE and is used only for K;
3. V is additionally projected into a separate 64 dims, which is concatenated onto the shared 192 dims.

This way, both K and V have head_dims=256, and the total KV Cache size is (192+64+64)*2=640, slightly larger than MLA's 512+64=576. We denote this version "GQA2-(192+64)-S1", where "S1" stands for "Shared-1".

Part IV

Another KV-Shared scheme is as follows:

1. The 192 part carries no RoPE and is shared between K and V;
2. The 64 part carries RoPE and is likewise shared between K and V;
3. When computing Attention, since V carries RoPE, this effectively becomes an absolute positional encoding;
4. To preserve relative positional encoding, the output is split into 192+64, and the 64 part has an additional inverse RoPE applied to it.

In this scheme, K and V are fully shared, and the KV Cache size is (192+64)2=512, slightly smaller than MLA. We call this version "GQA2-(192+64)-S2", where "S2" stands for "Shared-2". The underlying principle is the VO-RoPE method the author recently proposed, see Transformer Upgrade Road: 19, The Second Type of Rotary Position Embedding*].

Part V

In addition, following the same logic, we ran a few more experiments with GQA4 and GQA1. All the experimental results are summarized below:

$$\begin{array}{c|ccc|c} \hline & \text{Params} & \text{Loss} & \text{Cache} & \text{note} \\ \hline \text{MLA} & 894M & 2.721 & 576 & \\ \text{MLA-256} & 989M & 2.705 & 576 & \\ \text{GQA2-(192+64)-S1} & 946M & 2.714 & 640 & \\ \text{GQA2-(192+64)-S2} & 943M & 2.708 & 512 & \text{introduce VO-RoPE} \\ \text{GQA4-(64+64)-S2} & 842M & 2.738 & 512 & \\ \text{GQA4-(128+64)-S2} & 899M & 2.713 & 768 & \text{KV cache max} \\ \text{GQA1-(512+64)-S3} & 1171M & 2.677 & 576 & \text{head_dims max} \\ \hline \end{array}$$

Here, "GQA1-(512+64)-S3" is an MQA implemented according to MLA's inference-time form, a scheme that lies somewhere between S1 and S2, and its main feature is a large head_dims.

Interpretation of the results:

1. KV-Shared GQA inherently comes with Partial RoPE;
2. KV-Shared GQA2-256 can also surpass MLA;
3. The introduction of VO-RoPE seems to help performance (S1 ≲ S2);
4. Under the same KV Cache size, larger head_dims is better;
5. GQA2-(192+64)-S2 slightly outperforms GQA1-256-PR;
6. GQA4-(128+64)-S2 has the largest KV Cache but is not the best performer, once again showing that head_dims matters more.

Regarding KV-Shared, there are two more observations:

1. During training, GQA1-256-PR is clearly ahead of GQA2-(192+64)-S2 early on, but is caught up with — and even slightly surpassed — later in training. This suggests GQA1-256-PR may lack staying power in the later stages;
2. Without KV-Shared, GQA is at most GQA1-256, meaning head_dims tops out at 256. But with KV-Shared, GQA can reach GQA1-512-S — so purely in terms of head_dims, KV-Shared has a higher ceiling.

Part VI

Since parameter counts were not strictly aligned across experiments, readers might wonder "is it really increasing head_dims that matters, or just increasing parameter count?" So here we add a few experiments that control for parameter count.

Three methods for aligning parameter counts are considered here:

1. double-heads: Taking "GQA2-128 vs GQA1-256" as an example, doubling GQA2-128's num_heads makes its parameter count equal to that of GQA1-256;
2. Shrinking the MLP: Reducing the intermediate_size of the MLP (SwiGLU) can also roughly equalize the parameter counts of GQA1-256 and GQA2-128;
3. Q&O LoRA: Since most of GQA's parameters come from the Query and Output projection matrices, applying LoRA to these two matrices can also reduce GQA1-256's parameter count.

The experimental results are as follows:

$$\begin{array}{c|ccc|ccc} \hline & \text{Params} & \text{Loss} & \text{Cache} & \text{num_heads} & \text{intermediate_size} & \text{qo_lora} \\ \hline \text{MLA} & 894M & 2.721 & 576 & 16 & 5456 & \text{No}\\ \hline \text{GQA2-128} & 842M & 2.75 & 512 & 16 & 5456 & \text{No}\\ \text{GQA1-256} & 943M & 2.72 & 512 & 16 & 5456 & \text{No}\\ \hline \text{GQA2-128} & 943M & 2.723 & 512 & \color{red}{32} & 5456 & \text{No} \\ \text{GQA1-256} & 843M & 2.747 & 512 & 16 & \color{red}{4096} & \text{No} \\ \text{GQA1-256} & 842M & 2.726 & 512 & 16 & 5456 & \color{red}{\text{Yes}} \\ \hline \text{GQA4-(64+64)-S2} & 842M & 2.738 & 512 & 16 & 5456 & \text{No} \\ \text{GQA2-(192+64)-S2} & 943M & 2.708 & 512 & 16 & 5456 & \text{No} \\ \hline \text{GQA4-(64+64)-S2} & 943M & 2.711 & 512 & \color{red}{32} & 5456 & \text{No} \\ \text{GQA2-(192+64)-S2} & 843M & 2.733 & 512 & 16 & \color{red}{4096} & \text{No} \\ \text{GQA2-(192+64)-S2} & 842M & 2.708 & 512 & 16 & 5456 & \color{red}{\text{Yes}} \\ \hline \end{array}$$

The results fall into three main groups:

1. Compared to doubling head_dims, doubling the number of heads consistently results in a loss that is about 0.003 worse;
2. Compared to halving head_dims, shrinking the MLP consistently results in a loss that is about 0.004 better;
3. Q&O LoRA incurs the smallest performance loss — it allows doubling head_dims without increasing parameter count, and still achieves a noticeably lower loss.

Conclusion: from the perspective of increasing parameter count, increasing head_dims is likely the direction with the largest gains, and combined with Q&O LoRA, this can be achieved with almost no increase in parameter count, while still yielding substantial benefit.

Summary

The preliminary conclusions are:

1. Increasing head_dims yields the largest benefit;
2. Partial RoPE also helps the loss to some extent;
3. KV-Shared likely also plays some role.

It seems that our previous efforts to find MLA alternatives while stuck at head_dims=128 were disadvantaged from the very start — no wonder they could never quite match MLA. To match MLA, head_dims should probably start at 192, combined with Partial RoPE. As for KV-Shared, it may also help, but this likely needs verification at larger scale.

Significance

The real significance of all this depends on how strongly we're determined to replace MLA.

Suppose GQA2-(192+64)-S2 could serve as a substitute for MLA — but MLA can also be scaled up to 256, and currently GQA2-(192+64)-S2 still falls short of MLA-256. In that case, the only two benefits of replacing MLA would be:

1. A simpler structure, making it easier to add QK-Norm;
2. During decoding, head_dims changes from 512+64 to 256, and num_groups becomes 2, enabling tensor parallelism (TP).