DeepSeek-V3详解系列文章之：多头潜在注意力机制 (MLA)
本文是“DeepSeek-V3详解”系列的第一篇文章，我们将深入探讨DeepSeek最新开源模型DeepSeek-V3 [1, 2]。
本系列文章将涵盖两个主要主题：
本文主要关注多头潜在注意力机制 (MLA)，该机制最初在DeepSeek-V2的开发中提出，并在DeepSeek-V3中得到应用。
目录：
为了更好地理解MLA，并使本文自成一体，在深入探讨MLA细节之前，我们将回顾本节中几个相关的概念。
请注意，MLA的开发是为了加快自回归文本生成的推理速度，因此在此上下文中讨论的MHA是针对仅解码器的Transformer。
下图比较了三种用于解码的Transformer架构，其中(a)显示了原始“Attention is All You Need”论文中提出的编码器和解码器。其解码器部分随后被[6]简化，得到(b)所示的仅解码器的Transformer模型，该模型后来被许多生成模型（如GPT [8]）使用。
如今，大型语言模型更常选择(c)所示的结构以实现更稳定的训练，在输入而不是输出上应用归一化，并将LayerNorm升级为RMS Norm。这将作为我们本文中讨论的基线架构。
在此上下文中，MHA计算在很大程度上遵循[6]中的过程，如下图所示：
假设我们有n_h个注意力头，每个注意力头的维度表示为d_h，因此连接后的维度将为(h_n · d_h)。
对于具有l层的模型，如果我们将该层中第t个标记的输入表示为维度为d的h_t，则需要使用线性映射矩阵将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也将添加到缓存中以供以后使用。
然而，键值缓存带来的加速是以内存为代价的，因为键值缓存通常随批大小 × 序列长度 × 隐藏大小 × 头数而变化，当我们有更大的批大小或更长的序列时，这会导致内存瓶颈。
这进一步导致了两种旨在解决此限制的技术：多查询注意力和分组查询注意力。
下图显示了原始MHA、分组查询注意力 (GQA) [10]和多查询注意力 (MQA) [9]之间的比较。
MQA的基本思想是跨所有查询头共享单个键和单个值头，这可以显著减少内存使用，但也将影响注意力的准确性。
GQA可以看作是MHA和MQA之间的一种插值方法，其中只有一对键和值头将仅由一组查询头共享，而不是所有查询。但这仍然会导致比MHA更差的结果。
在后面的章节中，我们将看到MLA如何设法在内存效率和建模精度之间取得平衡。
我们需要提到的最后一个背景知识是RoPE [11]，它通过使用正弦函数旋转多头注意力中的查询和键向量来直接将位置信息编码到注意力机制中。
更具体地说，RoPE对每个标记的查询和键向量应用位置相关的旋转矩阵，并使用正弦和余弦函数作为其基，但以独特的方式应用它们来实现旋转。
要了解是什么使其成为位置相关的，请考虑一个只有4个元素的玩具嵌入向量，即(x_1, x_2, x_3, x_4)。
要应用RoPE，我们首先将连续的维度分组为对：
然后，我们应用旋转矩阵来旋转每一对：
其中θ = θ(p) = p ⋅ θ_0，而θ_0是基频。在我们的4维玩具示例中，这意味着(x_1, x_2)将旋转θ_0，而(x_3, x_4)将旋转2 ⋅ θ_0。
这就是为什么我们将此旋转矩阵称为位置相关的：在每个位置（或每一对）处，我们将应用一个不同的旋转矩阵，其中旋转角度由位置确定。
RoPE由于其在编码长序列方面的效率而被广泛用于现代大型语言模型，但正如我们从上述公式中看到的，它对Q和K都具有位置敏感性，这使得它在某些方面与MLA不兼容。
最后，我们可以继续讨论MLA部分。在本节中，我们将首先阐述MLA的高级思想，然后深入探讨为什么它需要修改RoPE。最后，我们还将介绍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}都是将共享潜在向量映射回高维空间的上投影矩阵。
类似地，我们还可以将查询映射到一个潜在的低维向量，然后将其映射回原始的高维空间：
### 为什么需要解耦RoPE
如前所述，RoPE是训练生成模型以处理长序列的常用选择。如果我们直接应用上述MLA策略，这将与RoPE不兼容。
为了更清楚地看到这一点，请考虑当我们使用Eqn. (7)计算注意力时会发生什么：当我们将转置的q与k相乘时，矩阵W^Q和W^{UK}将出现在中间，它们的组合相当于从d_c到d的单个映射维度。
在原始论文[3]中，作者将其描述为W^{UK}可以被“吸收”到W^Q中，因此我们不需要在缓存中存储W^{UK}，从而进一步减少内存使用。
但是，当我们考虑图(4)中的旋转矩阵时，情况并非如此，因为RoPE将在W^{UK}的左侧应用旋转矩阵，并且此旋转矩阵最终将位于转置的W^Q和W^{UK}之间。
正如我们在背景部分所解释的，此旋转矩阵是位置相关的，这意味着每个位置的旋转矩阵都不同。因此，W^{UK}**不能再被W^Q**吸收**。
为了解决这种冲突，作者提出了他们所谓的“解耦RoPE”，通过引入额外的查询向量以及共享键向量，并且仅在RoPE过程中使用这些额外向量，同时保持原始键与旋转矩阵某种程度上隔离。
整个MLA过程可以总结如下（方程编号从[3]的附录C中重复使用）：
其中
在此过程中，只需要缓存带有框的蓝色变量。此过程可以用下面的流程图更清晰地说明：
### MLA的性能
下表比较了KV缓存（每个标记）所需的元素数量以及MHA、GQA、MQA和MLA之间的建模能力，这表明MLA确实可以在内存效率与建模能力之间取得更好的平衡。
有趣的是，MLA的建模能力甚至超过了原始MHA。
更具体地说，下表显示了MHA、GQA和MQA在7B模型上的性能，其中MHA的性能明显优于MQA和GQA。
[3]的作者还对MHA与MLA进行了分析，结果总结在下表中，其中MLA总体上取得了更好的结果。
参考文献以上是DeepSeek-V3解释了1：多头潜在注意力的详细内容。更多信息请关注PHP中文网其他相关文章！
