解锁闪电般的推理：Deepseek-v3的多头潜在注意力革命！

Rifx.Online
Machine Learning , Natural Language Processing , Generative AI
05 Mar, 2025

DeepSeek-V3 解释 1：多头潜在注意力

DeepSeek-V2 和 DeepSeek-V3 背后的关键架构创新，以加快推理速度

这是我们新系列“DeepSeek-V3 解释”的第一篇文章，我们将试图揭开 DeepSeek-V3 [1, 2] 的神秘面纱，这是 DeepSeek 最新开源的模型。

在这个系列中，我们旨在涵盖两个主要主题：

DeepSeek-V3 中的主要架构创新，包括 MLA（多头潜在注意力）[3]、DeepSeekMoE [4]、无辅助损失负载均衡 [5] 和多标记预测训练。
DeepSeek-V3 的训练，包括预训练、微调和强化学习对齐阶段。

本文主要关注 多头潜在注意力，该概念最早在 DeepSeek-V2 的开发中提出，随后也应用于 DeepSeek-V3。

背景：我们从标准的 MHA 开始，解释为什么在推理阶段需要键值缓存，MQA 和 GQA 如何尝试优化它，以及 RoPE 的工作原理等。
多头潜在注意力：对 MLA 的深入介绍，包括其动机、为什么需要解耦的 RoPE 以及其性能。
参考文献。

背景

为了更好地理解 MLA，并使本文自成一体，我们将在本节中回顾几个相关概念，然后再深入探讨 MLA 的细节。

解码器仅 Transformer 中的 MHA

请注意，MLA 是为了加速自回归文本生成中的推理速度而开发的，因此我们在此上下文中讨论的 MHA 是针对仅解码器的 Transformer。

下图比较了用于解码的三种 Transformer 架构，其中 (a) 显示了在原始“注意力即一切”论文中提出的编码器和解码器。其解码器部分随后被 [6] 简化，形成了 (b) 中的仅解码器 Transformer 模型，该模型后来被许多生成模型如 GPT [8] 使用。

如今，LLM 更常选择 (c) 中所示的结构，以实现更稳定的训练，对输入应用归一化而不是输出，并将 LayerNorm 升级为 RMS Norm。这将作为我们在本文中讨论的基线架构。

在此上下文中，MHA 计算在很大程度上遵循 [6] 中的过程，如下图所示：

假设我们有 n_h 个注意力头，每个注意力头的维度表示为 d_h，因此连接的维度将是 (h_n · d_h)。

给定一个具有 l 层的模型，如果我们将该层中第 t 个标记的输入表示为 h_t，其维度为 d，我们需要使用线性映射矩阵将 h_t 的维度从 d 映射到 (h_n · d_h)。

更正式地，我们有（来自 [3] 的方程）：

其中 W^Q、W^K 和 W^V 是线性映射矩阵：

在这样的映射之后，q_t、k_t 和 v_t 将被分割成 n_h 个头，以计算缩放的点积注意力：

其中 W^O 是另一个投影矩阵，将维度从 (h_n · d_h) 反向映射到 d：

请注意，上述 Eqn.(1) 到 (8) 描述的过程仅针对单个标记。在推理过程中，我们需要对每个新生成的标记重复此过程，这涉及大量重复计算。这导致了一种称为键值缓存的技术。

键值缓存

顾名思义，键值缓存是一种旨在通过缓存和重用先前的键和值来加速自回归过程的技术，而不是在每个解码步骤中重新计算它们。

请注意，KV 缓存通常仅在推理阶段使用，因为在训练中我们仍需要并行处理整个输入序列。

KV 缓存通常实现为滚动缓冲区。在每个解码步骤中，仅计算新的查询 Q，而缓存中存储的 K 和 V 将被重用，因此注意力将使用新的 Q 和重用的 K、V 进行计算。同时，新标记的 K 和 V 也将被附加到缓存中以供后续使用。

然而，KV 缓存所带来的加速是以内存为代价的，因为 KV 缓存通常与 batch size × sequence length × hidden size × number of heads 成比例，导致当我们有更大的批量大小或更长的序列时出现内存瓶颈。

这进一步导致了两种旨在解决此限制的技术：多查询注意力和分组查询注意力。

多查询注意力 (MQA) 与分组查询注意力 (GQA)

下图显示了原始 MHA、分组查询注意力 (GQA) [10] 和多查询注意力 (MQA) [9] 之间的比较。

MQA 的基本思想是跨所有查询头共享一个单一的键和一个单一的值头，这可以显著减少内存使用，但也会影响注意力的准确性。

GQA 可以被视为 MHA 和 MQA 之间的一种插值方法，其中一对键和值头仅由一组查询头共享，而不是所有查询。但这仍然会导致与 MHA 相比的结果较差。

在后面的部分中，我们将看到 MLA 如何在内存效率和建模准确性之间寻求平衡。

RoPE（旋转位置嵌入）

我们需要提到的最后一个背景是 RoPE [11]，它通过使用正弦函数在多头注意力中直接将位置信息编码到注意力机制中，通过旋转查询和键向量来实现。

更具体地说，RoPE 在每个标记的查询和键向量上应用一个 位置依赖的旋转矩阵，并使用正弦和余弦函数作为其基，但以独特的方式应用以实现旋转。

为了了解什么使其位置依赖，考虑一个只有 4 个元素的玩具嵌入向量，即 (x_1, x_2, x_3, x_4)。

要应用 RoPE，我们首先将连续维度分组为对：

(x_1, x_2) -> 位置 1
(x_3, x_4) -> 位置 2

然后，我们应用一个旋转矩阵来旋转每一对：

其中 θ = θ(p) = p ⋅ θ_0，θ_0 是一个基频。在我们 4 维的玩具示例中，这意味着 (x_1, x_2) 将被旋转 θ_0，而 (x_3, x_4) 将被旋转 2 ⋅ θ_0。

这就是为什么我们称这个旋转矩阵为 位置依赖的：在每个位置（或每对）上，我们将应用不同的旋转矩阵，其中旋转角度由位置决定。

由于 RoPE 在编码长序列方面的效率，它在现代 LLM 中被广泛使用，但正如我们从上述公式中看到的，它对 Q 和 K 都是位置敏感的，这在某些方面使其与 MLA 不兼容。

多头潜在注意力

最后，我们可以进入 MLA 部分。在本节中，我们将首先阐述 MLA 的高层次思想，然后深入探讨为什么需要修改 RoPE。最后，我们展示 MLA 的详细算法及其性能。

MLA：高层次思想

MLA 的基本思想是将注意力输入 h_t 压缩为一个低维潜在向量，其维度为 d_c，其中 d_c 远低于原始的 (h_n · d_h)。稍后当我们需要计算注意力时，我们可以将这个潜在向量映射回高维空间以恢复键和值。因此，只需要存储潜在向量，从而显著减少内存。

这个过程可以用以下方程更正式地描述，其中 c^{KV}_t 是潜在向量，W^{DKV} 是将 h_t 的维度从 (h_n · d_h) 映射到 d_c 的压缩矩阵（这里上标 D 代表“下投影”，意味着压缩维度），而 W^{UK} 和 W^{UV} 都是将共享潜在向量映射回高维空间的上投影矩阵。