[latexpage]
最近，幻方发布的 DeepSeek-V2 引发了广泛关注。它不仅以每百万 token 仅需 1 块钱的价格震惊了大家，比现有的 API 便宜了两个数量级，甚至有人开玩笑说：「这个价格哪怕输出乱码，我也会认为是一种艺术。」 从技术报告来看，这样低廉价格背后的关键技术之一是新提出的 MLA（Multi-head Latent Attention），这是对 GQA 的改进。据说这种改进不仅更加高效，还能提高性能，引起了许多读者的兴趣。本文将带大家梳理从 MHA 、 MQA 、 GQA 到 MLA 的演变过程，并重点介绍 MLA 的设计思路。
MHA: 多头注意力
首先，我们来看看 MHA（Multi-Head Attention），即多头注意力。这是经典论文 《Attention is All You Need 》中提出的一种注意力机制，可以说是当前主流大规模语言模型 （LLM） 的基础。
简单来说，多头注意力将输入的向量序列分成多个部分，每部分单独计算注意力，然后再将结果拼接在一起。具体公式如下：
\[
\begin{aligned}
\boldsymbol{o}t &= \left[\boldsymbol{o}_t^{(1)}, \boldsymbol{o}_t^{(2)}, \cdots, \boldsymbol{o}_t^{(h)}\right] \ \boldsymbol{o}_t^{(s)} &= Attention\left(\boldsymbol{q}_t^{(s)}, \boldsymbol{k}{\leq t}^{(s)} ,\boldsymbol{v}_{\leq t}^{(s)}\right) \
\boldsymbol{q}_i^{(s)} &= \boldsymbol{x}_i\boldsymbol{W}_q^{(s)} \
\boldsymbol{k}_i^{(s)} &= \boldsymbol{x}_i\boldsymbol{W}_k^{(s)} \
\boldsymbol{v}_i^{(s)} &= \boldsymbol{x}_i\boldsymbol{W}_v^{(s)}
\end{aligned}
\]
其中，$\boldsymbol{q}_i^{(s)}$，$\boldsymbol{k}_i^{(s)}$，$\boldsymbol{v}_i^{(s)}$分别表示查询、键和值向量，它们是通过与不同的权重矩阵相乘得到的。
在实际应用中，我们经常设置 $d_k = d_v = d / h$，即将向量的维度平均分配到每个头。例如，在 LLAMA2-7b 模型中，$d=4096$，$h=32$，所以 $d_k = d_v = 128$。
KV 缓存的重要性
在自回归语言模型中，我们可以缓存已计算的 $\boldsymbol{k}$和 $\boldsymbol{v}$值以供后续使用，这就是所谓的 KV Cache 。这可以避免重复计算，提高推理效率。然而，KV Cache 的大小对 GPU 显存提出了很高的要求。因此，如何减少 KV Cache 的大小，同时尽可能保证模型效果，成为了一个重要的研究方向。
MQA, GQA 到 MLA 的演变
为了解决 KV Cache 的问题，研究人员提出了 MQA（Multi-Query Attention） 、 GQA（Grouped Query Attention） 等改进方法。 MQA 通过共享查询向量来减少计算量，而 GQA 通过分组的方式降低 KV Cache 的存储需求。这些改进在一定程度上解决了 KV Cache 的问题，但仍有优化空间。
MLA（Multi-head Latent Attention） 是最新的改进，它在 GQA 的基础上进一步优化。 MLA 的设计思路是通过引入潜在变量来更高效地管理注意力机制，从而在减少 KV Cache 大小的同时，依然保持甚至提升模型的性能。
结语
从 MHA 到 MLA，注意力机制不断发展，每一步的改进都在努力平衡计算效率和模型效果。 DeepSeek-V2 的 MLA 技术展示了在这一领域的最新进展，预示着未来大规模语言模型在性能和成本上的进一步突破。希望这篇文章能帮助大家更好地理解这些技术背后的原理和演变过程。
根据 DeepSeek-V2 论文 （https://arxiv.org/pdf/2405.04434），多头潜在注意力 （Multi-head Latent Attention, MLA） 机制旨在提高推理效率，同时保持高性能。这种机制通过低秩键值联合压缩 （Low-rank Key-Value Joint Compression） 显著减少了 KV 缓存，从而在推理过程中提高了效率。以下是对 MLA 机制的详细
1. 预备知识：标准多头注意力 （MHA）
在解释 MLA 之前，需要理解标准多头注意力 （Multi-Head Attention, MHA） 。 MHA 通过并行计算多个注意力头来捕捉不同的注意力模式，每个注意力头都有独立的查询 （Query） 、键 （Key） 和值 （Value） 向量。具体步骤如下：
- 输入分割：输入序列被分割成多个头，每个头对应不同的查询、键和值向量。
- 序列处理：每个头独立处理其对应的查询、键和值向量，通过点积计算注意力权重，并加权求和得到每个头的输出。
- 头合并：所有头的输出被拼接并通过线性变换生成最终输出。
MHA 的主要问题在于 KV 缓存的大小，尤其是在长序列推理时，这会显著降低效率。
2. 低秩键值联合压缩
MLA 通过低秩键值联合压缩来解决上述问题。具体步骤如下：
- 键值压缩：MLA 将键和值向量联合压缩成一个低秩潜在向量，以减少缓存的存储需求。具体实现是通过矩阵分解技术，将高维的键和值向量映射到低维空间。
- 潜在表示：在低维空间中，键和值向量被表示为潜在向量，这些潜在向量可以在推理阶段快速检索，从而减少计算开销。
3. 解耦旋转位置嵌入
为了进一步提高效率，MLA 引入了解耦旋转位置嵌入 （Decoupled Rotary Position Embedding, RoPE） 。这项技术允许在不同注意力头之间共享位置信息，从而减少了计算复杂度。
4. KV 缓存比较
MLA 与传统的 MHA 相比，显著减少了 KV 缓存的大小。论文中提到，MLA 能够在保持高性能的同时，将 KV 缓存减少 93.3% 。这一改进在长序列推理任务中尤为显著，因为 KV 缓存的减小直接提高了推理速度和效率。
结论
MLA 机制通过低秩键值联合压缩和解耦旋转位置嵌入，在减少 KV 缓存的同时，显著提高了推理效率。与传统的 MHA 相比，MLA 在不牺牲性能的前提下，实现了更高的推理效率。这使得 DeepSeek-V2 能够在保持高性能的同时，显著降低推理成本。
希望以上解析能帮助你更好地理解 MLA 机制。如果有更多问题，欢迎随时询问。
在 MLA（Multi-head Latent Attention） 中，通过低秩键值联合压缩实现了将高维的键 （Key） 和值 （Value） 向量压缩成低维的潜在向量，从而保留重要信息并减少存储和处理的数据量。以下是具体实现的步骤和方法：
1. 低秩矩阵分解
低秩矩阵分解是实现压缩的核心技术。它通过将高维矩阵分解为两个或多个低维矩阵的乘积来减少数据的维度。具体来说，假设我们有一个高维的键矩阵 ( K. 和一个值矩阵 ( V )，我们可以使用低秩矩阵分解将它们压缩为低维矩阵。✅
2. 键和值矩阵的联合压缩
在 MLA 中，键和值矩阵 ( K. 和 ( V ) 被联合压缩成潜在向量。这个过程可以通过以下步骤实现：✅
- 初始矩阵：
- ( K \in \mathbb{R}^{n \times d_k} ) ：键矩阵，维度为 ( n \times d_k ) 。
- ( V \in \mathbb{R}^{n \times d_v} ) ：值矩阵，维度为 ( n \times d_v ) 。
- 低秩分解：
- 将键矩阵 ( K. 和值矩阵 ( V ) 分解成低秩矩阵。假设我们使用秩 ( r ) 进行分解：✅
- ( K \approx K_L K_R. ，其中 ( K_L \in \mathbb{R}^{n \times r} )，( K_R \in \mathbb{R}^{r \times d_k} ) 。✅
- ( V \approx V_L V_R. ，其中 ( V_L \in \mathbb{R}^{n \times r} )，( V_R \in \mathbb{R}^{r \times d_v} ) 。✅
- 联合压缩：
- 将键和值矩阵联合表示为潜在向量 ( Z. ：✅
- ( Z = K_L = V_L \in \mathbb{R}^{n \times r} ) 。
通过这种方式，键和值矩阵被压缩到相同的低维潜在空间中。
3. 潜在向量的使用
在推理过程中，潜在向量 ( Z. 被用于计算注意力权重和输出：✅
- 注意力权重计算：
- 使用低维潜在向量 ( Z. 来计算注意力权重，而不是直接使用高维的键和值向量。这可以通过点积计算或其他注意力机制实现。✅
- 加权求和：
- 将计算得到的注意力权重应用于值矩阵 ( V. 的低秩表示 ( V_R )，并进行加权求和，得到最终的输出。✅
4. 优势
- 减少存储需求：通过压缩键和值矩阵，显著减少了 KV 缓存的大小，从而降低了存储需求。
- 提高计算效率：低维的潜在向量使得计算注意力权重和输出更加高效，减少了计算复杂度。
总结
通过低秩键值联合压缩，MLA 成功地将高维的键和值向量压缩成低维的潜在向量。这些潜在向量在保留重要信息的同时，显著减少了需要存储和处理的数据量，从而提高了推理效率和性能。具体而言，低秩矩阵分解技术是实现这一压缩过程的关键，通过将高维矩阵分解为低维矩阵的乘积，达到了有效的压缩效果。