终于把 deepseek 中的多头潜在注意力机制搞懂了！！

MLA 的核心优化方式：对 KV（Key & Value）进行低秩压缩，只存储一个小的潜在向量，在需要时再恢复 Key 和 Value，极大地减少 KV 缓存大小。对 Query 进行低秩压缩，虽然不会减少 KV 缓存，但能够减少训练过程中的激活存储，提高计算效率。这样可以在不显著降低模型性能的情况下，极大减少 Transformer 模型在推理和训练时的内存占用，提高推理速度！解耦旋转位置嵌入

Python_金钱豹

1421人浏览 · 2025-02-27 10:48:39

Python_金钱豹 · 2025-02-27 10:48:39 发布

今天给大家介绍 deepseek 中的一个关键创新，多头潜在注意力机制

尽管多头自注意力机制在 Transformer 中取得了很好的效果，但其计算和存储成本较高。因此，研究者提出了多种改进方法来优化多头注意力机制，包括多查询注意力（MQA）和分组查询注意力（GQA），它们主要用于优化推理速度和计算成本。

多头潜在注意力（Multi-Head Latent Attention，MLA）是 DeepSeek 模型中引入的一种创新注意力机制，旨在优化传统多头注意力（Multi-Head Attention，MHA）的计算效率和内存占用。MLA 通过对键（Key）和值（Value）进行低秩联合压缩，显著减少了推理过程中的键值缓存（KV Cache），在保持模型性能的同时降低了内存占用。

多头潜在注意力机制（MLA）

多头潜在注意力 (MLA) 是一种新的注意力机制，旨在解决 MHA 的内存问题。

它通过对键和值进行低秩联合压缩，这可以减小 KV 缓存的大小（减少了计算和存储需求），同时保持甚至提高性能。

MLA 引入了两项关键创新：

低秩键值压缩
解耦旋转位置嵌入 (RoPE)

低秩键值联合压缩

MLA 的关键是对键（Key）和值（Value）进行低秩联合压缩，以减少 KV 缓存的存储需求。

具体方法如下：

键值的低秩压缩

我们不直接存储原始的 Key 和 Value，而是首先将隐藏状态投影到一个更小的压缩潜在向量

其中：

然后，我们可以通过以下方式从重新计算 Key 和 Value：

其中：

这样，在推理时，我们只需缓存而不是完整的 Key 和 Value，因此 KV 缓存的存储需求仅为，其中 l 是 Transformer 层数。

是上投影（up-projection）矩阵，用于恢复 Key 和 Value。
是压缩后的 KV 表示，维度远小于原始的 Key-Value 维度，即存储需求大大降低。
是降维投影矩阵，用于将隐藏状态映射到较小的维度。

查询（Query）的低秩压缩

MLA 还对查询（Query）进行了低秩压缩，尽管这不会直接减少 KV 缓存的大小。

查询的压缩方式类似于 KV：

其中：

尽管 Query 的压缩不会减少 KV 缓存的大小，但它可以减少训练时的激活存储需求，从而降低计算成本。

是压缩后的查询向量，远小于原始查询维度。
是查询的降维投影矩阵。
是查询的上投影矩阵。

总结

MLA 的核心优化方式：

对 KV（Key & Value）进行低秩压缩，只存储一个小的潜在向量，在需要时再恢复 Key 和 Value，极大地减少 KV 缓存大小。
对 Query 进行低秩压缩，虽然不会减少 KV 缓存，但能够减少训练过程中的激活存储，提高计算效率。

这样可以在不显著降低模型性能的情况下，极大减少 Transformer 模型在推理和训练时的内存占用，提高推理速度！

解耦旋转位置嵌入 (RoPE)

为了在保证内存效率的同时高效地结合旋转位置编码（RoPE），MLA 采用了一种解耦 RoPE 策略。

传统 RoPE 简介

旋转位置编码（RoPE）是一种高效的相对位置编码方法。它通过将 token 的位置信息直接嵌入查询（Q）和键（K）向量，利用旋转变换在点积注意力计算时自然地保留相对位置信息。

传统位置编码需要额外加法运算，而 RoPE 通过旋转变换，将位置信息隐式地融合到 Q 和 K 向量的点积计算中。
具体而言，对于任意 d 维的向量，RoPE 将每对相邻的两个维度视为一个二维平面，并在该平面内执行旋转变换

其中，旋转角度与 token 位置 p 相关，通常设定为：

使得高维度的旋转角度变化较小，而低维度变化较大，从而在不同尺度上编码位置信息。

传统 RoPE 存在的问题

尽管 RoPE 可以有效地嵌入相对位置信息，并广泛应用于 LLM（如 LLaMA、GPT-4 等），但它存在计算和存储上的挑战。

计算开销

传统 RoPE 需要在每个注意力头的查询和键上分别执行旋转变换
由于 Transformer 具有多个注意力头，对于大规模模型，这种额外的计算会导致推理效率下降

内存占用

在推理过程中，RoPE 位置编码的结果通常需要存储在 KV 缓存中。
对于长序列输入（如 100K+ tokens），KV 缓存的大小会随着序列长度线性增长，占用大量 GPU 显存。

解耦 RoPE的核心思想

解耦 RoPE 旨在优化 RoPE 的计算和存储效率，同时保留其编码能力。

它的主要改进点如下：

额外引入“解耦查询”

在标准 MHA 计算中，每个注意力头的查询和键都需要执行 RoPE 变换：

解耦 RoPE 的创新点在于，将查询拆分为两个部分：

具体计算公式：

其中：

和是新的投影矩阵，专门用于生成解耦的查询和键。
和由普通部分与解耦部分组合
普通查询：不经过 RoPE 变换，代表非位置敏感的特征信息。
解耦查询：专门用于嵌入 RoPE 位置编码信息。

共享 RoPE 变换的 Key

在标准 RoPE 中，每个注意力头都需要独立计算。

解耦 RoPE 采用共享 Key 策略，即所有注意力头共用一个旋转变换后的，从而减少计算开销：

这样做的好处是：

减少计算量：不必为每个 Query 和 Key 计算 RoPE 变换。
减小 KV 缓存大小：只需存储共享的，从而降低 GPU 内存占用。

解耦 RoPE 的计算流程

计算普通查询和普通键
计算解耦查询和共享键
拼接查询和键
计算注意力分数

解耦 RoPE优点

拆分查询（Query）为普通查询和解耦查询，使得部分信息不受 RoPE 影响。
共享 Key 的 RoPE 变换，减少计算量和 KV 缓存大小。
降低 GPU 内存占用，提高推理速度，特别适用于长序列任务和大规模 Transformer。

MLA 如何在推理过程中减少计算

MLA 中的一个关键优化是权重矩阵的吸收。

在 MLA 中，为了减少计算复杂度并提升推理效率，模型采用了一种优化策略：将上投影矩阵吸收到其他计算中，以避免显式重构 Key 和 Value。

主要有两个关键优化：

1.吸收到

传统方法需要显式计算 Key 向量，然后计算 Query 和 Key 之间的注意力分数。
优化方式：MLA 直接将上投影矩阵吸收到里，这样就可以简化查询计算，并优化注意力分数的计算。

查询计算

Query 计算方式如下：

其中：

是查询的压缩潜在向量
是查询的上投影矩阵

注意力分数计算

传统计算方式：
优化方式（矩阵结合律）

这里，是一个新的合并矩阵，它直接将吸收到中，从而避免显式计算 Key 向量。

优化效果

减少计算步骤：不再需要独立计算
提升计算效率：减少矩阵运算的开销，使推理更快。

2.吸收到

传统方法需要先计算 Value 向量，然后再进行注意力计算并投影到最终的输出层。
优化方式：直接将吸收到里，简化最终的输出计算。

注意力输出计算

传统方法

其中，是通过上投影重构的 Value。

最终输出计算

代入

优化方式

这里，是一个新的合并矩阵，它直接将吸收到中，减少额外的计算步骤。

优化效果

避免显式计算：减少计算复杂度，提高推理效率。
减少矩阵运算的存储需求：降低显存占用，使模型更适用于大规模推理任务。

关键要点

内存效率：较小的 KV 缓存意味着模型可以处理更长的序列或更大的批量大小，而不会耗尽内存。
推理速度：存储和检索的数据更少，模型可以更快地生成文本。
可扩展性：MLA 使得在内存有限的设备（如手机或边缘设备）上部署大型transformer模型成为可能。
性能提升：尽管进行了内存和计算优化，MLA 的性能仍优于 MHA，证明了效率和有效性可以并驾齐驱。

如何学习大模型 AI ？

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

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

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

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

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

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

在这里插入图片描述