大模型算法面试笔记——多头潜在注意力（MLA）

注意力机制解决的问题：传统序列处理模型如RNN和LSTM，捕捉长距离依赖关系的难题。注意力机制允许模型在序列的不同位置之间建立直接联系，有效捕捉远距离依赖关系。

为了减少推理过程中KV Cache占用的显存，GQA和MQA通过head之间共享KV实现，这是一种牺牲性能对存储空间妥协的方案，而MLA通过对KV对做低秩联合压缩来减少推理中的KV缓存，目标是减少kv cache存储量的同时，保存模型的效果。

具体做法是，对于每个token，先通过一个低秩矩阵将KV联合压缩到一个低维向量$c^{KV}$中，然后通过两个升维矩阵$W^{UK}$，$W^{UV}$解压缩回高维，后续进行普通的多头注意力计算，这样每次只需要存这个低维向量。

这样做有个问题，就是压缩和解压操作使计算量增加了，而实际计算中，通过“矩阵吸收”操作，也就是矩阵运算过程中的结合律使多个矩阵合并，从而减少计算量。
具体计算过程如下（对$q$做相同压缩操作，于是也有了$c^Q$和$W^{UQ}$）：

$\mathrm{attention}=\mathrm{softmax}(\frac{q{k}^T}{\sqrt{d}})v{W}^O$
$=\mathrm{softmax}(\frac{c^Q{W}^{UQ}(c^{KV}{W}^{UK}{)}^T}{\sqrt{d}})c^{KV}{W}^{UV}{W}^O$
$=\mathrm{softmax}(\frac{c^Q(W^{UQ}(W^{UK}{)}^T)(c^{KV}{)}^T}{\sqrt{d}})c^{KV}(W^{UV}{W}^O)$
$=\mathrm{softmax}(\frac{c^Q{W}^{UQUK}(c^{KV}{)}^T}{\sqrt{d}})c^{KV}{W}^{UVO}$

如上所示，计算过程中，由于矩阵乘法结合律，$W^{UQ}$，$W^{UK}$合并成一个矩阵$W^{UQUK}$，同理，$W^{UV}$，$W^O$合并成$W^{UVO}$。

对比普通MHA计算公式：

$\mathrm{attention}=\mathrm{softmax}(\frac{q{k}^T}{\sqrt{d}})v{W}^O$
$=\mathrm{softmax}(\frac{h_t{W}^Q{K}^T}{\sqrt{d}})v{W}^O$

可知，两种注意力机制计算量相同，没有引入额外计算量，而缓存从两个高维$K$,$V$变成了一个低维$c^{KV}$

接下来是位置编码RoPE的处理，MHA中，RoPE可以通过对$q$,$k$向量乘以一个位置相关的变换矩阵$R_i$($i$为当前token所处的位置)。然而，在MLA中，如果做相同的处理将会如下所示：

$q_i{R}_i(k_j{R}_j{)}^T=c^Q{W}^{UQ}{R}_i(c_j^{KV}{W}^{UK}{R}_j{)}^T$
$=c^Q{W}^{UQ}{R}_i{R}_j^T(W^{UK}{)}^T(c_j^{KV}{)}^T$

由于$R_i$不是一个固定的矩阵，无法实现矩阵吸收来减少计算量。对于这个问题，deepseek的做法是将参与注意力计算的$q$，$v$分成两部分，一部分进行矩阵吸收操作，不带位置信息，一部分进行位置信息计算。

对于$q$，基于潜在向量$c^Q$通过矩阵$W^{QR}$变换为低维向量后进行RoPE变换得到$q^R$；对于$k$，直接将输入$h_t$也通过一个矩阵$W^{KR}$变换后做RoPE变换得到$k^R$，其中$k^R$按照MQA的处理方式，各个head之间共享，既减少了显存调用又保证了位置编码的全局一致。然后将$q_R$,$k^R$拼接到前面计算得到的$q$,$k$向量后面，得到最终用于计算注意力的$q=[q^C;q^R]$,$k=[k^C;k^R]$。
这样计算点积时如下，其中$t$,$j$表示token，$i$表示head：

$$q_{t,i}{k}_{j,i}^T=[q_{t,i}^C;q_{t,i}^R]\times [k_{j,i}^C;k_t^R]=q_{t,i}^C{k}_{j,i}^C+q_{t,i}^R{k}_t^R$$

这样不包含位置编码的部分$q_{t,i}^C{k}_{j,i}^C$就可以进行矩阵吸收的处理，每个head缓存一个$c_{t,i}^{KV}$；后一项按MQA的方式计算，所有head只需缓存一个共享的$k_t^R$。

学习参考资料：zhihu-冷面爸