本文约3100字,建议阅读6分钟
本文对DeepSeek-V3的主要架构创新——多头潜在注意(MLA)进行了简要介绍。
实现更快的推理速度的DeepSeek-V2和DeepSeek-V3背后的关键架构创新。
图片由作者使用ChatGPT制作
这是我们新系列“DeepSeek-V3解释”的第一篇文章,我们将尝试阐明DeepSeek-V3[1,2],DeepSeek最近的开源模型。
在本系列中,我们致力于两个主要话题:
本文主要关注多头潜在注意(MHLA),在DeepSeek-V2开发阶段首次提出,DeepSeek-V3中使用。
内容列表:
背景
为了更好地的地理解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