「DeepSeek-V3 技术解析」：多头潜在注意力机制（MLA）

最新推荐文章于 2025-03-28 07:00:42 发布

大模型研究院

最新推荐文章于 2025-03-28 07:00:42 发布

阅读量1.1k

点赞数 24

文章标签：开源 easyui bert 架构人工智能大模型 deepseek

本文链接：https://blog.youkuaiyun.com/l01011_/article/details/146333149

版权

1、技术背景

为了便于读者更好地理解 MLA 架构并保持本文的完整性，在深入探讨 MLA 的技术细节之前，我们将首先回顾几个相关技术概念。

1.1 Decoder-only Transformers 中的多头注意力机制

有一点需要特别说明，MLA 架构的研发初衷是加快自回归文本生成的推理效率，因此本文讨论的 MHA（多头注意力机制）特指 Decoder-only Transformer 架构。

图 1 对比了三种可用于解码的 Transformer 架构：(a) 展示的是《Attention is All You Need》论文提出的编码器-解码器架构。其解码器部分随后被文献 [6] 简化，形成如 (b) 所示的 decoder-only Transformer 架构，该架构被 GPT 等生成模型广泛采用[8]。

请添加图片描述
在这种情况下， MHA 计算大致遵循文献 [6] 中的流程，如下图所示：

请添加图片描述
图 2. 缩放点积注意力与多头注意力架构对比（图源文献[6]）假设模型包含 n_h 个注意力头，每个注意力头维度为 d_h，则拼接后总维度为 n_h·d_h。

对于具有 l 层的模型，若设某层第 t 个 token 的输入表示为维度 d 的向量 h_t，则需要通过线性映射矩阵（linear mapping matrices）将 h_t 的维度从 d 映射到（h_n - d_h）。

更形式化地表述如下（公式引用自文献[3]）：

请添加图片描述
其中，W^Q、WK 和 W^V 是线性映射矩阵：

请添加图片描述
完成映射后，q_t、k_t 和 v_t 将被分割为 n_h 个注意力头进行缩放点积注意力计算：

请添加图片描述
最终通过投影矩阵 W^O 将维度从 (h_n - d_h) 反向映射到 d：

请添加图片描述
需要强调的是，上述公式 (1)-(8) 描述的是单 token 处理流程。在推理过程中，每个新生成的 token 都需要重复这一计算流程，这产生了大量重复计算 —— 这正是催生键值缓存（Key-Value Cache）技术的关键动因。

1.2 Key-Value Cache

正如其名称所示，键值缓存（Key-Value Cache）是一种通过缓存和复用先前计算的键值对（Keys and Values）来加速自回归过程的技术，避免了在每个解码步骤重新计算这些参数的需求。

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

键值缓存通常以滚动缓冲区（rolling buffer）的形式实现。在每个解码步骤中，仅计算新的查询向量 Q（Query），而缓存中存储的 K（Keys）和 V（Values）会被复用，注意力机制将通过新计算的 Q 与复用的 K、V 进行运算。与此同时，新生成的 token 对应的 K 和 V 也会被追加到缓存中以供后续使用。

然而，键值缓存带来的加速是以内存消耗为代价的。由于键值缓存的大小通常与批处理量大小（batch size）× 序列长度（sequence length）× 隐藏层维度（hidden size）× 注意力头数（number of heads）成正比，当出现更大批处理量或更长序列时，极易形成内存瓶颈。

这种限制进一步催生了两项旨在突破该瓶颈的技术：多查询注意力机制（Multi-Query Attention）和分组查询注意力机制（Grouped-Query Attention）。

1.3 多查询注意力机制（MQA） vs 分组查询注意力机制（GQA）

下图展示了原始的多头注意力机制（MHA）、分组查询注意力机制（GQA）[10] 和多查询注意力机制（MQA）[9] 的对比。

请添加图片描述
MQA 的核心思想是让所有查询头（query heads）共享一个单独的键头（key head）和值头（value head）。这种做法能显著降低内存占用，但也会对注意力的计算精度产生负面影响。

GQA 可视为 MHA 与 MQA 的一种折中方案。在这种方法中，一组查询头仅共享一对键头和值头，而非所有查询头共享同一对。尽管如此，相较于原始的 MHA，其效果仍会有所逊色。

在后续章节中，我们将探讨 MLA 如何在内存效率与建模精度之间实现平衡。

1.4 旋转位置编码（RoPE）

最后需要提及的相关背景知识是旋转位置编码（RoPE）[11]。该方法通过在多头注意力机制中对查询向量（Query）和键向量（Key）施加基于正弦函数的旋转操作，将位置信息直接编码到注意力计算中。

具体而言，RoPE 会对每个词元（token）的查询向量和键向量应用一个与位置相关的旋转矩阵。该矩阵以正弦和余弦函数为基础，但以一种独特的方式应用它们来实现旋转。

为理解其“位置相关（position-dependent）”特性，我们以一个 4 维的嵌入向量（x₁, x₂, x₃, x₄）为例。

要应用 RoPE，首先需要将连续的维度两两配对：

(x₁, x₂) → 位置1
(x₃, x₄) → 位置2

然后，对其中的每一对应用旋转矩阵：

请添加图片描述

其中旋转角度 θ = θ§ = p ⋅ θ₀（θ₀为基频参数）。在此示例中，(x₁, x₂) 会被旋转 θ₀ 角度，而 (x₃, x₄) 则被旋转 2 · θ₀ 角度。

因此，我们称这种旋转矩阵为“位置相关（position-dependent）”：在每个位置（或每对），将应用不同的旋转矩阵，其中旋转角度由位置决定。

RoPE 因其在长序列编码中的高效性被现代大语言模型广泛采用。然而，从上述公式可见，该方法对查询向量和键向量的位置敏感性，使其在部分场景下与 MLA 机制不兼容。

2、Multi-head Latent Attention

最后我们终于可以进入 MLA 部分的讲解。在这一章节中，我们将首先阐述 MLA 的核心设计思想，然后深入探讨其为何需要修改 RoPE，最后展示 MLA 的具体算法实现及其性能表现。

2.1 MLA：核心设计思想

MLA 的核心设计思想是将注意力机制的输入 h_t 压缩为一个低维的潜在向量（维度 d_c，且 d_c 远小于原始维度（h_n·d_h））。当需要计算注意力时，我们可以将这个潜在向量重新映射回高维空间，以还原键向量和值向量。通过这种方式，只需存储潜在向量即可，从而实现显存占用的大幅降低。

这个过程可以通过以下公式更正式地进行描述。其中 c^{KV}_t 表示潜在向量，W^{DKV} 是压缩矩阵（上标 D 代表"下投影"，即降维操作），负责将 h_t 的维度从（h_n·d_h）压缩到d_c；而 W^{UK} 和 W^{UV} 则是上投影矩阵，负责将共享的潜在向量映射回高维空间。

请添加图片描述
类似地，我们也可以将查询向量映射到一个潜在的低维向量，然后再将其映射回原始的高维空间：

请添加图片描述

2.2 为何需要解耦式 RoPE

正如前文所述，RoPE 是训练生成模型处理长序列的常用位置编码方案。如果我们直接应用上述 MLA 策略，将会与 RoPE 不兼容。

为了更清晰地理解这一点，来看看使用公式 (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} 之间。

正如技术背景部分所述，这个旋转矩阵是位置相关（position-dependent）的，即每个位置对应的旋转矩阵都不同。因此，W^{UK} 无法再被 W^Q 吸收。

为解决这一矛盾，作者提出了“解耦式RoPE”方案：通过引入额外的查询向量和一个共享的键向量，并仅在 RoPE 过程中使用这些新增向量，同时保持原始键向量与旋转矩阵的隔离。

完整的 MLA 流程可总结如下（公式编号沿用文献 [3] 附录 C 的编排）：

请添加图片描述

其中：

公式 (37)-(40) 描述查询向量的处理流程
公式 (41)-(42) 描述键向量的处理流程
公式 (43)-(44) 说明如何将新增的共享键向量用于 RoPE，需注意公式 (42) 的输出不参与 RoPE 计算
公式 (45) 描述值向量的处理流程

在此过程中，仅需缓存蓝色变量。该流程可通过下图更直观地展示：

请添加图片描述

2.3 MLA 的性能表现

下表对比了 MHA、GQA、MQA 与 MLA 在 KV 缓存（每个 token）所需元素数量及建模能力上的差异，证明 MLA 确实能在显存效率与建模能力间实现更优平衡。

有趣的是，MLA 的建模能力甚至超越了原始 MHA 架构。

请添加图片描述
下表展示了 7B 参数规模下 MHA、GQA 与 MQA 的性能表现，其中 MHA 显著优于 MQA 和 GQA。

请添加图片描述
文献 [3] 的作者还对 MHA 与 MLA 进行了对比分析，结果如下表所示。数据显示 MLA 在整体指标上表现更优。

请添加图片描述

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

这份完整版的大模型 AI 学习资料已经上传优快云，朋友们如果需要可以微信扫描下方优快云官方认证二维码免费领取【保证100%免费】