MHA, MQA, CQA到MLA

MHA

Multi-head Attention，多头注意力机制，提出自Attention is all you need

在目前主流LLM采用的自回归中，token by token生成时，新预测出的第t+1个token，并不会影响之前已经计算好的k与v，因此为了避免重复计算，可将这部分结果缓存下来，这就是KVCache，后续所有的MQA，CQA与MLA都是围绕“如何减少KVCache同时尽可能保证效果”来发展出来的。

MQA

Multi-Query Attention，提出自Fast Transformer Decoding: One Write-Head is All You Need.

MQA思路很简单，对于不同的Header，直接公用他们的k与v，即W_k与W_v均相同

使用MQA模型包括PaLM，StarCoder，Gemini

GQA

有人担心MQA压缩过于严重，因此GQA(Grouped-Query Attention)诞生，出自GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

原理很朴素，将所有Head分为g个组，每组共享一对K、V

使用GQA的模型有LLAMA2-70B, LLAMA3全系列

MLA

Multi-head Latent Attention，在deepseek-V2从低秩投影角度引入MLA

缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA - 苏剑林的文章 - 知乎 https://zhuanlan.zhihu.com/p/700588653

MLA的核心思路是思考是否能够在GQA的基础上进一步降低kvcache

$\underset{{\mathbf{c}}_i\in {\mathbb{R}}^{g(d_k+d_v)}}{\underbrace{\left[{\mathbf{k}}_i^{(1)},\cdots ,{\mathbf{k}}_i^{(g)},{\mathbf{v}}_i^{(1)},\cdots ,{\mathbf{v}}_i^{(g)}\right]}}={\mathbf{x}}_i\underset{{\mathbf{W}}_c\in {\mathbb{R}}^{d\times g(d_k+d_v)}}{\underbrace{\left[{\mathbf{W}}_k^{(1)},\cdots ,{\mathbf{W}}_k^{(g)},{\mathbf{W}}_v^{(1)},\cdots ,{\mathbf{W}}_v^{(g)}\right]}}$

这里我们将所有 ${\mathbf{k}}_i^{(s)},{\mathbf{v}}_i^{(s)}$拼在一起记为 ${\mathbf{c}}_i$，相应的投影矩阵也拼在一起记为 ${\mathbf{W}}_c$，注意到一般都有 $d_c=g(d_k+d_v)<d$，所以 ${\mathbf{x}}_i$到 ${\mathbf{c}}_i$的变换就是一个低秩投影。所以，MLA的本质改进不是低秩投影，而是低秩投影之后的工作。

Part 1：MLA如何做低秩投影来减少kvcache

GQA在投影之后做了什么呢？首先它将向量对半分为两份分别作为K、V，然后每一份又均分为份，每一份复制次，以此来“凑”够个Attention Head所需要的K、V。我们知道分割、复制都是简单的线性变换,所以MLA的第一个想法是将这些简单的线性变换换成一般的线性变换，以增强模型的能力：

$\begin{array}{c}{\mathbf{o}}_t=\left[{\mathbf{o}}_t^{(1)},{\mathbf{o}}_t^{(2)},\cdots ,{\mathbf{o}}_t^{(h)}\right]\\ {\mathbf{o}}_t^{(s)}=Attention\left({\mathbf{q}}_t^{(s)},{\mathbf{k}}_{\le t}^{(s)},{\mathbf{v}}_{\le t}^{(s)}\right)\triangleq \frac{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right){\mathbf{v}}_i^{(s)}}{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right)}\\ {\mathbf{q}}_i^{(s)}={\mathbf{x}}_i{\mathbf{W}}_q^{(s)}\in {\mathbb{R}}^{d_k},{\mathbf{W}}_q^{(s)}\in {\mathbb{R}}^{d\times d_k}\\ {\mathbf{k}}_i^{(s)}={\mathbf{c}}_i{\mathbf{W}}_k^{(s)}\in {\mathbb{R}}^{d_k},{\mathbf{W}}_k^{(s)}\in {\mathbb{R}}^{d_c\times d_k}\\ {\mathbf{v}}_i^{(s)}={\mathbf{c}}_i{\mathbf{W}}_v^{(s)}\in {\mathbb{R}}^{d_v},{\mathbf{W}}_v^{(s)}\in {\mathbb{R}}^{d_c\times d_v}\\ {\mathbf{c}}_i={\mathbf{x}}_i{\mathbf{W}}_c\in {\mathbb{R}}^{d_c},{\mathbf{W}}_c\in {\mathbb{R}}^{d\times d_c}\end{array}$

然而，理论上这样是能增加模型能力，但别忘了GQA的主要目的是减少KV Cache，出于节省计算和通信成本的考虑，我们一般会缓存的是投影后 ${\mathbf{k}}_i,{\mathbf{v}}_i$的而不是投影前 ${\mathbf{c}}_i$的或 ${\mathbf{x}}_i$，而MLA的这个做法，通过不同的投影矩阵再次让所有的K、V Head都变得各不相同，那么KV Cache的大小就恢复成跟MHA一样大了，违背了GQA的初衷。

对此，MLA发现，我们可以结合Dot-Attention的具体形式，通过一个简单但不失巧妙的恒等变换来规避这个问题。首先，在训练阶段还是照常进行，此时优化空间不大；然后，在推理阶段，我们利用

${\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }=\left({\mathbf{x}}_t{\mathbf{W}}_q^{(s)}\right)\left({\mathbf{c}}_i{\mathbf{W}}_k^{(s)}\right){}^{\top }={\mathbf{x}}_t\left({\mathbf{W}}_q^{(s)}{\mathbf{W}}_k^{(s)}{}^{\top }\right){\mathbf{c}}_i^{\top }$

这意味着推理阶段，我们可以 ${\mathbf{W}}_q^{(s)}{\mathbf{W}}_k^{(s)}{}^{\top }$将合并起来作为Q的投影矩阵，那么 ${\mathbf{c}}_i$则取代了原本的 ${\mathbf{k}}_i$，同理，在 ${\mathbf{o}}_t$后面我们还有一个投影矩阵，于是的 ${\mathbf{v}}_i^{(s)}={\mathbf{c}}_i{\mathbf{W}}_v^{(s)}$的 ${\mathbf{W}}_v^{(s)}$也可以吸收到后面的投影矩阵中去，于是等效地${\mathbf{v}}_i$也可以用 ${\mathbf{c}}_i$代替，也就是说此时KV Cache只需要存下所有的 ${\mathbf{c}}_i$就行，而不至于存下所有的 ${\mathbf{k}}_i^{(s)}$、${\mathbf{v}}_i^{(s)}$。注意到 ${\mathbf{c}}_i$跟 ${}^{(s)}$无关，也就是说是所有头共享的，即MLA在推理阶段它可以恒等变换为一个MQA。

再次强调，本文的主题是一直都是减少KV Cache，那到目前为止，MLA做到了什么呢？答案是通过不同的投影矩阵来增强了GQA的能力，并且推理时可以保持同样大小的KV Cache。那么反过来，如果我们只需要跟GQA相近的能力，那么是不是就可以再次减少KV Cache了？换言之，$d_c$没必要取 $g(d_k+d_v)$，而是取更小的值（DeepSeek-V2取了512），从而进一步压缩KV Cache，这就是MLA的核心思想。

💡

这里有一点个人的猜测，采用512这样的小lantent给他映射回去，会增加额外的计算量，即过去的k和v仍然需要基于 ${\mathbf{c}}_i$重新进行一点计算，但是考虑到推理场景对kvcache的依赖，这实际上是一种tradeoff；

（注：这里有一个细节，就是 ${\mathbf{W}}_q^{(s)}{\mathbf{W}}_k^{(s)}{}^{\top }$合并成一个矩阵的恒等变换，理论上只有在无限精度下才成立，实际上如果我们使用单精度尤其是BF16的话，经过变换后的精度损失往往还是挺明显的，经过多层累积后可能放大到比较可观的程度，这里可能要根据实际误差看要不要做一些后处理。）

Part 2：兼容RoPE

然而，发现这种MLA将q和k合并为位置无关投影矩阵的策略，实际上无法兼容RoPE，具体的数学证明如下：

$$\begin{gathered} {\mathbf{q}}_i^{(s)}={\mathbf{x}}_i{\mathbf{W}}_q^{(s)}{\mathcal{R}}_i,{\mathbf{k}}_i^{(s)}={\mathbf{c}}_i{\mathbf{W}}_k^{(s)}{\mathcal{R}}_i \\ {\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }=\left({\mathbf{x}}_t{\mathbf{W}}_q^{(s)}{\mathcal{R}}_t\right)\left({\mathbf{c}}_i{\mathbf{W}}_k^{(s)}{\mathcal{R}}_i\right){}^{\top }={\mathbf{x}}_t\left({\mathbf{W}}_q^{(s)}{\mathcal{R}}_{t-i}{\mathbf{W}}_k^{(s)}{}^{\top }\right){\mathbf{c}}_i^{\top } \end{gathered}$$

这里的 ${\mathbf{W}}_q^{(s)}{\mathcal{R}}_{t-i}{\mathbf{W}}_k^{(s)}{}^{\top }$就无法合并为一个固定的投影矩阵了（跟位置差 $t-i$相关），从而MLA的想法无法结合RoPE实现。

最后发布的MLA，采取了一种混合的方法——每个Attention Head的Q、K新增 $d_r$个维度用来添加RoPE，其中K新增的维度每个Head共享：

$\begin{array}{c}{\mathbf{o}}_t=\left[{\mathbf{o}}_t^{(1)},{\mathbf{o}}_t^{(2)},\cdots ,{\mathbf{o}}_t^{(h)}\right]\\ {\mathbf{o}}_t^{(s)}=Attention\left({\mathbf{q}}_t^{(s)},{\mathbf{k}}_{\le t}^{(s)},{\mathbf{v}}_{\le t}^{(s)}\right)\triangleq \frac{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right){\mathbf{v}}_i^{(s)}}{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right)}\\ {\mathbf{q}}_i^{(s)}=\left[{\mathbf{x}}_i{\mathbf{W}}_{qc}^{(s)},{\mathbf{x}}_i{\mathbf{W}}_{qr}^{(s)}{\mathcal{R}}_i\right]\in {\mathbb{R}}^{d_k+d_r},{\mathbf{W}}_{qc}^{(s)}\in {\mathbb{R}}^{d\times d_k},{\mathbf{W}}_{qr}^{(s)}\in {\mathbb{R}}^{d\times d_r}\\ {\mathbf{k}}_i^{(s)}=\left[{\mathbf{c}}_i{\mathbf{W}}_{kc}^{(s)},{\mathbf{x}}_i{\mathbf{W}}_{kr}^{\bcancel{(s)}}{\mathcal{R}}_i\right]\in {\mathbb{R}}^{d_k+d_r},{\mathbf{W}}_{kc}^{(s)}\in {\mathbb{R}}^{d_c\times d_k},{\mathbf{W}}_{kr}^{\bcancel{(s)}}\in {\mathbb{R}}^{d\times d_r}\\ {\mathbf{v}}_i^{(s)}={\mathbf{c}}_i{\mathbf{W}}_v^{(s)}\in {\mathbb{R}}^{d_v},{\mathbf{W}}_v^{(s)}\in {\mathbb{R}}^{d_c\times d_v}\\ {\mathbf{c}}_i={\mathbf{x}}_i{\mathbf{W}}_c\in {\mathbb{R}}^{d_c},{\mathbf{W}}_c\in {\mathbb{R}}^{d\times d_c}\end{array}$

这样一来，没有RoPE的维度就可以重复“Part 1”的操作，在推理时KV Cache只需要存 ${\mathbf{c}}_i$，新增的带RoPE的维度就可以用来补充位置信息，并且由于所有Head共享，所以也就只有在K Cache这里增加了$d_r$个 维度，原论文取了 $d_r=d_k/2=64$，相比原本的 $d_c=512$，增加的幅度不大。

Part 3：Final version

最后有一个细节，就是MLA的最终版本，还将Q的输入也改为了低秩投影形式，这与减少KV Cache无关，主要是为了减少训练期间参数量和相应的梯度（原论文说的是激活值，个人表示不大理解）所占的显存：

$\begin{array}{c}{\mathbf{o}}_t=\left[{\mathbf{o}}_t^{(1)},{\mathbf{o}}_t^{(2)},\cdots ,{\mathbf{o}}_t^{(h)}\right]\\ {\mathbf{o}}_t^{(s)}=Attention\left({\mathbf{q}}_t^{(s)},{\mathbf{k}}_{\le t}^{(s)},{\mathbf{v}}_{\le t}^{(s)}\right)\triangleq \frac{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right){\mathbf{v}}_i^{(s)}}{{\sum }_{i\le t}\exp \left({\mathbf{q}}_t^{(s)}{\mathbf{k}}_i^{(s)}{}^{\top }\right)}\\ {\mathbf{q}}_i^{(s)}=\left[{\mathbf{c}}_i^{\prime }{\mathbf{W}}_{qc}^{(s)},{\mathbf{c}}_i^{\prime }{\mathbf{W}}_{qr}^{(s)}{\mathcal{R}}_i\right]\in {\mathbb{R}}^{d_k+d_r},{\mathbf{W}}_{qc}^{(s)}\in {\mathbb{R}}^{d_c^{\prime }\times d_k},{\mathbf{W}}_{qr}^{(s)}\in {\mathbb{R}}^{d_c^{\prime }\times d_r}\\ {\mathbf{k}}_i^{(s)}=\left[{\mathbf{c}}_i{\mathbf{W}}_{kc}^{(s)},{\mathbf{x}}_i{\mathbf{W}}_{kr}^{\bcancel{(s)}}{\mathcal{R}}_i\right]\in {\mathbb{R}}^{d_k+d_r},{\mathbf{W}}_{kc}^{(s)}\in {\mathbb{R}}^{d_c\times d_k},{\mathbf{W}}_{kr}^{\bcancel{(s)}}\in {\mathbb{R}}^{d\times d_r}\\ {\mathbf{v}}_i^{(s)}={\mathbf{c}}_i{\mathbf{W}}_v^{(s)}\in {\mathbb{R}}^{d_v},{\mathbf{W}}_v^{(s)}\in {\mathbb{R}}^{d_c\times d_v}\\ {\mathbf{c}}_i^{\prime }={\mathbf{x}}_i{\mathbf{W}}_c^{\prime }\in {\mathbb{R}}^{d_c^{\prime }},{\mathbf{W}}_c^{\prime }\in {\mathbb{R}}^{d\times d_c^{\prime }}\\ {\mathbf{c}}_i={\mathbf{x}}_i{\mathbf{W}}_c\in {\mathbb{R}}^{d_c},{\mathbf{W}}_c\in {\mathbb{R}}^{d\times d_c}\end{array}$

一个比较好理解的图如下：

其中两个标红的部分为kvcache需要储存的部分，左侧的 $k_i^R$用作RoPE，而右侧的 $c_i^{KV}$则为上文所讲的低秩投影后的latent结果；