Attention进阶史(MHA, MQA, GQA, MLA)

在深度学习领域,注意力机制(Attention Mechanism)自诞生以来便成为推动自然语言处理和计算机视觉等任务发展的核心动力。从最初的多头注意力(MHA)到如今的高效变体,如多查询注意力(MQA)、分组查询注意力(GQA)和多层注意力(MLA),注意力机制不断演进,旨在解决计算效率、内存占用和模型性能之间的平衡问题。与此同时,KV Cache等技术的引入进一步优化了推理阶段的效率,使得大规模模型在实际应用中的部署成为可能。本文将带您回顾注意力机制的进阶史,深入探讨这些技术的设计思想、优势及其对深度学习领域的深远影响。
1. Transformer
Transformer是典型的编码器-解码器模型,该模型主要由两个块组成,如下图所示

- 编码器(encoder):负责把输入信息编码成特征向量的算法组件,输入信息根据不同的任务而不同,可是是文本也可以是图像。注意,文本和图像本身不能直接输入编码器,必须把它们向量化(Embedding)成向量才能送入编码器。
- 解码器(decoder):解码器是服务于生成任务的,如果是判别任务,可以没有解码器结构。解码器需要对编码器的输出结果进行翻译和解释,生成想要的目标序列。因为最初的Transformer是用于生成任务,所以是编码器-解码器架构。
Transformer算法之所以基于编码器-解码器框架,是因为该模型最开始是进行翻译或QA问题的,此时编码器负责接受问题,提取问题中的语言与意图给到解码器,解码器在对应生成相关答案。实际上,不管是编码器还是解码器都可以独立或联合使用,具体取决于任务,以文本信息为例:
(1)仅使用编码器的模型:适用于需要理解输入的任务,例如句子情感取向分类和命名实体识别等。
(2)仅使用解码器的模型:适用于文本生成等生成任务。
(3)使用编码器-解码器的模型:适用于需要输入的生成任务,例如语言翻译或摘要提取等任务。
1.1 MHA: 多头自注意力机制(Multi-heads Self-Attention)
Transformer模型中最关键部分就是自注意力(Self-Attention)机制,正如 Transformer 的论文的标题是“Attention Is All You Need”!以文本问题为例来讲解这个机制。在处理文本问题时,自注意力机制会告诉模型在处理每个单词的表示时,特别关注在句子中传递给它的某些单词,并或多或少地忽略其它单词。简单来说,就是给句子中不同单词分配不同的权重。这是符合常理的,因为一句话中的每个单词重要程度是不一样的,从语法角度说,主谓宾语比其它句子成分更重要,self-attention机制就是模型尝试学习句子成分重要程度的方法。自注意力机制权重的计算过程如下。
计算自注意力的第一步是为编码器的每个输入向量(在本例中是每个词的特征向量)创建三个新向量。它们分别被称为查询向量(Queries)、键向量(Keys)和值向量(Values),简称 $\boldsymbol{q}, \boldsymbol{k}, \boldsymbol{v}$ 向量,如下图所示。这些向量是通过与在训练过程中训练的三个权重矩阵 ( $\boldsymbol{W}^{\mathrm{Q}}, ~ \boldsymbol{W}^{\mathrm{K}}, ~ \boldsymbol{W}^{\mathrm{V}}$ )相乘而创建的。这三个向量的创建过程在模型实现时非常简单,通过神经网络层的映射即可得到。具体来说,输入数据为 token 本身(假设 64 维),而映射后的输入向量可以是 192 维,此时第 0-63 维作为 $\boldsymbol{q}$ 向量,64-127 维作为 $\boldsymbol{k}$ 向量,而 128-192 维作为 $\boldsymbol{v}$ 向量。请注意,这些新向量的维度一般小于输入向量。例如,在原论文中新向量的维度是 64 维,而输入向量的 512 维,这可以一定程度上节省后续的计算开销。

查询向量、键向量和值向量是为计算和思考注意力机制而抽象出的概念,或者说是我们对模型的学习期望。因为这三个新向量在刚创建时是随机初始化的,没有特殊含义,是经过模型训练分别得到了类似查询、回复、存值等向量功能,一个词向量可以通过它们与其它词向量进行互动来建模词与词之间的相关性。在读者阅读完接下来的全部计算过程之后,就会明白它们名字的由来。
计算自注意力的第二步是计算一个相关性分数(score)。假设正在计算本例中第一个词“Thinking”的相关性分数。这个分数就决定着“Thinking”这个词在某句话中与其它词的关联程度,所以“Thinking”这个词要与其它所有词都计算一个分数。
分数是通过查询向量与正在评分的相应单词的键向量的点积计算得出的。点积的公式:
$a \times b=|a| \times|b| \times \cos \theta$ 。其意义就是比较两个向量的相关程度,相关性越高,分数越大,如图

第三步和第四步是将相关性分数除以 8 ( 8 是论文中使用的查询向量维度的平方根,即 $\sqrt{64}=8)$ 。这会使模型训练时的梯度更稳定。然后通过 Softmax 函数映射出最后的结果。 Softmax 函数可以对分数进行归一化处理,使它们都为正且加起来为 1 ,计算过程如下图所示。

Softmax映射后的分数决定了每个词在句子中某个位置的重要性。显然,当词语处于该位置时,其Softmax分数最高。因为它的查询向量、键向量和值向量都是来源于这个词本身的,具体来说,q,k,v向量是当前词向量经过神经网络层映射得到的,是比较相近的三个向量,所以 $q_1, k_1$ 的点积结果就大,表示该词与自身是最相关的。但是,这个词是存在于一个句子中的,因此它与其它词之间应该也存在一个相关性,但这个相关性一般没有与自身的相关性大。所以,当前词的查询向量和别的单词的键向量做点积时,结果就会相对较小。这时的结果表示该词与别的单词的相关性。值得注意的是自注意力机制中“自”的含义是:q,k,v向量都是来源于当前词本身的,是“自己”通过神经网络层映射得到的。
第五步是将每个值向量乘以对应的Softmax分数,目的是做每个单词重要程度的重新分配。最后是加权值向量求和的操作。这会在该位置产生自注意层的输出,例如,词“Thinking”经过自注意力层处理后的输出为output = 0.88x$v_1$ +0.12x$v_2$,即当前这句话经过自注意力层处理后,词“Thinking”的含义包含了88%的自身含义和12%的下一词”Machines”的含义,这样处理就体现了文本上下文的关系。当前这句话中的其它词也要做相同的处理,如下图所示。

自注意力计算到此为止,关于 q,k,v 的含义和作用在上文也有解释。实际上,一个粗略的类比是将其想象为在文件柜中搜索。查询向量就像一张便签纸,上面写着正在研究的主题。键向量就像柜子内文件夹的标签,当标签与便签匹配时,我们取出该文件夹的内容,这些内容就是值向量。不过不只要查找一个值,还要从多个文件夹中进行相关内容的查找,如图

在实际的实现中,会将输入向量打包成矩阵,以矩阵形式完成此计算,以便更快地在计算机中计算处理:

其公式表示如下:
$$
\operatorname{Attention}(\boldsymbol{Q}, \boldsymbol{K}, \boldsymbol{V})=\operatorname{Softmax}\left(\frac{\boldsymbol{Q} \boldsymbol{K}^{\mathrm{T}}}{\sqrt{d_k}}\right) \boldsymbol{V}
$$
其中, $\boldsymbol{Q}, ~ \boldsymbol{K}$ 和 $\boldsymbol{V}$ 是输入矩阵,分别代表查询矩阵,键矩阵和值矩阵,$d_k$ 是向量维度。公式 (6-1)的作用是通过对 $\boldsymbol{Q}$ 和 $\boldsymbol{K}$ 的相似度进行加权,来得到对应于输入的 $\boldsymbol{V}$ 的加权和。
具体来说,这个公式分为三个步骤:
(1)计算 $\boldsymbol{Q}$ 和 $\boldsymbol{K}$ 之间的相似度,即 $\boldsymbol{Q} \boldsymbol{K}^{\mathrm{T}}$ 。
(2)由于 $\boldsymbol{Q}$ 和 $\boldsymbol{K}$ 的维度可能很大,因此需要将其除以 $\sqrt{d_k}$ 来缩放。这有助于避免在 Softmax 计算时出现梯度消失或梯度爆炸的问题。
(3)对相似度矩阵进行 Softmax 操作,得到每个查询向量与所有键向量的权重分布。然后,将这些权重与值矩阵 $\boldsymbol{V}$ 相乘并相加,得到自注意力机制的输出矩阵。
最后, 该论文通过添加一种称为“多头注意力”(Multi-heads self-attention)的机制进一步细化了自注意力层。对于多头注意力,其中有多组查询向量、键向量和值向量,这里把一组 q,k,v 称之为一个头,Transformer原论文中使用八个注意力头。每组注意力头都是可训练的,经过训练可以扩展模型关注不同位置的能力。
举一个形象的类比:把注意力头类比成小学生,那么多个小学生在学习过程中会形成不同的思维模式,对同样的问题会产生不同的理解。这就是为什么要使用多头的原因,就是希望模型可以从不同的角度思考输入信息,如下图:

但是,多头注意力机制也给带来了一个问题。如果使用八个头,经过多头注意力机制后会得到
8
个输出,但是,实际上只需要一个输出结果。所以需要一种方法将这八个输出压缩成一个矩阵,方法也很简单,将它们乘以一个额外的权重矩阵即可。这个操作可以通过一个神经网络层的映射完成。

1.2 Transformer 编码器结构
每个编码器中的自注意力层周围都有一个残差连接,然后是层归一化步骤。归一化的输出再通过前馈网络(Feed Forward Network,FFN)进行映射,以进行进一步处理。前馈网络本质上就是几层神经网络层,其中间采用ReLU激活函数,两层之间采用残差连接。编码器的结构如下图所示。

其中,残差连接可以帮助梯度的反向传播,让模型更快更好地收敛。层归一化用于稳定网络,减轻深度学习模型数值传递不稳定的问题。
至于Transformer中的FNN是一个MLP,它在自注意力机制之后对序列中的每个向量单独应用。FNN 起到两个主要作用:
(1)引入非线性:虽然自注意力机制能捕捉序列中不同位置的向量之间的依赖关系,但它本质上是线性的。通过引入FNN层,Transformer可以学习到输入序列的非线性表示,这有助于模型捕捉更复杂的模式和结构。
(2)局部特征整合:FNN层是一个MLP,对序列中每个位置的向量独立作用。这意味着它可以学习到局部特征并整合这些特征,以形成更丰富的特征表示。这种局部特征整合与自注意力机制中的全局依赖关系形成互补,有助于提高模型性能。换句话说,自注意力机制学习的是向量之间的关系,而FNN学习的是每个向量本身更好的特征表示。
FNN层由两个全连接层和一个非线性激活函数(如 ReLU 或 GELU)组成。假设有一个输入向量 $\boldsymbol{x} \in \mathbb{R}^d$ ,那么 FNN 层的计算过程如下:
$$
\operatorname{FNN}(\boldsymbol{x})=\boldsymbol{W}_2 \cdot \operatorname{ReLU}\left(\boldsymbol{W}_1 \cdot \boldsymbol{x}+b_1\right)+b_2
$$
其中 $\boldsymbol{W}_1, b_1, \boldsymbol{W}_2$ 和 $b_2$ 是需要学习的权重矩阵和偏置向量,ReLU 是激活函数。
以上这些加起来就是编码器的结构组成,可以将编码器堆叠N次以进一步编码信息,其中每一层都有机会学习不同的注意力表示,因此有可能提高Transformer网络的预测能力。最后,在文本信息送进编码器之前,往往需要对输入信息添加位置编码。这是因为 Transformer模型中的自注意力机制是一种全局操作,它在计算时并不考虑输入序列中元素的顺序。然而,在自然语言处理任务中,单词之间的顺序是非常重要的。为了让Transformer 能够捕捉到这种顺序信息,需要为输入添加位置编码。
位置编码是一种表示序列中每个位置信息的向量。位置编码的维度与输入向量的维度相同,因此可以将它们逐元素相加,从而保留位置信息。位置编码可以是固定的(如基于正弦和余弦函数的编码),也可以是可学习的(通过训练得到的向量)。在原始的Transformer 论文中,使用了一种基于正弦和余弦函数的固定位置编码。对于一个给定位置pos和编码维度 i,位置编码的计算公式如下:

通过上述公式可以为输入序列中的每个位置生成一个位置编码向量,该向量具有一定的模式,能够表示该位置在序列中的位置信息。
为了将位置编码添加到输入序列中,可以将输入序列中的每个词语向量与对应位置编码向量相加,得到包含位置信息的输入向量。如下图所示,这样 Transformer 模型就可以更好地处理输入序列中的信息,从而提高模型的性能。

如果这些位置编码是可学习的,那么在模型刚开始训练时,会进行随机初始化,我们期望模型能自己通过学习,找到词语之间的位置相关性。就像期望模型可以通过学习,让词向量映射生成的三个新向量 q, k, v 可以分别执行查询、回复、存值等向量功能一样。
1.3 Transformer 解码器结构
解码器的工作是生成文本序列。解码器具有与编码器类似的子层。它有两个多头注意层、一个前馈网络、残差连接以及每个子层之后的层归一化。这些子层的行为类似于编码器中的层,这里不再做重复赘述。解码器的顶层是一个充当分类器的线性层和一个用于获取单词概率的Softmax 所覆盖。编解码器完整的结构如下图所示。注意,在一些判别任务中可能没有解码器结构,编码器-解码器结构经常存在于一些生成任务中。

以“我是学生”→ “i am a student” 的语言翻译任务为例,如下图所示。

在解码器中,输入包括当前解码器的自注意力层输出和编码器的输出。具体来说,我们将当前解码器的自注意力层输出作为查询向量 q ,这表示解码器当前正在尝试生成的元素,需要哪些信息来进行准确预测。编码器的输出作为键向量 k 和值向量 v ,这为解码器提供了一个关于编码器输入序列的丰富上下文信息。然后使用 self-attention公式来计算注意力分布。通过将解码器的查询向量 q 与编码器的键向量 k 和值向量 v 结合,模型可以决定在生成每个新元素时,应该给予编码器输出中的哪些部分更多的重视。这种方法使得解码器在生成输出时不仅考虑到它自己之前生成的内容,而且还考虑到整个输入序列的内容。这对于生成与输入密切相关的准确和连贯的输出至关重要。
另外,与编码器类似,解码器中也包含了前馈神经网络,用于在生成序列的过程中增强模型的表达能力。需要注意的是,在解码器中会使用遮掩机制。Transformer 中的掩码(masking)机制用于防止模型在处理序列时访问不应该访问的信息。掩码机制在自注意力计算过程中起作用,主要有两种类型:填充掩码(padding mask)和序列掩码(sequence mask,也称为查找掩码或解码器掩码)。
- 填充掩码
在自然语言处理任务中,为了将不同长度的句子输入到模型中,通常需要对较短的句子进行填充,使其与最长句子的长度相同。填充通常使用特殊的符号(如<pad>)表示。填充掩码的目的是在自注意力计算过程中忽略这些填充位置。这样做是因为这些填充符号实际上并不携带任何有意义的信息,我们不希望它们影响其它单词之间的注意力权重计算。填充掩码通过将填充位置对应的注意力logits设置为一个非常大的负数来实现。这样,在应用 Softmax 函数时,填充位置对应的注意力权重会接近于零。
- 序列掩码
序列掩码主要应用于Transformer 解码器。掩码防止模型访问未来的输出,确保训练时只依赖当前及之前的信息。这在训练序列模型(如机器翻译、文本生成等)中至关重要,主要原因如下:
1. 保持模型训练的合理性
在推理阶段,模型只能基于当前及之前的信息生成输出,无法访问未来信息。如果在训练时允许模型访问未来信息,会导致训练与推理条件不一致,模型在实际应用中表现不佳。序列生成任务通常遵循因果关系,即当前输出只能依赖过去的信息,不能依赖未来的信息。
2. 避免模型作弊
如果模型在训练时能访问未来信息,它可能会直接“抄袭”未来的正确答案,而不是学习如何基于当前信息合理预测。这会导致模型无法真正学会生成合理的序列。模型如果依赖未来信息,其泛化能力会显著下降,因为在实际应用中未来信息是不可用的。
3. 确保自回归生成的有效性
自回归模型:在文本生成、机器翻译等任务中,模型通常是自回归的,即逐个生成序列中的元素。每一步的生成都依赖于之前生成的元素,而不能依赖未来的元素。通过掩码(如Transformer中的解码器掩码),确保模型在生成当前元素时只能看到之前的部分,而不能看到未来的部分。
4. 注意力机制的需要
在Transformer等模型中,解码器的自注意力层需要掩码来防止当前时间步关注未来的时间步。如果不加掩码,模型会直接“看到”整个序列,破坏自回归生成的过程。掩码确保训练时解码器的行为与推理时一致,即只能基于已生成的部分预测下一步。
5. 实际应用中的不可见性
在实际应用中,模型无法提前知道未来的信息。如果在训练时允许模型访问未来信息,会导致模型在实际场景中失效。在评估模型性能时,必须确保模型没有“偷看”未来信息,否则评估结果会失真。
我们举个例子来理解一些序列掩码的实现:假设序列长度为4,某时刻模型只有两个token作为输入,并且我们正在观察第二个token。在这种情况下,最后两个token被屏蔽,其对应的权重分数也将是0,如下图所示。

序列掩码通过在注意力 logits矩阵中添加一个下三角矩阵(其上三角部分为负无穷)来实现。这样,在应用
Softmax 函数时,当前位置之后的单词对应的注意力权重将接近于零。这使得解码器在每个时间步只能关注当前及之前的单词。下面进行详细的示例计算推导,这种掩蔽通常通过被称为注意掩蔽的矩阵来实现,如下图。

想象一个由四个单词组成的序列(例如“robot must obey orders”)。先可视化其注意力分数的计算:

相乘之后使用注意掩蔽矩阵,它将我们想要屏蔽的单元格设置为负无穷大或一个非常大的负数,如图:

然后,在每一行上应用 Softmax ,产生用于自注意力的实际分数,如图:

此分数表的含义如下:
(1)当模型处理数据集中的第一个token(第1行)时,该token仅包含一个单词(“机器人”),其100%的注意力将集中在该单词上。
(2)当模型处理数据集中的第二个token(第2行)时,其中包含单词(“机器人 必须”),当模型处理单词“必须”时,48%的注意力将集中在“机器人”上,而52%的注意力将集中在“必须”上。
总结一下,掩码机制在Transformer模型中具有重要作用。填充掩码用于忽略填充符号的影响,而序列掩码确保解码器在生成过程中遵循自回归原则。掩码的存在使得模型在处理序列时更加稳定、可靠。
1.4 Transformer 顶层结构
Transformer的顶层包含线性层(神经网络层)与Softmax层,是用于生成目标序列中每个位置上的词语概率分布的关键组件,如下图。

Softmax函数是一个常用的激活函数,可以将任意实数向量转换为概率分布。在Transformer中,使用Softmax函数将线性顶层输出的向量转换为目标序列中每个位置上的词语概率分布。
Transformer的顶层使用一个全连接层将特征向量映射到一个大小为词表大小的向量,然后使用Softmax函数将其转换为概率分布。这样就可以预测目标序列中每个位置上的词语,并且可以使用这些预测来计算损失函数,从而训练模型。
以上就是Transformer的计算机制。Transformer利用注意力机制可以做出更好的预测。之前的循环神经网络试图实现类似的事情,但因为它们受到短期记忆的影响,效果不如Transformer算法好,特别是当编码或生成较长序列。因为Transformer架构在计算注意力时的计算范围是全局的,即查询向量和其它所有的键向量都要做点积来计算相关性,因此Transformer可以捕获长距离的信息依赖,理论上这个距离是可以覆盖当前处理的整句话的。
2. KV Cache:Attention机制进阶的核心
KV Cache 是一种优化技术,旨在提高模型在推理阶段的效率。它通过缓存键(K)和值(V)的值,减少重复计算,从而加速解码器中的矩阵运算。这种技术在处理大规模语言模型时尤为重要,因为它可以显著减少推理时间,尽管会带来额外的内存开销。后续的 MQA, GQA, MLA 等都是基于KV Cache的改进优化。
在讲解KV Cache之前,首先明确其四个最重要的核心:
- KV Cache 应用于推理阶段
KV Cache 主要用于模型的推理阶段。在这个阶段,键(K)和值(V)的值保持不变,因此可以通过缓存这些值来避免重复计算。
- KV Cache 只存在于解码器中
KV Cache 技术通常应用于解码器部分。这是因为解码器在生成输出时需要进行多次矩阵运算,而 KV Cache 可以显著优化这一过程。
- 加速 Q@K@V 的矩阵相乘
KV Cache 的主要目的是加速 Q@K@V 的两次矩阵相乘操作。通过缓存 K 和 V 的值,模型可以减少计算量,从而提高推理速度。
- KV Cache 会增加内存占用
尽管 KV Cache 可以加速推理过程,但它也会增加内存的占用。这是因为需要额外的存储空间来保存缓存的 K 和 V 值。
以一个长度为4的序列为例,在不使用 KV Cache 的情况下,长度为4的序列的注意力计算过程需要进行四次独立的计算,每次计算对应序列中的一个元素。
1.第一次计算(处理 $x_1$ ):
- 计算 $Q_1=W_Q \cdot x_1$
- 计算 $K_1=W_K \cdot x_1$ 和 $V_1=W_V \cdot x_1$
- 计算 $Q_1$ 与所有 $K_j(j=1)$ 的点积,得到注意力分数
- 通过 softmax 计算注意力权重,并加权求和得到 Output ${ }_1$
2.第二次计算(处理 $x_2$ ):
- 计算 $Q_2=W_Q \cdot x_2$
- 计算 $K_2=W_K \cdot x_2$ 和 $V_2=W_V \cdot x_2$
- 计算 $Q_2$ 与所有 $K_j(j=1,2)$ 的点积,得到注意力分数
- 通过 softmax 计算注意力权重,并加权求和得到 Output ${ }_2$
3.第三次计算(处理 $x_3$ ):
- 计算 $Q_3=W_Q \cdot x_3$
- 计算 $K_3=W_K \cdot x_3$ 和 $V_3=W_V \cdot x_3$
- 计算 $Q_3$ 与所有 $K_j(j=1,2,3)$ 的点积,得到注意力分数
- 通过 softmax 计算注意力权重,并加权求和得到 Output ${ }_3$
4.第四次计算(处理 $x_4$ ):
- 计算 $Q_4=W_Q \cdot x_4$
- 计算 $K_4=W_K \cdot x_4$ 和 $V_4=W_V \cdot x_4$
- 计算 $Q_4$ 与所有 $K_j(j=1,2,3,4)$ 的点积,得到注意力分数
- 通过 softmax 计算注意力权重,并加权求和得到 Output ${ }_4$
如下图所示:

在自回归模型中(autoregressive models),会逐个生成文本的每个token,这个过程可能比较慢,因为模型一次只能生成一个token,而且每次新的预测都依赖于之前的上下文。这意味着,要预测第4个token,你需要用到前3个token的信息,这通常涉及到对这些token的表示进行一系列矩阵乘法运算。当序列长度很大时,例如要预测第1001个token,你不仅需要前999个token的信息,还要加上第1000个token的信息,这使得整个QK矩阵非常巨大(1000×1000),进而导致attention的计算量巨大。KV cache就是在这里发挥作用,通过存储之前K , V 的计算结果,并在后续的token生成时复用这些结果,从而避免重复计算。如下图所示:

更具体地说,KV cache在自回归生成模型中充当一个内存库的角色,模型会存储来自自注意力层(self-attention layers)的之前处理过的token的键值对。在Transformer架构中,Self Attention层通过将Queries(Q)与Keys(K)相乘来计算注意力分数,然后产生Values(V)向量的加权和作为输出。通过存储这些信息,模型可以避免重复的计算,而是直接从缓存中检索之前token的K 和V 。
需要注意的是,KV cache只在多个token生成步骤中发生,并且仅在decoder进行(例如,在decoder only的模型如GPT,或在encoder-decoder模型如T5的解码部分)。像BERT这样的encoder only模型,不是生成型的,因此不涉及KV cache。
由于decoder是causal的(即,一个token的注意力attention只依赖于它前面的token),在每一步生成过程中,我们实际上是在重复计算相同的前一个token的注意力,而我们真正需要做的是仅计算新token的注意力。这就是KV cache发挥作用的地方。通过缓存之前的K和V,我们可以专注于只计算新token的注意力。
为什么只存储K和V,而不存储Q?
- 由于Q是当前时间步的输入,每次生成新token时都会重新计算Q,不需要缓存。
- 而K和V可以复用之前时间步的计算结果,通过缓存K和V,可以避免在每个时间步重新计算它们,从而提升效率。
KV Cache的占用内存会随着QA的进行线性增长
KV Cache通常用于存储模型在处理序列数据时生成的键值对,以便在生成后续内容时能够快速访问之前的信息。
在实际实验时,每次问答(QA)记录都会增加KV Cache的存储需求,因为模型需要保留之前问答的上下文信息以生成连贯的响应。这种线性增长的内存占用可能会导致在处理长对话或大量问答时,内存需求显著增加,从而影响系统的性能和效率。

由于内存限制,LLM的回复能力确实会随着问答的进行逐渐变差,这种现象可以理解为模型具备一定的“遗忘性”。具体而言,由于硬件资源的限制,KV Cache的大小是有限的。当缓存达到其容量上限时,旧的信息可能会被新的信息覆盖或丢弃。其表现为随着问答的进行,早期的对话内容可能会因为KV Cache的容量限制而被移除或覆盖,导致模型逐渐“遗忘”之前的上下文。由于模型无法访问完整的对话历史,其生成的回复可能会变得不够准确或连贯,尤其是在需要依赖早期信息的情况下。所以,在长对话或多轮问答中,模型的性能可能会显著下降,因为它无法有效地利用整个对话历史。
简而言之,KV Cache是一种用内存换取推理速度的技巧。
3. MQA: Multi-Query Attention
多查询注意力机制 (MQA) 是 Transformer 中使用的传统多头自注意力机制(MHA)的一种变体。在传统的多头注意力机制中,每个注意力头都使用自己的一组查询、键和值,这可能需要大量计算,尤其是在注意力头数量增加的情况下。MQA 通过在多个注意力头之间共享同一组键和值,同时为每个注意力头维护不同的查询,简化了这一过程。这种方法减少了计算和内存开销,而不会显著影响模型的性能。

多查询注意力机制的关键概念:
- 共享键和值:与传统的多头注意力(其中每个头都有自己的键和值)不同,MQA 对所有注意力头使用相同的键和值。
- 不同查询:MQA 中的每个注意力头仍然使用自己的一组查询,从而允许它从不同方面关注输入序列。
- 效率:通过共享键和值,MQA 减少了所需的计算量和内存,使其比传统的多头注意力更高效。
多查询注意力机制的好处:
- 降低计算复杂度:通过共享键和值,MQA 显著减少了所需的操作数量,使其比传统的多头注意力更高效。
- 更低的内存使用率:MQA 通过存储更少的键和值矩阵来减少内存使用率,这对于处理长序列特别有益。
- 保持性能:尽管效率有所提高,MQA 仍保持与传统多头注意力机制相媲美的竞争性能,使其成为大规模 NLP 任务的可行选择。
4. GQA:Group Query Attention
组查询注意力 (GQA) 是对 Transformer 中使用的传统多头自注意力机制和多查询注意力机制的折中。在标准多头自注意力中,每个注意力头独立处理整个序列。这种方法虽然功能强大,但计算成本高昂,尤其是对于长序列。而MQA虽然通过在多个注意力头之间共享同一组键和值简化了这一过程,但其简化也不可避免的带来了一些精度的损失。GQA 通过将查询分组在一起来解决此问题,从而降低了计算复杂性,而不会显著影响性能。

组查询注意力机制的关键概念:
- 分组查询:在 GQA 中,查询根据其相似性或其他标准分组。这允许模型在类似的查询之间共享计算,从而减少所需的总体操作数量。
- 共享键和值表示:GQA 不会为每个查询计算单独的键和值表示,而是为每个组计算共享键和值表示。这进一步减少了计算负载和内存使用量。
- 高效计算:通过分组查询和共享计算,GQA 可以更高效地处理更长的序列,使其适合需要处理大量文本或数据的任务。
实际上,MHA和MQA都可以看作是GQA的特殊情况:
- 当组数 g 与 头数 head 相等时,GQA = MHA。
- 当组数 g 为 1 时,GQA = MQA。
大家如果对卷积算法比较熟悉的话,MHA, MQA, GQA 的关系与卷积,逐通道卷积,组卷积的关系是一致的。
5. MLA: Multi Head Latent Attention

多头潜在注意力 (MLA) 将潜在特征表示纳入注意力机制,以降低计算复杂度并改善上下文表示。MLA的核心是对KV进行压缩后,再送入标准的MHA算法中,用一个更短的k,v向量来进行计算,进而减少KV Cache的大小。

上图中的公式描述了多头注意力机制(Multi-Head Attention, MHA)中潜在向量(latent vectors)的计算过程。以下是对这些公式的逐步解读:
1.计算代表 KV Cache 的潜在向量
$$
\left[c_t^{K V}\right]=W^{D K V} \mathbf{h}_t
$$
-解释:这里,$c_t^{K V}$ 是在时间步 $t$ 计算的键值缓存潜在向量。 $W^{D K V}$ 是一个权重矩阵,用于将隐藏状态 $\mathbf{h}_t$ 映射到键值缓存空间,这一步可以通过神经网络映射得到。$c_t^{K V}$ 相对与原来的$\mathbf{h}_t$ 要小很多。
2.计算键(Key)和值(value)的潜在向量
$$
\left[k_{t, 1}^C ; k_{t, 2}^U ; \ldots ; k_{t, n_h}^U\right]=k_t^C=W^{U K} c_t^{K V}
$$
$$
\left[v_{t, 1}^C ; v_{t, 2}^U ; \ldots ; v_{t, n_h}^U\right]=v_t^C=W^{U K} c_t^{K V}
$$
-解释:$k_t^C$ 是键的潜在向量,通过将 $c_t^{K V}$ 与权重矩阵 $W^{U K}$ 相乘得到,这一步是做上采样,通过潜向量特征$c_t^{K V}$映射得到较大的$k_t^C$用于后续的注意力计算。$v_t^C$的计算同理。
3.计算旋转位置编码(RoPE)
$$
k_t^{R}=\operatorname{RoPE}\left(W^{K R} \mathbf{h}_t\right)
$$
-解释:RoPE(Rotary Position Embedding)是一种位置编码方法,用于在键向量中引入位置信息。
4.组合潜向量k和位置编码k得到最终的键向量
$$
k_{t, i}=\left[k_{t, i}^C ; k_t^R\right]
$$
-解释:最终的键向量 $k_{t, i}$ 是通过将内容相关的键向量 $k_{t, i}^C$ 和位置编码 $k_t^R$ 连接起来得到的。
5.计算查询(Query)的潜在向量

-与K向量的计算类似,通过潜在向量计算得到参与后续MHA计算的查询向量q。
6.注意力计算
最终的注意力输出 $u_t$ 是通过将查询 $\left(q_{t, i}\right)$ ,键 $\left(k_{j, i}\right)$ 和值 $\left(v_{j, i}^C\right)$ 结合起来计算的:

- 解释:
- $\mathbf{o}_{t, i}$ 是第 $i$ 个注意力头的输出。
- Softmax 函数用于计算注意力权重,其中 $q_{t, i}^T k_{j, i}$ 是查询和键的点积,$\sqrt{d_h+d_h^R}$ 是缩放因子,用于稳定梯度。
-$v_{j, i}^C$ 是第 $j$ 个时间步的值向量。
上述的公式计算中,需要注意以下几个问题:
- 相较与MHA而已,MLA是再MHA之前多了一些qkv的处理,这看上去是多了一些计算,为什么会变的更高效?
- 位置编码信息RoPE为什么不直接加在$k_{t, i}^C$上,而要新创建一些$ k_t^R$?
- 为什么对于查询向量q,也要进行潜向量的计算?
针对问题1,虽然看上去是多了一些计算,但是这种方式使得我们的KV Cache变小了。具体来说,我们只需要缓存一个较短的$c_t^{KV}$向量,而不是$k_t^{C}$和$v_t^{C}$。用白话说就是缓存长的变成缓存短的,缓存两个变成缓存一个。
针对问题2,是因为旋转位置编码RoPE与潜向量的计算不兼容,为了同时使用潜向量计算和旋转位置编码RoPE两个技术,只能多创建一个新的向量$ k_t^R$来编码位置信息,将来通过向量合并将位置信息带入键向量。
针对问题3,将Q的输入也改为了低秩投影形式(潜向量),这与减少KV Cache无关。个人认为主要是为了特征的对齐,如果只对键 k 和值 v 进行潜在向量计算,而忽略查询 q,会导致 q 和 k、v 的特征空间不一致,影响注意力机制的效果。在注意力机制中,q、k、v是平等的输入,对它们进行相同的潜在向量计算可以保持模型的对称性和一致性。
最后,大家如果对残差网络ResNet熟悉的话,也可以参照瓶颈残差模块来理解MLA的潜向量计算。瓶颈残差模块通过“压缩-提取-扩展”的结构(1×1卷积降低通道数、3×3卷积提取特征、1×1卷积恢复通道数)在减少计算量的同时高效提取特征,类似于MLA通过多层注意力机制逐步生成潜在向量,对输入信息进行压缩和抽象,从而减少后续注意力计算的开销。两者都通过分层结构逐步提取更高级的特征表示,同时保留原始信息,在计算效率和特征表达能力之间取得了平衡。
2 条评论
发表回复 取消回复
作者
arwin.yu.98@gmail.com
写的好详细,前馈神经网络好像写错了,FFN,你这边写的是FNN
哈哈哈 是的 笔误 抱歉啦