【大模型】DeepSeek核心技术之MLA (Multi-head Latent Attention)

在讲解MLA之前，需要大家对几个基础的概念（KV Cache， Grouped-Query Attention (GQA), Multi-Query Attention (MQA)，RoPE）有所了解，这些有助于理解MLA是怎么工作的，为什么需要这么做。

酒酿小圆子～

1489人浏览 · 2025-03-01 15:50:19

酒酿小圆子～ · 2025-03-01 15:50:19 发布

文章目录

1. Multi-Head Attention (MHA)
2. Multi-head Latent Attention (MLA)
3. 总结
参考资料

在讲解MLA之前，需要大家对几个基础的概念（KV Cache， Grouped-Query Attention (GQA), Multi-Query Attention (MQA)，RoPE）有所了解，这些有助于理解MLA是怎么工作的，为什么需要这么做。
这里给出概念及对应的讲解博客：

MHA，MQA，GQA及KV Cache：【大模型】MHA，MQA，GQA及KV Cache详解
旋转位置编码RoPE：【大模型】旋转位置编码（Rotary Position Embedding，RoPE）

1. Multi-Head Attention (MHA)

首先跟着DeepSeek V2的论文简单回顾一下Multi-Head Attention（MHA）的计算过程，首先给出各个变量的含义如下：

$d$ 代表输入维度（input dim）
$n_h$ 代表头数（head数）
$d_h$ 代表每个头的维度
$h_t$ 代表输入的第 t 个向量
$l$ 代表 transformer 的层数

主要公式如下：
在这里插入图片描述

下面来介绍下公式的含义：

$W_Q, W_K, W_V \in R^{d_hn_h*d}$ 表示输入维度，公式(1)-(3)中我们只使用一个矩阵来处理多头（multi-head）
公式(4)-(6)表示对 $q_t, k_t, v_t \in R^{d_hn_h}$ 进行分割，可以得到每个头对应的 q, k, v
公式(7)表示对q, k 进行softmax操作，然后再乘上v
公式(8)表示对多头输出的结果进行拼接操作，再乘上 $W^o$ 得到最终的输出

下面我们分析下，KV Cache的占用量：

对于标准的MHA而言，对于每一个token，KV Cache占用的缓存的大小为 $2n_hd_hl$ 。

后续我们要介绍的MLA就是致力于在推理过程中降低 $n_{h} d_{h}$ 。

2. Multi-head Latent Attention (MLA)

参考资料：全网最细！DeepSeekMLA 多头隐变量注意力：从算法原理到代码实现

2.1 低秩压缩

MLA的核心是对KV做了低秩压缩（Low-Rank Key-Value Joint Compression），在送入标准的MHA算法前，用一个更短的向量来表示原来的长向量，从而大幅减少KV Cache空间。

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

这里给出MLA的整体结构图：
在这里插入图片描述

这里先给出图中一些字母和符号的含义，方便我们后续理解。

q: query

k: key

v: value

$h_t$ : 输入的第 t 个向量

C: compress 压缩

R: RoPE 旋转位置编码

D: down 下采样，降维

U: up 上采样，升维

MLA的核心是对KV做了低秩压缩（Low-Rank Key-Value Joint Compression）来减少KV cache，公式如下：
在这里插入图片描述

公式(9)中通过下采样矩阵，对输入 $h_t$ 进行压缩得到中间表示 ${C^{KV}_t}$ ，再基于公式(10)和(11)进行上采样升维度还原KV。
KV cache占用空间大幅下降。从MLA的架构图上可以看到，需要缓存的元素为 ${C^{KV}_t}$ 和 ${k^{R}_t}$ 。这里我们先主要关注 ${C^{KV}_t} \in R^{d_c}$ ，且 $d_c<<{n_h}{d_h}$ 。前面我们提到，对于标准的MHA而言，每一个token的KV Cache大小为 $2n_hd_hl$ 。而对MLA而言，每一步token的推理产生的缓存变成 $d_c l$ ，缓存的矩阵大小相比于原始KV做了压缩，因此缓存量大幅下降。（补充：deepseek v2中 $d_{c}$ 被设置为 $4 d_{h}$ ）
在MLA中，同时也压缩了 query 向量。我们知道在KV Cache中，Q的作用只发生在当下，无需缓存。但是在模型训练的过程中，每个输入的token会通过多头注意力机制生成对应的query、key和value，这些中间数据的维度往往非常高，因此占用的内存量也相应很大。所以论文中也提到为了降低训练过程中的激活内存activation memory，DeepSeek-V2还对queries进行低秩压缩。对Q的压缩方式和K、V一致，依然是先降维再升维，这个操作并不能降低KV Cache，而是降低内容占用，另外一方面也可以使得query 和key, value 能在同一个低维空间进行一致性表示。

2.2 应用RoPE

上面这种低秩压缩的计算方式，对于RoPE旋转位置编码是有影响的，因为压缩操作可能已经丢失了某些信息，使得位置编码不能直接和有效地反映原始Q和K的位置关系。因此，不能直接在压缩后的向量上应用RoPE。

那么可不可以在解压后的向量上应用RopE呢？ 可以，但是影响效率，因为前面已经说过不显式计算解压后的向量，而是直接应用压缩后的向量。

如何解决呢？Deepseek-V2设计了两个新的向量（ ${q^{R}_{t,i}}$ 和 ${k^R_t}$ ），单独应用RoPE，将位置信息写入这个新的向量中。

在这里插入图片描述
所以到目前为止，我们得到的 $q, k$ 包括两部分拼接而成：一部分是做了低秩压缩得到的 ${q^C, k^C}$ 向量，一部分是增加了RoPE位置编码的 ${q^R, k^R}$ 向量。（后面这部分 ${q^R, k^R}$ 向量是基于MQA方式计算得到的，所有Head共享1个 $k$ ）。

如何理解上述的操作过程？这也是MLA方法的核心。

我们先来看看DeepSeek-V2论文中有一段原文解释（中文翻译）：
位置编码使用RoPE，但RoPE与低秩KV不兼容。具体来说，RoPE对Q和K都是位置敏感的。如果我们为应用RoPE，那么公式（42）的 $W$ （K的权重矩阵）将与位置敏感的RoPE矩阵耦合。因此，在推理过程中， $W^{UK}$ 无法再被吸收到 $W^{UQ}$ （Q的权重矩阵）中，因为与当前生成的token相关的RoPE矩阵将位于 $W^{UK}$ 和 $W^{UQ}$ 之间，而矩阵乘法不满足交换律。因此，我们必须在推理过程中重新计算所有前缀token的k，这将极大地降低推理效率。

注意：MLA不是不能加RoPE，而是加了RoPE之后无法用恒等变换技巧来减少KV Cache。

需要注意的是，在对 ${k^R_t}$ 进行编码时，它是直接从input hidden $h_t$ 上来的，也就是 $k$ 向量不需要进行先降维、后升维的操作。
压缩完、且RoPE编码完之后，最后将这4个变量（ $q_{t}^{C}=W^{U Q} c_{t}^{Q}$ 、 $\mathbf{k}_{t}^{C}=W^{U K} \mathbf{c}_{t}^{K V}$ 、 $\mathbf{q}_{t}^{R}$ 、 $\mathbf{k}_{t}^{R}$ ）分别拼接起来，形成 带信息压缩 的和 带位置信息 的向量。
- 带信息压缩：Query—— $\mathbf{q}_{t}^{C}$ ，Key—— $\mathbf{k}_{t}^{C}$
- 带位置信息：Query—— $\mathbf{q}_{t}^{R}$ ，Key—— $\mathbf{k}_{t}^{R}$
最后将拼接后的 $q_{t,i}$ 和 $k_{t,i}$ ，结合 ${v}_{t}^{C}$ 来进行后续的multi-head attention的计算（也就是seft-attention的常规计算那一套流程）。

在这里插入图片描述
MHA公式：

在这里插入图片描述

2.3 矩阵融合

从前面的整体结构图中，我们看到向量 $\mathbf{c}_{t}^{K V}$ 、 $\mathbf{k}_{t}^{R}$ 需要缓存以进行生成。在推理过程中，常规做法需要从 $\mathbf{c}_{t}^{K V}$ 中恢复 $\mathbf{k}_{t}^{C}$ 和 $\mathbf{v}_{t}^{C}$ 以进行注意力计算。

在DeepSeek V2中巧妙地利用了矩阵融合操作，将上采样矩阵 $W^{UK}$ 融合到 $W^{UQ}$ 中，并将 $W^{UV}$ 融合到 $W^{O}$ 中。也就是说不需要显示地去计算得到 ${k}_{t}^{C}$ 和 ${v}_{t}^{C}$ ，而可以直接基于 ${C_t}^{KV}$ 进行计算，避免了在推理过程中重复计算 ${k}_{t}^{C}$ 和 ${v}_{t}^{C}$ 的开销。

这里解释一下什么是 矩阵融合(can be absorbed into) 操作。后续计算的时候甚至都不需要显示进行融合操作，而是由神经网络自动通过训练进行的，我们仅需要对压缩后的隐向量操作即可。

在这里插入图片描述

最终，MLA单个Token产生的缓存包含了两个部分，即 $\left(d_{c}+d_{h}^{R}\right) l$ ，实现了计算量小且效果优于MHA的结果。

在这里插入图片描述

3. 总结

MLA可被视为GQA的一般化，它用投影矩阵的方式替代了GQA的分割、重复，并引入了一个恒等变换技巧来可以进一步压缩KV Cache，同时采用了一种混合方法来兼容RoPE。总的来说，MLA称得上是一种非常实用的注意力变体。

参考资料

DeepSeek技术社区

欢迎加入DeepSeek 技术社区。在这里，你可以找到志同道合的朋友，共同探索AI技术的奥秘。

更多推荐

2025年程序员转行大模型全攻略：职业前景、热门岗位与选择策略解析！

DeepSeek技术社区

HAproxy服务器带日志

我现在有两台WindowsAD域控 10.10.10.1 和10.10.10.2，我现在有一个准入系统10.10.11.100需要提供ldap认证，我现在想搭建一个haproxy服务器IP 10.10.10.200做ldap业务的负载均衡，这样在准入系统上只要填haproxy服务器地址即可，所有通过haproxy服务器获取ladp的信息在haproxy需要有日志，haproxy服务已经安装，配置