专栏名称: 数据派THU

本订阅号是“THU数据派”的姊妹账号，致力于传播大数据价值、培养数据思维。

一文通透DeepSeek V2——通俗理解多头潜在注意力MLA：改进MHA，从而压缩KV缓存，提高推理速度

数据派THU · 公众号 · 大数据 · 2025-03-16 17:00

正文

来源：七月在线实验室
本文约6000字，建议阅读10分钟
本文介绍DeepSeek LLM、DeepSeekMoE、DeepSeekMath。

前言

成就本文有以下三个因素：

24年5.17日，我在我司一课程「大模型与多模态论文100篇」里问道：大家希望我们还讲哪些论文一学员朋友小栗说：幻方旗下公司深度求索于24年5.7日发布的deepseek-v2
24年5.24日，我司一课程「大模型项目开发线上营1」里的一学员朋友问我：校长最近开始搞deepseek了吗？刚看了论文，没搞懂MLA那块的cache是怎么算的，我总觉得他的效果应该类似MQA才对，但是反馈是挺好的我当时回复他道：目前团队项目上的事情太多，然后近期在写那个KAN

确实还没来得及看这个deepseek，我近期看下。

我们在继英文层面的论文翻译、审稿、对话、idea提炼之后(其中的审稿和翻译已于24年8月底上线七月官网 )

打算再整一下中文层面的「硕士论文修订助手(不限学科，CS和非CS的都涵盖)」，预计24年9月份上线七月官网。

对于模型的选择，在闭源之外，在开源模型上有两个选择，一个chatglm4，一个DeepSeek(后来实际上线时，一开始用的deepseek v2，后来改成了deepseek v3)

而搞DeepSeek之前——近几天，会先写一下它的论文解读(当然，因为DeepSeek-V2从DeepSeek LLM、DeepSeekMoE迭代而来，且用到了DeepSeekMath中的GRPO算法，故第一部分会先讲DeepSeek LLM、DeepSeekMoE、DeepSeekMath)，故本文就来了，且DeepSeek也算证明了在国内也可以做出有效果、有影响力的创新。

且一如既往做到——对于几乎每一个主题，都如本博客万千读者或七月学员所说的：“还是看校长的文章好理解”，而其中的关键之一是做好图、文、公式的一一对应，不一笔带过、不自以为然，本文亦如此

同时本文也见证了自己从技术人到产品人定位的过渡

23上半年侧重原理，系统大量读paper，深挖原理。
23下半年侧重工程，和项目组不断优化各种工程问题。
24上半年侧重研究，横跨或综合多个领域(比如llm+机器人)，继续各种抠paper。
24下半年侧重产品，把世界级前沿落地成产品给用户使用，以发挥更大的价值和影响力(希望有机会早日达到世界级影响力)。

第一部分从DeepSeek LLM、DeepSeekMoE到DeepSeekMath

友情提醒，如不需要透彻深入理解，或者想直接看DeepSeek-V2的，可以直接跳到本文的第二部分，本文第二部分也是本文的最精华所在。

当然，如果你就是想先从本第一部分开始看则非常好，但里面的数学公式比较多，喜欢抠公式的可以细抠，不喜欢抠公式的则不用抠太细，不影响对DeepSeek-V2的整体理解。

第二部分 DeepSeek-V2：提出多头潜在注意力MLA且改进MoE

DeepSeek-V2属于DeepSeek的第二代版本，参数规模虽然达到了庞大的236B，但由于其MoE的结构，使得其中每个token激活仅21B的参数，且支持128K的上下文(It is equipped with a total of 236B parameters, of which 21B are activated for each token, and supports a context length of 128K tokens)

其对应论文为《DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model》，发布于24年5.7日。

他们首先在完整的预训练语料库上预训练DeepSeek-V2。
然后，收集了150万个对话会话，涵盖了数学、代码、写作、推理、安全等各个领域，以对DeepSeek-V2 Chat（SFT）进行监督微调（SFT）。
最后，他们遵循DeepSeekMath的方法，采用组相对策略优化（GRPO）进一步使模型与人类偏好对齐，并生成DeepSeek-V2 Chat（RL）。

DeepSeek-V2主要有两大创新点，其在Transformer架构「一个注意力模块和一个前馈网络(FFN)，如对transformer还不够熟练，请看此文：Transformer通俗笔记：从Word2Vec、Seq2Seq逐步理解到GPT、BERT」的基础上

1. 改造注意力模块

其通过创造性的提出多头潜在注意力：Multi-head Latent Attention(简称MLA），替代传统多头注意力(Multi Head Attention) 。

具体而言，MLA利用低秩键值联合压缩(low-rank key-value joint compression)来降低推理时的KV Cache开销——相当于low-rank joint compression for keys and values to reduce KV cache，且性能不输于MHA(论文中说的是性能比MHA还更好) 。

我个人认为，MLA本质上也是受到了LoRA和Stable Diffusion的启发「前者详见此文《LLM高效参数微调方法：从Prefix Tuning、Prompt Tuning、P-Tuning V1/V2到LoRA、QLoRA(含对模型量化的解释)》，后者详见此文《AI绘画原理解析：从CLIP、BLIP到DALLE、DALLE 2、DALLE 3、Stable Diffusion(含ControlNet详解)》」

所以MLA 是多头潜在注意力，在多头注意力上加上了低秩KV压缩。

2. 改造FFN

其把FFN的结构改成DeepseekMoE——是对传统MoE结构的改进(下图各种结构、表示很多，初看难以一看就懂，没事，下文会逐一详解，任何一个符号都不会放过——最后，反复琢磨之后，你也可以和我一样：脱离本文，手绘下图、手推下图背后的公式) 。

值得一提的是，他们还一块发布了 DeepSeek-V2-Lite，相当于配备 MLA 和 DeepSeekMoE 的较小模型，它总共有15.7B参数，其中每个token激活2.4B参数(we also release DeepSeek-V2-Lite, a smaller model equipped with MLA and DeepSeekMoE, for the open-source community. It has a total of 15.7B parameters, where 2.4B are activated for each token)

DeepSeek-V2-Lite 有 27 层，隐藏维度为 2048。它还采用 MLA，并具有 16 个注意力头，每个头的维度为 128 。

其 KV 压缩维度为 512，但与 DeepSeek-V2 略有不同，它不压缩查询。

对于解耦查询和键，每头维度为 64。

DeepSeek-V2-Lite 还采用 DeepSeekMoE，除了第一层外，所有前馈神经网络 (FFNs) 都被 MoE 层替换每个 MoE 层由 2 个共享专家和 64 个路由专家组成，每个专家的中间隐藏维度为 1408。在这些路由专家中，每个token将激活6个专家。

DeepSeek-V2-Lite 也在与 DeepSeek-V2 相同的预训练语料库上从头开始训练，该语料库未被任何 SFT 数据污染。

它使用 AdamW 优化器，超参数设置为 𝛽1 = 0.9， 𝛽2 =0.95，权重衰减 =0.1。

学习率使用预热和阶梯衰减策略进行调度最初，在前2000步期间，学习率从0线性增加到最大值。

随后，在训练了大约80%的tokens后，学习率乘以0.316，并在训练了大约90%的tokens后再次乘以0.316。

最大学习率设置为4.2 × 10−4，梯度裁剪范数设置为1.0 没有采用批量大小调度策略，而是以恒定的批量大小4608个序列进行训练在预训练期间，将最大序列长度设置为4K，并在5.7T tokens上训练DeepSeek-V2-Lite。

2.1 DeepSeek-V2提出MLA的背景与作用

2.1.1 KV Cache所导致的显存消耗大，需要尽可能降低

众所周知，KV Cache是大模型标配的推理加速功能——也是推理过程中，显存资源巨大开销的元凶之一。如下图所示，在模型推理时，KV Cache在显存占用量可达30%以上。

目前大部分针对KV Cache的优化工作：

1. 比如著名的vLLM「这是其介绍页面、这是其对应的GitHub、其论文则为：Efficient Memory Management for Large Language Model Serving with PagedAttention，当然了，我也写了一篇专门介绍vLLM的博客，详见《一文通透vLLM与其核心技术PagedAttention：减少KV Cache碎片、提高GPU显存利用率(推理加速利器)》」，其基于paged Attention，最大限度地利用碎片化显存空间，从而提升了空间利用率。

2. 再比如GQA、MQA。

GQA是query数不变，但多个query(比如2个)组成一个group以共享一个key value MQA则query也不变，但所有query(比如8个)共享一个key、一个value

至于更多，详见此文：一文通透各种注意力：从多头注意力MHA到分组查询注意力GQA、多查询注意力MQA。

这些方案的问题是什么呢？在于：

第一类方案并没有从根本上改变KV Cache占用空间巨大的问题。
而第二类方案中的MQA虽然较大降低了KV cache计算量，但性能相比MHA下降太多了至于第二类方案中的GQA则取了个折中：不好的是缓存下降的不够多、好的是相比MHA性能没有下降太多，毕竟我们追求的是缓存下降、性能不降。

那KV Cache到底是什呢？

1. 对此，我们先来回顾下transformer当中的注意力计算公式：

2. GPT预测下一个token时，其只能看到待预测token之前的所有token，故在最终生成Q1，Q2，Q3，Q4整个序列的过程中，会涉及到如下计算过程：

3. 然后把上面的softmax结果和对应的V值一相乘，便可得到：

$\begin{array}{l} \operatorname{Att}_{1}(Q, K, V)=\operatorname{softmaxed}\left(Q_{1} K_{1}^{T}\right) \overrightarrow{V_{1}} \\ \operatorname{Att}_{2}(Q, K, V)=\operatorname{softmaxed}\left(Q_{2} K_{1}^{T}\right) \overrightarrow{V_{1}}+\operatorname{softmaxed}\left(Q_{2} K_{2}^{T}\right) \overrightarrow{V_{2}} \\ \operatorname{Att}_{3}(Q, K, V)=\operatorname{softmaxed}\left(Q_{3} K_{1}^{T}\right) \overrightarrow{V_{1}}+\operatorname{softmaxed}\left(Q_{3} K_{2}^{T}\right) \overrightarrow{V_{2}} + \operatorname{softmaxed}\left(Q_{3} K_{3}^{T}\right) \overrightarrow{V_{3}} \\ \operatorname{Att}_{4}(Q, K, V)=\operatorname{softmaxed}\left(Q_{4} K_{1}^{T}\right) \overrightarrow{V_{1}}+\operatorname{softmaxed}\left(Q_{4} K_{2}^{T}\right) \overrightarrow{V_{2}} + \operatorname{softmaxed}\left(Q_{3} K_{3}^{T}\right) \overrightarrow{V_{3}} + \operatorname{softmaxed}\left(Q_{4} K_{4}^{T}\right) \overrightarrow{V_{4}} \end{array}$

可以很明显的看到，上述计算过程中，有不少的KV重复计算，比如：

上面这句话值得反复品味三遍！

如果序列长度越长，类似这样的KV重复计算会越多，从而势必将白白消耗那么大的显存，所以才说需要降低这种KV重复计算

2.1.2 Multi-head Latent Attent：致力于在推理中降低

MLA是对传统多头注意力做的改进，其目的有两个：首先是，降低推理过程中的KV Cache资源开销，其次，缓解MQA、MGA对性能的损耗。

如上文所说，KV Cache中，提到每一步都需要将K和V缓存下来。

具体而言，对于单个Attention Block块中的多头注意力(下图来自上文提到过的Transformer通俗笔记)。

比如，举个例子，假设：

(注意，是针对每个token，总之，如论文中所说，During inference, all keys and values need to be cached to accelerate inference, so MHA needs to cache elements for each token)

对Key和Value进行了一个低秩联合压缩(即Low-Rank Key-Value Joint Compression，通过低秩转换为一个压缩的KV，使得存储的KV的维度显著减小)

如上图所示(在MHA GQA中大量存在于keys values中的KV缓存——带阴影表示，到了MLA中时，只有一小部分的被压缩Compressed的Latent KV了) 那，MLA具体如何做压缩呢，详看下节。

2.2 详解MLA的两个部分：一部分做压缩、一部分做RoPE编码

2.2.1 MLA对Q K V的压缩：先对KV联合压缩后升维，再对Q压缩后升维

2.2.1.1 先对KV联合压缩(Low-Rank Key-Value Joint Compression)、后升维

$\begin{array}{l} \mathbf{c}_{t}^{K V}=W^{D K V} \mathbf{h}_{t} \\ \mathbf{k}_{t}^{C}=W^{U K} \mathbf{c}_{t}^{K V} \\ \mathbf{v}_{t}^{C}=W^{U V} \mathbf{c}_{t}^{K V} \end{array}$

可以看到针对KV先一块降维(公式表示时用D表示降维)，再K、V各自升维(公式表示用U表示升维)，具体而言

且在推理过程中

（即如论文中所说，during inference, since 𝑊𝑈𝐾 can be absorbed into 𝑊𝑄, and 𝑊𝑈𝑉 can be absorbedinto 𝑊𝑂, we even do not need to compute keys and values out for attention)

注，为方便大家更好的理解，我于24年8月27日晚上特地花了个把小时(反复琢磨、反复修改)，用手头的iPad Pro + apple pencil画了下上述过程，以方便大家一目了然的对比前后维度的变化(至于下图底部中“除了之外”的怎么回事，不急，下文很快会逐一阐述)

2.2.1.2 再对Q压缩降维、后升维

其次，对于上图左下角的

之前提到KV Cache中，Q的作用只发生在当下(预测下一个token时，其只能看到待预测token之前的所有token)，但是在模型训练的过程中，每个输入的token会通过多头注意力机制生成对应的query、key和value。这些中间数据的维度往往非常高，因此占用的内存量也相应很大。
所以论文中也提到为了降低训练过程中的激活内存activation memory，DeepSeek-V2还对queries进行低秩压缩——即便这并不能降低KV Cache，而其对Q的压缩方式和K、V一致，依然是先降维再升维。