The Road to Transformer Upgrades: 21. What Makes MLA Good? (Part 2)
In the article Transformer升级之路：20、MLA好在哪里?（上）, we ran ablation studies on the differences between MLA and the more common MHA, GQA, and MQA, covering changes such as "increasing head_dims," "Partial RoPE," and "KV sharing." The preliminary experimental results suggested that all three of these changes are likely contributors to MLA's excellent performance.
In this article, we'll take a more theoretical perspective to try to understand what makes MLA successful.
Partial Rotation
Let's put the final conclusion up front:
Under equal training and inference cost, MLA may well be the best-performing Full Attention variant.
Clearly, this statement places MLA in a very high position. It's a conclusion reached under fairly idealized and simplified assumptions, based on the experimental results from the previous article as well as the theoretical analysis that follows in this one. Since real-world training and inference involve many complicating factors, this conclusion will likely deviate somewhat in practice — but at the very least, we can say that MLA is heading in the right direction.
There's a major precondition for MLA's strong performance: that Partial RoPE — applying RoPE only partially — performs no worse, and possibly even better, than the full version of RoPE. "Partial RoPE" here can mean two things: first, when adding RoPE to attention's $\boldsymbol{Q}$ and $\boldsymbol{K}$, we can apply it to only a small fraction of the dimensions, leaving the rest unchanged; second, we can alternate RoPE layers with NoPE layers across depth, with NoPE layers potentially making up the majority.
In plain terms, RoPE can be applied "just a little," but it can't be dropped entirely — going fully without it hurts performance. If a theoretical explanation is needed, the author tends to agree with the account given in Transformer升级之路：18、RoPE的底数选择原则, which roughly argues that Partial RoPE lets retrieval results balance position and semantics more effectively. There are also newer works like FoX and SBA that show some promise, but as far as MLA is concerned, these variants are essentially equivalent to NoPE, so they don't change the conclusion.
The conclusion that "Partial RoPE works just as well" lets us place most of attention's computational complexity onto the NoPE part, which creates more room to maneuver — and this is exactly what MLA benefits from.
Key-Value Sharing
The evolution of Full Attention has roughly gone from MHA, MQA, GQA, and then to MLA. Although MQA can be viewed as a special case of GQA, chronologically GQA did indeed come later. After MLA, two more variants emerged: MFA and TPA. Fundamentally, all of these variants are attempts to squeeze the KV Cache as much as possible to speed up generation, while trying to preserve performance as much as possible.
Simply put, the complexity of an attention model can be split into training, prefill, and decoding. Since training and prefill are similar, this essentially comes down to prefill versus decoding. Prefill refers to the stage where the model processes the input and produces the first token; we'll discuss this in the next section. Decoding refers to the token-by-token generation stage, which can be accelerated via the KV Cache mechanism — but this also means that KV Cache size is almost the sole bottleneck for decoding speed.
So compressing the KV Cache is equivalent to speeding up decoding. Now let me ask a question: In an NoPE setting, given a fixed KV Cache size, what is the best-performing attention variant? If we ignore differences in parameter count, and restrict the discussion to single-layer MHA/GQA/MQA (we'll come back to TPA and MFA later), the answer is:
An MQA with head_dims equal to the KV Cache size, and K and V shared.
Surprising, isn't it? Actually it's not hard to understand once you think it through. Since MHA and MQA can both be viewed as special cases of GQA, we only need to analyze GQA. As shown in 《缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA》, GQA can be re-expressed as a model where K and V are concatenated:
\begin{equation}\underbrace{\left[\boldsymbol{k}_i^{(1)},\cdots,\boldsymbol{k}_i^{(g)},\boldsymbol{v}_i^{(1)},\cdots,\boldsymbol{v}_i^{(g)}\right]}_{\boldsymbol{c}_i\in\mathbb{R}^{g(d_k+d_v)}} = \boldsymbol{x}_i \underbrace{\left[\boldsymbol{W}_k^{(1)},\cdots,\boldsymbol{W}_k^{(g)},\boldsymbol{W}_v^{(1)},\cdots,\boldsymbol{W}_v^{(g)}\right]}_{\boldsymbol{W}_c\in\mathbb{R}^{d\times g(d_k+d_v)}}\end{equation}
Here $g(d_k+d_v)$ is exactly the total KV Cache size per token. Then, when computing attention, we can absorb the transformation from $\boldsymbol{c}$ to $\boldsymbol{k},\boldsymbol{v}$ into $\boldsymbol{W}_q$ and $\boldsymbol{W}_o$ respectively, which gives us an MQA whose K and V are both $\boldsymbol{c}$-dimensional. So "an MQA with head_dims equal to the KV Cache size, with K and V shared" is actually a "superset" of MHA/GQA/MQA for a given KV Cache size, which naturally makes it the theoretically optimal choice.
Dual Projection
Putting this together: if we want the best possible performance under a fixed decoding speed, we should train an MQA with a specified head_dims and shared KV — for instance, if we cap the KV Cache at 512, then an MQA with head_dims=512 and shared KV is the optimal choice. As it happens, MLA during the decoding phase is precisely a KV-shared MQA (in its NoPE part), which is one manifestation of it being on the right track.
However, while raising head_dims to 512 is fine for decoding, it's hard to accept for training and prefill, since their bottleneck is compute, and the main factors affecting compute speed are num_heads and head_dims. To preserve performance, there's not much room to adjust num_heads, so head_dims essentially becomes the sole indicator of compute cost — raising head_dims to 512 means quadrupling the compute (compared to head_dims=128).
Now let me ask another question: Still under NoPE, given fixed num_heads and head_dims, what is the best-performing attention variant? I think everyone can agree on the answer here: MHA, because it imposes the fewest constraints. So, purely from the standpoint of training and prefill cost, what we want is to train an MHA with head_dims=128.
How do we reconcile these two conflicting goals for prefill and decoding? This is where MLA's "master stroke" comes in: it obtains K and V through a two-step projection — first projecting the input into a single 512-dimensional vector, then projecting that vector into multiple 128-dimensional vectors. By exploiting the inherent identity-transform property of "attention + NoPE," this lets the model switch freely between MHA-128 and MQA-512.
$$\require{cancel}\begin{array}{c|c} \text{training/Prefill} & \text{Decoding} \\ \\ \begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\right] \\[10pt] \boldsymbol{o}_t^{(s)} = \frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)\boldsymbol{v}_i^{(s)}}{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{(s)}{}^{\top}\right)} \\[15pt] \boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_q^{(s)}\in\mathbb{R}^{d\times d_k}\\ \boldsymbol{k}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_k^{(s)}\in\mathbb{R}^{d_k},\quad \boldsymbol{W}_k^{(s)}\in\mathbb{R}^{d_c\times d_k} \\ \boldsymbol{v}_i^{(s)} = \boldsymbol{c}_i\boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_v},\quad \boldsymbol{W}_v^{(s)}\in\mathbb{R}^{d_c\times d_v} \\[10pt] \boldsymbol{c}_i = \boldsymbol{x}_i \boldsymbol{W}_c\in\mathbb{R}^{d_c},\quad \boldsymbol{W}_c\in\mathbb{R}^{d\times d_c} \end{gathered} & \begin{gathered} \boldsymbol{o}_t = \left[\boldsymbol{o}_t^{(1)}\boldsymbol{W}_v^{(1)}, \boldsymbol{o}_t^{(2)}\boldsymbol{W}_v^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\boldsymbol{W}_v^{(h)}\right] \\[10pt] \boldsymbol{o}_t^{(s)} = \frac{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}}{}^{\top}\right)\boldsymbol{v}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}} }{\sum_{i\leq t}\exp\left(\boldsymbol{q}_t^{(s)} \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}}{}^{\top}\right)} \\[15pt] \boldsymbol{q}_i^{(s)} = \boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\boldsymbol{W}_k^{(s)}{}^{\top}\in\mathbb{R}^{d_c}\\ \boldsymbol{k}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}} = \boldsymbol{v}_i^{\color{#ccc}{\smash{\bcancel{(s)}}}} = \boldsymbol{c}_i= \boldsymbol{x}_i \boldsymbol{W}_c\in\mathbb{R}^{d_c} \end{gathered} \end{array}$$
In Summary
Let's summarize the reasoning above:
1. Main premise: Partial RoPE performs no worse than, and possibly better than, full RoPE, which lets us focus our main efforts on the NoPE part;
2. Decoding's main bottleneck is the KV Cache; the theoretically optimal model is an MQA with head_dims = KV Cache size and shared KV;
3. Training and prefill's main bottleneck is head_dims; the theoretically optimal model is an MHA with the desired head_dims;
4. Under NoPE, attention has an identity-transform property that lets us use LoRA to accommodate both ideal directions as much as possible — which is exactly what MLA does.
What remains is to append a shared low-dimensional RoPE to K, supplying positional information to MLA at minimal cost — and this achieves two goals at once: appending RoPE happens to align with "Partial RoPE," while also increasing head_dims, which matches the conclusions of the previous article. In other words, whether intentionally or not, MLA ends up using Partial RoPE and increased head_dims, and this is the main reason it can still rival MHA even under extreme compression.
From MQA's perspective, MLA adds a rank-128 LoRA to Q; from MHA's perspective, MLA adds a rank-512 LoRA to K and V. You could say MLA is an extreme "magic show" combining NoPE with LoRA, and MHA with MQA, successfully achieving a "meeting in the middle" between prefill and decoding.
Of course, the reasoning above certainly involves some oversimplifications. For example, real-world training and inference involve many more detailed factors, and reducing everything to head_dims and KV Cache isn't entirely accurate — for instance, MQA can't use tensor parallelism (TP) during decoding, which may introduce new efficiency issues. Also, our analysis didn't pay particular attention to matching parameter counts — for instance, at head_dims=128 we could also consider increasing the projection complexity of Q, K, V to improve performance, rather than necessarily increasing head_dims; and so on.
In short, this two-part article aims to provide some experiments and reasoning to argue for MLA's optimality within a certain scope. Of course, MLA was first proposed by DeepSeek, and third parties using MLA will always seem to be "copying" DeepSeek to some extent. But until a better variant appears, or until serious flaws are discovered, MLA remains a quite competitive choice. Avoiding MLA purely to show that you're not "following" DeepSeek would be a rather unwise decision.
For example, hybrid models combining Linear Attention and Softmax Attention currently show great competitiveness. But if we mix Linear Attention with the GQA8-128 used by LLaMA at a 3:1 ratio, the resulting KV Cache is roughly reduced to 1/4 of GQA8-128's — yet MLA itself already reduces the KV Cache to 1/4 of GQA8-128's.
Supplementary Discussion
So far we've been discussing MHA, GQA, MQA, and MLA. In this section, let's briefly talk about two attention variants that get less attention: TPA and MFA.
TPA stands for Tensor Product Attention — a name that sounds rather impressive, but in reality it's an intermediate product between GQA and MLA. Taking a target KV Cache of 512 as an example: TPA first projects to a 512-dimensional vector, then reshapes it into (4, 128), then splits it into two (2, 128) tensors representing the K Cache and V Cache respectively. Up to this point, TPA behaves identically to GQA2-128.
Next, TPA borrows an idea from MLA and re-projects the (2, 128) K/V into multiple heads. But instead of projecting from the whole vector as MLA does, it projects along the dimension of size "2" — in other words, it forms head_dims different linear combinations of the two 128-dimensional vectors. Clearly, this gives TPA a lower ceiling than MLA, which projects directly from the full 512-dimensional vector. To mitigate this, TPA additionally introduces data-dependent combination coefficients to enhance the expressiveness of K and V. Even so, the author still believes its ceiling is lower than MLA's.
Why is TPA designed this way? Largely to remain compatible with RoPE, which is arguably its biggest "advantage" over MLA. However, this "advantage" deserves scare quotes, because given that Partial RoPE performs no worse — and possibly even better — than full RoPE, being compatible with RoPE feels a bit like a moot point. Also, this design choice blocks TPA from increasing head_dims: for example, if you wanted to raise head_dims to 256, the K Cache and V Cache would just be shape (1, 256), leaving no freedom for linear combination within a single vector.
Now let's look at MFA, whose full name is "Multi-matrix Factorization Attention" — another somewhat imposing-sounding name. In reality, it's just an MQA with head_dims=256 plus a Q-LoRA. Does this configuration look familiar? It's exactly in line with the conclusion from the previous article — increasing head_dims to 256 to improve MQA's performance, while keeping the KV Cache close to MLA's, and using Q-LoRA to control the parameter count.
So it's no surprise to the author that MFA can "compete" with MLA — we already experimented with a roughly similar approach in the previous article. Furthermore, in the previous article we proposed two other directions for improving MQA's performance: one is Partial RoPE, which we've mentioned repeatedly in this article; the other is achieving full KV sharing via QKVO-RoPE, turning MQA into GQA2-256. Stacking these two improvements on top, MFA should be able to gain a bit more.
Article Summary
Building on the experimental results from the previous article, this article has offered a more theoretical line of reasoning to argue for MLA's optimality within a certain scope. Overall, under the Partial RoPE setting, MLA appears to be an extremely hard-to-beat attention variant.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.