DeepSeek-V3架构解析:多头潜在注意(MLA)技术详解

DeepSeek-V3创新:多头潜在注意(MLA)通过压缩注意力输入并解耦RoPE,在内存效率和建模能力间取得突破性平衡,甚至超越传统MHA。

原文标题:独家|DeepSeek-V3解释1:多头潜在注意

原文作者:数据派THU

冷月清谈:

本文是“DeepSeek-V3解释”系列的首篇文章,深入探讨了DeepSeek-V3中的一项关键架构创新——多头潜在注意(MLA)。文章首先回顾了标准多头注意力(MHA)机制,解释了推理阶段中键值缓存的需求及其局限性,并对比了多队列注意(MQA)和群队列注意(GQA)等优化方法。接着,详细介绍了MLA的核心逻辑,包括如何通过将注意力输入压缩到低维潜在向量来减少内存占用,以及如何通过解耦RoPE(旋转位置嵌入)来解决MLA与RoPE的不兼容问题。文章还提供了MLA的具体算法流程图,并展示了MLA在内存效率和建模能力之间的平衡,以及它在性能上与MHA、GQA和MQA的对比结果,强调MLA在某些情况下甚至超越了原始MHA的性能。

怜星夜思:

1、DeepSeek-V3中引入的MLA(多头潜在注意)在实际应用中,除了文章中提到的内存效率提升和建模能力增强之外,还可能带来哪些潜在的优势或改进?比如在模型训练速度、泛化能力或者对于不同类型数据的适应性方面?
2、文章提到RoPE(旋转位置嵌入)与MLA存在不兼容的问题,并通过“解耦的RoPE”来解决。那么,除了这种解耦方式,是否还存在其他可能的解决方案?或者,这种不兼容性是否暗示了RoPE本身存在某种局限性?
3、文章提到MLA在某些情况下建模能力甚至超过了最初的MHA。那么,你认为MLA在哪些具体的应用场景下,相比MHA、MQA和GQA等其他注意力机制,能够更好地发挥其优势?

原文内容

作者:Shirley Li
翻译:陈超
校对:吴振东

本文约3100字,建议阅读6分钟

本文对DeepSeek-V3的主要架构创新——多头潜在注意(MLA)进行了简要介绍。

实现更快的推理速度的DeepSeek-V2和DeepSeek-V3背后的关键架构创新。

图片由作者使用ChatGPT制作

这是我们新系列“DeepSeek-V3解释”的第一篇文章,我们将尝试阐明DeepSeek-V3[1,2],DeepSeek最近的开源模型。

在本系列中,我们致力于两个主要话题:

  • DeepSeek-V3的主要架构创新,包括MLA(多头潜在注意)[3],DeepSeekMoE[4],辅助-无损失负载平衡[5],以及多标记预测训练。
  • DeepSeek-V3训练,包括预训练,微调和RL(强化学习)校准阶段。

本文主要关注多头潜在注意(MHLA),在DeepSeek-V2开发阶段首次提出,DeepSeek-V3中使用。   

内容列表:

  • 背景:我们从标准MHA开始,解释我们在推理阶段为什么需要键值缓存,如何尝试优化MQA(多队列注意)和GQA(群队列注意)如何尝试优化,以及RoPE如何工作等等。
  • 多头潜在注意:在深入介绍MLA时,包括它的动机,为什么需要解耦的RoPE,及其表现。
  • 参考文献。
           
背景

为了更好地的地理解MLA并保持使本文的完整性,我们将在深入MLA细节前先回顾几个相关概念。

仅解码Transformer中的MHA

MLA被开发用于提高自回归文本生成的推理速度,我们将在本文探讨的MHA是用于仅解码Transformer。

下图比较了三种用于解码的Transformer架构,(a)表明编码和解码在最初的“注意力就是你所需要的”论文中提出。它的解码部分之后被[6]简化,催生(b)中所示的只有仅解码的Transformer模型,后来常用于众多像GPT的生成式模型[8]。

现在,大语言模型为了更稳定的训练常选择(c)中的结构,将标准化应用于输入而不是输出,LayerNorm升级为RMS Norm。这将作为我们将在本文中讨论的基准体系结构。   

图1. Transformer架构. (a) [6]中提到的编码-解码. (b) [7] 中提出以及GPT中使用的仅解码Transformer [8]. (c) 一个优化版本的(b)以及注意前的RMS 规范 [3]

在本文中,MHA计算很大程度上遵循[6]中的过程,如下图所示:

图2. 缩放点积注意vs.多头注意。图片来自[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的投影矩阵:

图片

注意到等式(1)-(8)所描述的过程仅仅是单一符号的。在推理过程中,我们需要对每一个新生成的符号重复这一过程,包括大量的重复计算。这催生了一种技术叫做键值缓存。

键值缓存

正如它的名字所示,键值缓存是一种用于加速自回归进程的技术,通过缓冲和重用以前的键和值,而非在每一次解码阶段重新进行计算。

注意到键值缓存通常仅用于推理阶段,在训练阶段我们仍然需要同时加工所有的输入序列。

键值缓存通常作为旋转缓冲器。在每个解码阶段,只有新队列Q被计算,而存储在缓冲器当中的K和V将会被重用,所以使用新的Q和重用的K、V来计算注意。同时,新的符号K和V也将在后续使用中被添加至缓冲器。

然而,由键值缓存实现的加速是以内存损失为代价的,键值缓存通常对批容量×序列长度×隐藏容量×头数进行缩放,如果有更大的批容量或者更长的序列将导致内存瓶颈。

这进一步催生了两个技术致力于解决这一局限:多队列注意和群队列注意。

多队列注意(MQA)vs群队列注意(GQA)

下图展示了原始MHA,群队列注意(GQA)[10]以及多队列注意(MQA)[9]的比较。

图3.MHA[6],GQA[10]以及MQA[9]. 图片来自[10].

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

然后,我们应用旋转矩阵来对每个对进行旋转:

图4. 旋转矩阵应用于一对符号的说明。图片来自作者

θ = θ(p) = p ⋅ θ_0, and θ_0 是一个基础频率. 在我们的4维例子中,这意味着(x_1, x_2) 将由θ_0 旋转,(x_3, x_4) 将由2 ⋅ θ_0旋转。

这就是为什么我们说这个旋转矩阵是位置依赖:在每个位置(或者每个对),我们都将使用不同的旋转矩阵,旋转角度由位置决定。   

RoPE在现代大语言模型中广泛使用,因为它在编码长序列当中的效率高,但是正如我们从以上公式当中所看到的,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}都是将共享潜在向量映射回高维空间的上投影矩阵。


同样地,我们也可以将队列映射到一个潜在的,低维的向量,然后将其映射回原始的高维空间。   

图片

为什么需要解耦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}之间。

正如我们在背景部分所解释的,这个旋转矩阵是位置依赖的,这意味着每个位置的旋转矩阵是不同的。因此,W^{UK}不能再被W^Q吸收。

为了解决这个冲突,作者提出了他们所谓的“解耦的RoPE”,通过在共享键向量的同时引入额外的查询向量,并且只在RoPE过程中使用这些额外的向量,同时保持原始键与旋转矩阵的隔离。

MLA的全部过程可以被总结如下(等式数字重用了[3]中的附录C)    


图5.MLA过程。图像由作者根据[3]中的等式编辑。

  • 等式. (37) to (40)描述如何加工队列符号。
  • 等式. (41) and (42)描述如何加工键符号。
  • 等式. (43) and (44) 描述如何为RoPE使用额外的共享键,注意RoPE中不涉及(42)的输出。
  • 等式. (45) 描述如何加工键符号。

在此进程当中,仅有带框的蓝色变量需要被缓存。这一过程可以用流程图更好地说明: 
 

6. MLA流程图.图片来自[3]


MLA的表现

下表比较了键值缓存(每个符号)需要的元素数量以及MHA,GQA和MQA和MLA的建模能力,表明MLA在内存效率和建模能力之间实现了更好的平衡。

有趣的是,MLA的建模能力甚至超过了最初的MHA。

[3]中的表 1


更具体地说,下表展示了MHA,GQA和MQA在7B模型上的表现,MHA表现显著超过了MQA和GQA。   


[3]中的表8


[3]的作者也进行了MHA和MLA的对比分析,结果总结于下表,MLA整体结果更好。

[3]中的表 9


参考文献

[1] DeepSeek

[2] DeepSeek-V3 Technical Report

[3] DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model

[4] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models

[5] Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts

[6] Attention Is All You Need

[7] Generating Wikipedia by Summarizing Long Sequences

[8] Improving Language Understanding by Generative Pre-Training

[9] Fast Transformer Decoding: One Write-Head is All You Need

[10] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

[11] RoFormer: Enhanced Transformer with Rotary Position Embedding



原文标题:

DeepSeek-V3 Explained 1: Multi-head Latent Attention

原文链接:

https://towardsdatascience.com/deepseek-v3-explained-1-multi-head-latent-attention-ed6bee2a67c4/


编辑:王菁




译者简介




陈超,北京大学应用心理硕士,数据分析爱好者。本科曾混迹于计算机专业,后又在心理学的道路上不懈求索。在学习过程中越来越发现数据分析的应用范围之广,希望通过所学输出一些有意义的工作,很开心加入数据派大家庭,保持谦逊,保持渴望。

翻译组招募信息

工作内容:需要一颗细致的心,将选取好的外文文章翻译成流畅的中文。如果你是数据科学/统计学/计算机类的留学生,或在海外从事相关工作,或对自己外语水平有信心的朋友欢迎加入翻译小组。

你能得到:定期的翻译培训提高志愿者的翻译水平,提高对于数据科学前沿的认知,海外的朋友可以和国内技术应用发展保持联系,THU数据派产学研的背景为志愿者带来好的发展机遇。

其他福利:来自于名企的数据科学工作者,北大清华以及海外等名校学生他们都将成为你在翻译小组的伙伴。


点击文末“阅读原文”加入数据派团队~



转载须知

如需转载,请在开篇显著位置注明作者和出处(转自:数据派ID:DatapiTHU),并在文章结尾放置数据派醒目二维码。有原创标识文章,请发送【文章名称-待授权公众号名称及ID】至联系邮箱,申请白名单授权并按要求编辑。

发布后请将链接反馈至联系邮箱(见下方)。未经许可的转载以及改编者,我们将依法追究其法律责任。






关于我们

数据派THU作为数据科学类公众号,背靠清华大学大数据研究中心,分享前沿数据科学与大数据技术创新研究动态、持续传播数据科学知识,努力建设数据人才聚集平台、打造中国大数据最强集团军。



新浪微博:@数据派THU

微信视频号:数据派THU

今日头条:数据派THU

点击“阅读原文”拥抱组织


从理论上讲,MLA通过引入潜在空间,能够学习到更抽象、更本质的特征表示。因此,在一些需要进行复杂推理和语义理解的任务中,MLA可能表现更好。例如,在进行知识图谱推理、常识推理等任务时,MLA可能能够更好地捕捉实体之间的关系和语义信息,从而提高推理的准确性。此外,MLA的潜在空间还可以被用于进行跨模态学习,例如将图像和文本映射到同一个潜在空间中,从而实现图像描述、视觉问答等任务。

从略微学术的角度来看,MLA可以被视为一种特殊的低秩分解方法,它将原始的注意力矩阵分解为两个低秩矩阵的乘积。这种分解可以降低模型的参数量,从而减少过拟合的风险,并提高模型的泛化能力。此外,MLA中的潜在向量可以被视为一种全局上下文信息,它可以帮助模型更好地理解输入序列的整体语义,从而提高模型在各种自然语言处理任务上的性能。但是,MLA的有效性也取决于潜在空间的维度选择和训练策略,需要进行仔细的实验和调优。

我觉得这得看场景。如果对内存要求特别苛刻,那MLA肯定有优势,毕竟它更省内存。如果在长文本生成或者需要处理大量上下文信息的任务中,MLA可能也会表现更好,因为它能更有效地利用历史信息。但如果对精度要求极高,而且不差钱,那还是得仔细对比各种方案,跑benchmark才知道。

抛开技术细节,我觉得这个问题更深层地反映了模型设计的trade-off。RoPE在长序列处理上有优势,但MLA在内存效率上有亮点。为了追求整体性能的提升,我们需要在不同的模块之间找到平衡点,甚至可能需要牺牲某些模块的局部最优。解耦RoPE是一种妥协,但也许未来的研究会发现更优雅的解决方案。

个人认为,MLA在以下场景中可能更具优势:

1. 资源受限的设备: 例如移动设备、嵌入式系统等,这些设备通常内存较小,MLA的内存效率优势可以帮助模型在这些设备上运行。
2. 长文本处理: MLA通过潜在向量压缩信息,可以更有效地处理长文本,例如长篇小说生成、文档摘要等。
3. 在线学习: MLA的键值缓存机制可以加速推理过程,适用于需要实时响应的在线学习场景,例如对话机器人、在线翻译等。

当然,具体效果还需要根据具体的应用场景和数据进行实验验证。

这个问题很有意思!除了文章里提到的,我觉得MLA可能还会间接提升训练速度。因为降低了内存占用,就能支持更大的batch size,并行计算能力更强,自然训练就更快了。至于泛化能力,理论上MLA引入的潜在空间,也许能让模型学到更本质的特征,从而提高泛化性。但具体效果还得看实验数据说话。

从技术角度分析,RoPE的不兼容性源于其位置依赖的旋转矩阵与MLA中潜在向量的维度映射之间的冲突。除了解耦RoPE,还可以考虑以下替代方案:

1. 自适应旋转: RoPE的旋转角度是固定的,可以尝试引入自适应的旋转角度,使其能够根据MLA的潜在向量进行调整,从而实现更好的兼容性。
2. 混合位置编码: 可以将RoPE与其他位置编码方法(如绝对位置编码)结合使用,利用其他方法的优势来弥补RoPE的不足。
3. 简化RoPE: 可以对RoPE进行简化,例如减少旋转矩阵的参数量或者使用更简单的旋转方式,从而降低其复杂度和计算成本。

总的来说,RoPE本身并没有绝对的局限性,但其与MLA的结合需要更加精细的设计和优化。

我觉得这个问题问到了点子上!解耦RoPE确实是一种方案,但可能不是唯一的。比如,能不能尝试其他的位置编码方式?或者对RoPE进行一些魔改,让它能更好地和MLA融合。另外,RoPE的本质是用旋转来编码位置信息,这种方式在处理长序列时可能效率更高,但同时也可能引入一些固有的问题,比如旋转角度的选择、旋转频率的设置等等,这些都可能限制了它的应用场景。

同意楼上的观点,MLA降低了内存占用,这允许我们在相同的硬件条件下使用更大的batch size,从而加速训练过程。此外,通过引入潜在向量,MLA可能迫使模型学习更加紧凑和高效的表示,这在一定程度上可以看作是一种正则化手段,有助于提升模型的泛化能力。但是,MLA的引入也可能增加模型的复杂性,如果潜在空间的维度选择不当,可能会导致信息损失或者过拟合。因此,MLA在不同类型数据上的适应性仍然需要仔细评估和调整。