【AIGC魔童】DeepSeek核心创新技术(二):MLA
- 1. MLA框架的定义与背景
- 2. MLA框架的技术原理
- (1)低秩联合压缩
- (2)查询的低秩压缩
- (3)旋转位置嵌入(RoPE)
- 3. MLA框架的优势
- 4. MLA框架的核心价值
DeepSeek 的 MLA(Multi-head Latent Attention)框架凭借其独特的技术原理和显著优势,吸引了众多关注。下面将详细解读 MLA 框架。
1. MLA框架的定义与背景
DeepSeek 是一家专注于人工智能技术的公司,其开发的 MLA(Multi-Head Latent Attention)框架是 DeepSeek-V3 模型中用于高效推理的核心注意力机制。MLA 通过低秩联合压缩技术,减少了推理时的键值(KV)缓存,从而在保持性能的同时显著降低了内存占用。这一技术的出现,是为了应对传统 Transformer 模型在大规模语言模型(LLM)推理过程中面临的内存瓶颈问题。
在标准的 Transformer 模型中,多头注意力(Multi-Head Attention, MHA)机制通过并行计算多个注意力头来捕捉输入序列中的不同特征。每个注意力头都有自己的查询(Query, Q)、键(Key, K)和值(Value, V)矩阵,计算过程如下:
-
查询矩阵 Q:用于计算输入序列中每个位置的注意力权重。
-
键矩阵 K:用于与查询矩阵 Q 计算注意力分数。
-
值矩阵 V:用于根据注意力分数加权求和,得到最终的输出。
然而,这种机制在处理长序列时,会面临巨大的内存开销。例如,对于一个长度为 S 的序列,每个头的维度为 d ,则每个头的 KV 缓存大小为2 x S x d 。对于大规模模型,这会导致显存占用过高,限制了模型的推理效率。
为了解决这一问题,MLA 框架应运而生。它通过低秩联合压缩技术,将 KV 缓存的存储需求显著降低,同时保持了模型的性能。这一技术的核心在于,通过低秩分解和矩阵变换,将原本需要存储的大量 KV 值压缩为更小的维度,从而减少了显存的使用量。
2. MLA框架的技术原理
MLA 框架本质上是一种优化后的注意力机制。在理解它之前,我们先来简单了解一下什么是注意力机制。在大语言模型处理信息时,比如处理一段文本,它需要知道文本中哪些部分是重要的,哪些部分相对次要,注意力机制就像是模型的 “聚焦器”,帮助模型把重点放在关键信息上。而 MLA 框架则是在这个基础上,进一步优化,让模型在处理信息时更加高效。
(1)低秩联合压缩
- 核心思想:MLA 的一个关键技术是对注意力机制中的键(Key)和值(Value)进行低秩联合压缩。简单来说,就是把原本较大的数据量通过一定的方式变小,这样在推理的时候,需要缓存的键值(KV)对数量就会减少。
低秩联合压缩技术是 DeepSeek MLA 框架的核心,它通过将高维的键(Key)和值(Value)矩阵压缩到低维空间,从而显著减少存储需求。在传统的多头注意力机制中,每个头的键和值矩阵都需要单独存储,这在处理长序列时会导致巨大的内存开销。例如,对于一个长度为 S 的序列,每个头的维度为 d ,则每个头的 KV 缓存大小为2 x S x d 。对于大规模模型,这会导致显存占用过高,限制了模型的推理效率。
MLA 框架通过低秩联合压缩技术解决了这一问题。它首先将输入数据压缩到一个低秩空间,然后再通过上投影矩阵将其恢复到原始维度。这种压缩方式不仅减少了存储需求,还保持了模型的性能。具体来说,MLA 的低秩联合压缩过程如下:
低秩压缩:首先对输入进行低秩压缩,将维度为 d 的输入压缩到维度为 r(其中r << d ),通过一个低秩变换矩阵 Wr 实现:
其中, Wr是一个d x r 的矩阵,将输入压缩到低秩空间。
扩展维度:然后通过两个变换矩阵Wk和Wv ,将低秩的 Latent_KV 扩展回原始维度d ,得到每个头的 K 和 V:
其中,Wk和Wv是r x d的矩阵,用于将低秩表示恢复到原始维度。
计算注意力:最后,通过查询矩阵 Q 与 K 计算注意力分数,并使用 V 进行加权求和,得到最终的输出:
通过这种方式,MLA 框架不仅减少了 KV 缓存的存储需求,还保持了模型的性能,使得大规模语言模型的推理变得更加高效。想象一下,原本模型需要一个很大的 “仓库” 来存放键值对信息,现在通过低秩压缩,“仓库” 变小了,在推理过程中内存使用就减少了,推理效率也就提升了。
(2)查询的低秩压缩
-
优化目的:除了对键和值进行压缩,MLA 还对注意力查询(Query)进行低秩压缩,以减少训练过程中的激活内存。查询可以理解为模型在寻找信息时提出的问题,对查询进行压缩,能让模型在训练时更节省内存资源。
-
实现方式:查询的低秩压缩通过类似的投影操作实现,具体公式如下:
其中,WQdown是查询的下投影矩阵, qi是第 i 个 token 的查询向量。通过这个投影操作,将查询向量也进行了低秩压缩。
- 性能保持:尽管 MLA 通过低秩压缩减少了 KV 缓存和激活内存,但它仍然能够保持与标准多头注意力(MHA)相当的性能。这就好比一辆车,经过改装后,不仅更省油(减少内存占用),速度还没有变慢(性能相当)。
(3)旋转位置嵌入(RoPE)
位置信息处理:在处理长序列时,位置信息非常重要。比如 “我今天去了北京” 和 “今天我去了北京”,虽然词语相同,但表达的意思可能因为位置不同而有所差异。MLA 架构结合了旋转位置嵌入(RoPE)来有效处理长序列中的位置依赖问题。
作用:RoPE 通过旋转操作将位置信息嵌入到键和查询中。具体来说,对于位置n和维度2i、2i + 1 ,RoPE 的操作如下:
其中,qn是位置n的查询向量。通过这样的旋转操作,模型能够更好地捕捉长距离依赖关系,从而提升对长序列的处理能力。
3. MLA框架的优势
-
内存占用少:低秩联合压缩和查询的低秩压缩,减少了 KV 缓存和激活内存,降低模型在推理和训练时对内存的需求,利于在资源有限的设备上运行模型。
-
推理效率高:内存占用减少,模型处理信息速度加快,能更高效生成结果,比如在对话系统中能更快回复用户问题。
-
长序列处理能力强:结合 RoPE,模型能更好处理长序列,理解文本中长距离的依赖关系,处理长篇文档时表现更出色。
4. MLA框架的核心价值
MLA(Multi-Head Latent Attention)框架通过低秩联合压缩技术,解决了传统 Transformer 模型在大规模语言模型推理过程中面临的内存瓶颈问题。其核心优势在于显著减少了 KV 缓存的存储需求,同时保持了模型的性能。具体来说,MLA 框架通过低秩压缩和矩阵变换,将高维的键(Key)和值(Value)矩阵压缩到低维空间,再通过上投影矩阵将其恢复到原始维度,从而减少了显存的使用量。这一技术不仅显著降低了内存占用,还提高了推理效率,使得大规模语言模型的推理变得更加高效。此外,MLA 框架具有很强的兼容性,可以无缝集成到现有的 Transformer 模型中,无需对模型架构进行大规模的修改,这使得其在实际应用中具有广泛的应用前景。