Multi-Head Latent Attention (MLA)详解

2025-08-05

Multi-Head Latent Attention (MLA)详解

论文 DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
github: DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model

参考博客：

洞见：
1.位置编码目前是添加到Q和K中，不再是直接添加到embedding中
2.RoPE虽然叫位置编码，但是自变量除了位置之外还有embeding维度。所以如果嵌入维度变了，位置编码也会变。
	所以可以只在固定维度添加位置编码？保证维度变了，位置编码不变？

1. Multi-Head Attention (MHA) 回顾

标准的多头注意力： $d$ 表示嵌入维度， $n_h$ 表示注意力头的数量， $d_h$ 表示每个头的维度， $h_t \in R^d$ 表示在注意力层中第t个token的注意力输入。

标准MHA机制通过三个矩阵 $W^Q、W^K、W^V \in \mathbb{R}^{d_h n_h \times d}$ 分别生成输入 $h_t$ 的查询向量 $q_t\in \mathbb{R}^{𝑑_ℎ}$ 、键向量 $k_t\in \mathbb{R}^{𝑑_ℎ}$ 和值向量 $v_t\in \mathbb{R}^{𝑑_ℎ}$ 。公式如下：

$q_t = W^Q h_t \tag{1}$

$k_t = W^K h_t \tag 2$

$v_t = W^V h_t \tag 3$

接下来， $𝑞_𝑡$ ， $k_𝑡$ ， $v_𝑡$ 将被切分成 $n_h$ 个头，用于多头注意力计算：

$\begin{equation} [𝑞_{𝑡,1}; 𝑞_{𝑡,2}; \ldots; 𝑞_{𝑡,𝑛_ℎ}] = 𝑞_𝑡 \quad \tag{4} \end{equation}$

$[𝑘_{𝑡,1}; 𝑘_{𝑡,2}; \ldots; 𝑘_{𝑡,𝑛_ℎ}] = 𝑘_𝑡 \quad \tag{5}$

$[𝑣_{𝑡,1}; 𝑣_{𝑡,2}; \ldots; 𝑣_{𝑡,𝑛_ℎ}] = 𝑣_𝑡 \quad \tag 6$

其中：

$𝑞_{𝑡,i}, 𝑘_{𝑡,i}, 𝑣_{𝑡,i} \in \mathbb{R}^{𝑑_ℎ}$ 是第 $𝑡$ 个 token 的第i个头的 query、key、value 向量。
$𝑛_ℎ$ 是注意力头数，QKV分头前的特征维度为 $d$ ，分头后每个头的维度是 $𝑑_ℎ$ 。

公式含义：将原始的 $𝑞_𝑡, 𝑘_𝑡, 𝑣_𝑡$ 按列切分成 $𝑛_ℎ$ 个头的子向量。

每个头的注意力计算公式为：

$attention\_weight_{t,i} = \text{Softmax}_{j=1}^{t}\left(\frac{𝑞_{𝑡,𝑖}^\top 𝑘_{𝑗,𝑖}}{\sqrt{𝑑_ℎ}}\right)$

$𝑜_{𝑡,𝑖} = \sum_{𝑗=1}^{𝑡} \text{Softmax}\left(\frac{𝑞_{𝑡,𝑖}^\top 𝑘_{𝑗,𝑖}}{\sqrt{𝑑_ℎ}}\right) 𝑣_{𝑗,𝑖} \tag 7$

最后拼接所有头的输出，再做输出投影：

$𝑢_𝑡 = 𝑊^𝑂 [𝑜_{𝑡,1}; 𝑜_{𝑡,2}; \ldots; 𝑜_{𝑡,𝑛_ℎ}] \tag 8$

其中 $𝑊^𝑂 \in \mathbb{R}^{𝑑 \times 𝑑_ℎ 𝑛_ℎ}$ 是输出映射矩阵。

推理时问题：

所有 key 和 value 需要缓存，KV cache大小为 $2𝑛_ℎ 𝑑_ℎ 𝑙$ （ $𝑙$ 为序列长度），当 batch size 和 $𝑙$ 很大时会占用大量显存。

2. Low-Rank Key-Value Joint Compression（低秩 KV 联合压缩）

MLA 核心是通过低秩联合压缩减少 KV cache大小。

压缩输入:

将输入 $ℎ_𝑡$ 压缩到KV低维共享表示 $𝑐_{𝑡}^{𝐾𝑉}$ ：

$𝑐_{𝑡}^{𝐾𝑉} = 𝑊^{𝐷𝐾𝑉} ℎ_𝑡 \tag 9$
- $𝑐_{𝑡}^{𝐾𝑉} \in \mathbb{R}^{𝑑_𝑐}$ ：压缩后的 latent 表示， $𝑑_𝑐 \ll 𝑑_ℎ 𝑛_ℎ$ 。
- $𝑊^{𝐷𝐾𝑉} \in \mathbb{R}^{𝑑_𝑐 \times 𝑑}$ ：降维矩阵，D的含义是降维down_sample。
解码 key 和 value:

从 $𝑐_{𝑡}^{𝐾𝑉}$ 解码出 key 和 value：

$𝑘_{𝑡}^{𝐶} = 𝑊^{𝑈𝐾} 𝑐_{𝑡}^{𝐾𝑉} \tag {10}$

$𝑣_{𝑡}^{𝐶} = 𝑊^{𝑈𝑉} 𝑐_{𝑡}^{𝐾𝑉} \tag {11}$
- $𝑊^{𝑈𝐾}, 𝑊^{𝑈𝑉} \in \mathbb{R}^{𝑑_ℎ 𝑛_ℎ \times 𝑑_𝑐}$ ：升维矩阵，U的含义是升维up_sample。
**低秩压缩后的注意力权重计算方式 **：

$𝑞_{𝑡}^\top 𝑘_{j}^{𝐶} = (W^Q h_t)^\top 𝑊^{𝑈𝐾} 𝑐_{𝑡}^{𝐾𝑉} = h_t^\top (W^Q)^\top 𝑊^{𝑈𝐾} 𝑐_{𝑡}^{𝐾𝑉} = \\ h_t (𝑊^{𝑈𝐾})^\top W^Q 𝑐_{𝑡}^{𝐾𝑉}$

最后的输出:

多头注意力相当于把 注意力权重 $attention\_weight$ 从一维向量变为了二维向量，shape: (seq_len,) -> (seq_len, head_num)，但本质上还是一个张量，只是得到这个张量的方式注意力机制计算比较复杂。

$[𝑜_{𝑡,1}; 𝑜_{𝑡,2}; \ldots; 𝑜_{𝑡,𝑛_ℎ}] = v_t @ attention\_weight = \\ 𝑊^{𝑈𝑉} 𝑐_{𝑡}^{𝐾𝑉} @ attention\_weight$

$𝑢_𝑡 = 𝑊^𝑂 [𝑜_{𝑡,1}; 𝑜_{𝑡,2}; \ldots; 𝑜_{𝑡,𝑛_ℎ}] \\ 𝑢_𝑡 = 𝑊^𝑂 𝑊^{𝑈𝑉} (𝑐_{𝑡}^{𝐾𝑉} @ attention\_weight)$

优势：

推理时仅需缓存 $𝑐_{𝑡}^{𝐾𝑉}$ （每 token 只需 $𝑑_𝑐$ 维），KV 缓存大小从 $2𝑑_ℎ 𝑛_ℎ 𝑙$ 缩减到 $𝑑_𝑐 𝑙$ 。
可将 $𝑊^{𝑈𝐾} \in \mathbb{R}^{𝑑_ℎ 𝑛_ℎ \times 𝑑_𝑐}$ 与 $𝑊^𝑄 \in \mathbb{R}^{d_h n_h \times d}$ 合并，将 $𝑊^{𝑈𝑉} \in \mathbb{R}^{𝑑_ℎ 𝑛_ℎ \times 𝑑_𝑐}$ 与 $𝑊^𝑂 \in \mathbb{R}^{𝑑 \times 𝑑_ℎ 𝑛_ℎ}$ 合并，无需显式生成或存储 key/value。

矩阵合并分析：

上面公式看似需要使用 $𝑐_{𝑡}^{𝐾𝑉}$ 重新生成历史的KV，但是通过权重吸收(合并)可以将重新生成的步骤合并其他向量的投影中，从而无需真正重新生成。

$𝑊^{𝑈𝐾} \in \mathbb{R}^{𝑑_ℎ 𝑛_ℎ \times 𝑑_𝑐}$ 与 $𝑊^𝑄 \in \mathbb{R}^{d_h n_h \times d}$ 的合并，维度变换是 $d -> d_hn_h -> d_c$ ，吸收后是 $d -> d_c$ ，如果是原始的MHA不进行维度压缩 $d -> d_hn_h$ ，可以看出来吸收后相比MHA计算量减少，不吸收则增大计算量。

$𝑊^{𝑈𝑉} \in \mathbb{R}^{𝑑_ℎ 𝑛_ℎ \times 𝑑_𝑐}$ 与 $𝑊^𝑂 \in \mathbb{R}^{𝑑 \times 𝑑_ℎ 𝑛_ℎ}$ 合并，维度变换是 $d_c -> d_hn_h -> d$ , 吸收后是 $d_c -> d$ ，同样吸收后相比MHA计算量减少，不吸收则增大计算量。

3. Low-Rank Query Compression（低秩 Query 压缩）

虽然压缩 Query 不能减少 KV 缓存，但是也能节省计算。

压缩 Query：
$𝑐_{𝑡}^{𝑄} = 𝑊^{𝐷𝑄} ℎ_𝑡 \quad \tag {12}$

$𝑞_{𝑡}^{𝐶} = 𝑊^{𝑈𝑄} 𝑐_{𝑡}^{𝑄} \quad \tag {13}$
- $𝑐_{𝑡}^{𝑄} \in \mathbb{R}^{𝑑_𝑐'}$ ：压缩后的 query 表示， $𝑑_𝑐' \ll 𝑑_ℎ 𝑛_ℎ$ 。
- $𝑊^{𝐷𝑄} \in \mathbb{R}^{𝑑_𝑐' \times 𝑑}$ ， $𝑊^{𝑈𝑄} \in \mathbb{R}^{𝑑_ℎ 𝑛_ℎ \times 𝑑_𝑐'}$ ：降维和升维矩阵。

作用：

降维后再升维，减少全连接计算量，尤其适用于输入维度 $𝑑$ 较大的场景。

4. Decoupled Rotary Position Embedding（解耦 RoPE）

背景

前面的推理没有考虑位置编码，如果考虑RoPE位置编码，会发现权重吸收和位置编码不兼容：

$𝑞_{𝑡}^\top 𝑘_{j}^{𝐶} => RoPE(𝑞_{𝑡}^\top)RoPE(𝑘_{j}^{𝐶}) \tag {-1} \\$

$RoPE(𝑞_{𝑡}^\top)RoPE(𝑘_{j}^{𝐶}) = (R_t𝑞_{𝑡})^\top R_j 𝑘_{j}^{𝐶} = \\ 𝑞_{𝑡}^\top R_t^\top R_j 𝑘_{j}^{𝐶} = 𝑞_{𝑡}^\top R_{j-t} 𝑘_{j}^{𝐶} = h_t^\top (W^Q)^\top R_{j-t} 𝑊^{𝑈𝐾} 𝑐_{𝑡}^{𝐾𝑉} \tag {-2}$

前面 $𝑊^{𝑈𝐾}$ 与 $𝑊^𝑄$ 可以合并，是因为$(W^Q)\top 𝑊^{𝑈𝐾} $是常量，但是$ (W^Q)\top R_{j-t} 𝑊^{𝑈𝐾} $是与query向量的位置$ t $相关的变量$ R_{j-t}$，即随推理过程中query的位置变化而变化，导致旧的缓存不能直接使用，所以无法直接合并。

为什么GQA的低秩投影没有受RoPE影响，因为GQA只减少了头数，计算时又恢复了头数，但是隐式恢复头数直接通过广播实现，不涉及特征维度改变序列长度改变；而MLA减小的是每个头的特征维度，计算时为了避免增加计算量，只能隐式恢复特征维度，即需要借助权重吸收，而两个权重之间又夹着RoPE，没办法权重吸收。

RoPE（Rotary Position Embedding）对 key 和 query 都是位置敏感的。

若直接将 RoPE 应用于压缩后的 key，升维矩阵 $𝑊^{𝑈𝐾}$ 无法与 $𝑊^𝑄$ 合并，影响推理效率？

一个简单的逻辑应该是把位置编码都应用到压缩后的特征上。

原因

RoPE 是依赖位置的旋转矩阵，矩阵乘法不满足交换律。
若在升维后的 key 上应用 RoPE，需重新计算所有历史 key 的位置编码，无法仅用缓存恢复。

解决方案

将 query 和 key 分为两部分，一部分就用MLA不使用RoPE 位置编码，另一部分用MQA使用位置编码：

对 query 和 key 分别应用 RoPE：

将输入的查询内容向量 $c^Q_t$ 通过矩阵 $W^{QR}$ 投影到 RoPE 编码后的空间，得到分离后的旋转位置编码查询向量 $q^R_t$ ，并按多头 $n_h$ 分块。

$[𝑞_{𝑡,1}^𝑅; 𝑞_{𝑡,2}^𝑅; \ldots; 𝑞_{𝑡,𝑛_ℎ}^𝑅] = \text{RoPE}(𝑊^{𝑄𝑅} 𝑐_{𝑡}^{𝑄}) \quad \text{(14)}$

将隐藏状态 $ h_t $ 通过矩阵 $W^{KR}$ 投影，并应用 RoPE 得到分离后的旋转位置编码键向量$ k^R_t $。

$𝑘_{𝑡}^𝑅 = \text{RoPE}(𝑊^{𝐾𝑅} ℎ_𝑡) \quad \text{(15)}$
拼接压缩部分和 RoPE 部分：

将内容查询向量与旋转位置编码查询向量进行拼接，得到最终的第 $i$ 个头的查询向量。

$𝑞_{𝑡,𝑖} = [𝑞_{𝑡,𝑖}^𝐶, 𝑞_{𝑡,𝑖}^𝑅] \quad \text{(16)}$

将内容键向量与旋转位置编码键向量拼接，得到最终的第 $i$ 个头的键向量。

$𝑘_{𝑡,𝑖} = [𝑘_{𝑡,𝑖}^𝐶, 𝑘_{𝑡,𝑖}^𝑅] \quad \text{(17)}$
注意力计算：

通过点积计算第 $i$ 个头在时间步 $t$ 的查询向量与历史键向量的相关性，除以 $\sqrt{d_h + d^R_h}$ 进行缩放，然后使用 Softmax 得到注意力权重，对对应的值向量 $v^C_{j,i}$ 进行加权求和。

$𝑜_{𝑡,𝑖} = \sum_{𝑗=1}^{𝑡} \text{Softmax}\left(\frac{𝑞_{𝑡,𝑖}^\top 𝑘_{𝑗,𝑖}}{\sqrt{𝑑_ℎ + 𝑑_ℎ^𝑅}}\right) 𝑣_{𝑗,𝑖}^𝐶 \quad \text{(18)}$
最终输出：

将所有头的输出 $o_{t,i}$ 拼接起来，并通过输出权重矩阵 $W^O$ 得到最终的输出向量 $u_t$ 。

$𝑢_𝑡 = 𝑊^𝑂 [𝑜_{𝑡,1}; 𝑜_{𝑡,2}; \ldots; 𝑜_{𝑡,𝑛_ℎ}] \quad \text{(19)}$

其中 $W^{QR} \in \mathbb{R}^{d^R_h n_h \times d'_c}$ 和$ W^{KR} \in \mathbb{R}^{dR_h \times d} $是分别用于生成分离查询向量和分离键向量的矩阵；$ RoPE(·) $表示应用旋转位置编码的操作；符号$ [\cdot ;\cdot]$表示拼接操作。在推理阶段，分离的键向量也需要被缓存。因此，MLA需要一个包含 $(d_c + d^R_h)l$ 元素的 KV 缓存。

参数说明：

$𝑑_ℎ^𝑅$ ：RoPE 部分的维度。
推理时只需缓存 $𝑐_{𝑡}^{𝐾𝑉}$ （大小约为 $𝑑_𝑐 + 𝑑_ℎ^𝑅$ ）。

5. 总结 MLA 思路

MLA 的核心目标是降低推理时的显存占用和计算延迟，同时保持注意力效果。具体方法包括：

低秩联合压缩 KV：通过共享压缩表示 $𝑐_{𝑡}^{𝐾𝑉}$ 减少 KV 缓存大小。
低秩 Query 压缩：减少 Query 的计算量。
解耦 RoPE：在保留位置编码效果的同时，使压缩策略兼容 RoPE，避免破坏缓存复用。

最终效果：

显存占用从 $2𝑑_ℎ 𝑛_ℎ 𝑙$ 降至 $𝑑_𝑐 𝑙$ 。
计算效率提升，适用于大规模模型部署。