大模型LLMs组件系列：深入浅出解析注意力Attention机制

原地漫游

已于 2024-08-22 19:07:01 修改

阅读量789

点赞数 22

分类专栏：大模型LLMs组件系列文章标签：架构人工智能语言模型自然语言处理

于 2024-08-20 19:09:20 首次发布

本文链接：https://blog.csdn.net/jpw41/article/details/141366095

版权

大模型LLMs组件系列专栏收录该内容

7 篇文章 0 订阅

订阅专栏

💡 注意力机制萌芽于上世纪80年代，最早是CNN视觉检测领域的技术，灵感来源于生物注意力，目的是从纷繁的信息中提取高价值信息，舍弃低价值信息。2014年，注意力机制被用于NLP领域，并于2017年开始作为Transformer架构的核心机制几乎一统NLP领域方法论。当前绝大多数大模型使用的Attention机制执行的是缩放点积注意力（Scaled Dot-Product Attention, SDPA），基本原理是对任意输入X，使其与Q、K、V三个投影矩阵相乘得到三个对应的隐状态向量，然后将Q与K状态向量进行点积后除以缩放因子得到注意力评分/权重。最后用这个权重与值向量点积得到输出。网上将Transformer架构的文章数量众多而良莠不齐，常常因信息超载而令人劳神。本篇的目的在于使用尽可能少的公式的同时深入浅出解析Transformer架构的机制，并回答一些常见问题。

1 简单self-Attention：注意力机制的原理与优势

1.1 最简单、单Query的、无Batch_size维度的Attention机制图解

简单attention
如上图，假设我们的输入的batch_size=1、seq_len=3的token序列input_ids在经过embedding，前归一化层，位置编码后，以input_embed张量的形式进入Attention机制。在这个假设中，input_embed.shape=torch.Size([3, 8])，分别对应token序列长度和embedding空间维数hidden_size（直观起见没有batch_size维度），简化的代码如下。

inputs = "你好啊"
input_ids = tokenizer( inputs )               # torch.tensor( [8, 4, 7] )
input_embed = self.embedding( input_ids )     # torch.shape = torch.Size([3, 8])
'''
在经过前归一化层、位置编码后，input_embed = tensor([ 
  [ 0.23, 0.81, 0.45, 0.33, 0.44, 0.12, 0.17, 0.88 ],
  [ 0.90, 0.78, 0.32, 0.55, 0.28, 0.31, 0.49, 0.73 ],
  [ 0.12, 0.31, 0.56, 0.71, 0.02, 0.80, 0.34, 0.36 ],
])
'''

注意力机制的核心就是让所有token都用自己的状态向量与包括自己在内的所有token进行点积，相互打分。所以让input_embed中的三个长度为8的张量两两相乘？不行，太简单了，还会使得结果方阵对角线下一半是镜像。那么怎么办？

答：先用两个形状一样但值不一样的可学习权重矩阵对input_embed进行投影变换，再两两相乘，这下不仅在Attention机制中加入了神经网络权重参数，还解决了input_embed自乘带来的镜像问题。这两个权重矩阵就是投影矩阵q_proj和k_proj，为了消除量纲，还需要除以一个基于head_dim的缩放因子，简化代码如下：

self.q_proj = torch.nn.Linear( hidden_size, head_dim )
self.k_proj = torch.nn.Linear( hidden_size, head_dim )
# d_q和d_k分别是两个投影变换矩阵的输出维度，相等，均为head_dim，因为两者的输出结果要进行点积。再上图例子中，两者值都等于6。
query_states = self.q_proj( input_embed )
key_states = self.k_proj( input_embed )
attn_weights = torch.matmul(query_states, key_states.transpose(0, 2)) / math.sqrt(self.head_dim)
# attn_weights.shape = torch.Size([3, 3])

这样就够了吗？当然不行，这里的运算结果将是一个seq_len*seq_len的评分方阵，输入token的隐藏状态不见了，那么怎么办？

答：要想再次获得token的状态信息，最直接的方式是再用这个评分矩阵乘以输入input_embed，理论上也可以这样做，只不过是效果如何的问题。Transformer机制给出的方案是再来一个投影矩阵v_proj对input_embed进行变换，然后再用上述评分矩阵乘以这个变换结果，不论如何，在计算效率允许的范围内，加一个可学习层总共是利大于弊的吧？代码简化如下：

self.v_proj = torch.nn.Linear( hidden_size, d_v )
# d_v和head_dim（d_q和d_k）可以不一样，上图中就设置为4，但很多大模型都设为一样，可能是为了计算方便。
value_states = self.v_proj( input_embed )
attn_weights = nn.functional.softmax(attn_weights, dim=-1)
attn_output = torch.matmul(attn_weights, value_states)
# 或者，使用下面torch封装好的SDPA方法一步到位，
attn_output = torch.nn.functional.scaled_dot_product_attention(
            query_states,
            key_states,
            value_states,
            attn_mask=causal_mask,
            dropout_p=self.attention_dropout if self.training else 0.0,
            is_causal=is_causal,
        )
# 计算完毕后attn_output.shape=torch.Size([seq_len=3, d_v=4])

上面代码实际上在执行下式

在这里插入图片描述

可以看到，理论上Attention机制的运算结果是一个shape和输入seq_len一样但状态维度可能不一样的张量，而我们知道所有结果是要在不同Transformer Block之间迭代的，这就要求一个Block的输出结果的shape必须等于输入的shape。为了实现这一目的，Block中Attention模块后会有一个MLP模块，可以将上述d_v维度映射到hidden_size维度。当然，实现这个目的只是附带，MLP还有更重要的任务（详情见本系列MLP部分），这个目的也不只有MLP能完成（比如一开始就设置d_v和head_dim一样）。

1.2 Attention机制揭秘：继任社长选拔原理

流程说完了，下面列出几个关键点和困惑点。开始之前，我们要明确的一个事实是，不论大模型的机制多花哨，其总是每次迭代都会生成一个next token的，而整个Transformer Block机制都是在做选择下一个token的事儿，选择的方法是最终生成一个hidden_states并将其输入一个shape=[hidden_size, vocab_size]的分类器（如lm_head层），得到取词表中每个token的概率。为了直观起见，我们将Transformer机制类比为某兴趣社团选拔下一任社长，词表空间中每个token代表这个社团中的每个人，每个token的hidden_states则表示人的人格特质，可以理解为经典及其学习中的特征向量；一个输入token序列就是这个社团的历任社长，而结果要求的next token则是该社团将从全体成员中选出的下一任社长，选拔的决定权在序列中最后一个token（现任社长）手中，当然也会综合考虑之前社长的影响。

问题1：Attention机制里的Q、K、V三个投影矩阵分别代表什么？——社长选拔机制

答：self_attention机制中，每个输入都要经过q、k、v三个投影矩阵变为三个不同的状态向量，其中q、k、v的概念其实来源于检索领域的Query查询Key-Value键值对。在这个机制中，输入序列的每个token都同时具备query，key，value三个状态向量，其中Query代表这个token主动的需求，内含主体的偏好；Key来自当token作为被观察的对象，具备被观察者自身的各异的表面的特质；Value则是token表象下隐藏的实质，是被查询到键后获得的反馈值。因此注意力点积的结果（如缩放点积注意力SDPA）代表每个有偏好的观察者token在观察到具有特质的对象token后获得的注意力反馈值。

在上面的例子中，q、k、v三个投影矩阵代表个人综合特质input_embed在选社长这个任务上可变的特质投射。其中，Query代表每任社长对社长人选在d_q=6个方面的期待和偏好，Key代表每任社长自身在这d_q=6个方面的可见的表现，Value则代表其作为社长的内敛的特质。QKV运算获得的结果和输入shape一样。如果是这个选拔机制第一次运行（第一次迭代），所有社长都要参与，但决定权在现任社长手中；之后只需要现任社长参与就行（kv_cahce）。每次迭代的结果就是现任社长在历史背景影响下产生的对继任的特质的偏好，特质表现为一个维数与input_embed一致的状态张量。这个张量将进入lm_head分类器对全体社团成员投票，最高票者当选下任社长。

attn_weights

问题2：Q和K点积得到的attn_weights矩阵代表什么意思？——历任社长之间的相互评价

如上，attn_weights是query_states与key_states的转置点积，并缩放、Softmax后的结果，方阵的每行代表不同token的query向量对所有token的key向量的打分。其中，q_1行k_0列的值0.90代表第2个token作为主体观察作为客体的第1个token时的评分为0.90。因此可以简单理解为，attn_weights矩阵代表所有token对包含自己在内的其他token的有向评分。

用上面的例子，attn_weights矩阵就是历任社长根据d_q=6个方面对自己和他人打分的结果，如果对这些方面的期待比较大而对方又表现得好，就会获得高分；如果不是那么期待和看重某一方面，那么就算对方在该方面表现得好也不会很高分。但实际上，这只是选拔机制第一次运行（第一次迭代）中走的形式，这个attn_weights矩阵只有最后一行，也就是现任社长（最后一个token）对所有社长的评分会被后续使用；而在之后的选拔（之后的迭代中），这个形式也不走了。

注意力权重点积值向量

问题3：attn_weights矩阵与V点积的结果代表什么意思？——现任社长决定的继任可能的特质

如上图，QK缩放点积后获得的注意力权重矩阵attn_weights与value_states执行点积后获得一个shape与后者一样的状态张量attn_outputs（第一次迭代，之后就只有一行了）。该张量每一行的状态张量代表一个token最终的注意力结果，而每行的结果其实都是value_states的三行状态张量加权得到的，权就是该行所代表的token对所有token的注意力打分，也就是attn_weights中的对应行。例如，attn_outputs中第一行就是attn_weights中第三行([0.23, 0.48, 0.29])与整个value_states矩阵点积的结果；更细致的，attn_outputs中第三行第一列的绿色块值就是attn_weights中第三行([0.23, 0.48, 0.29])与value_states矩阵第一列([v_00, v_10, v_20])点积（逐位相乘后相加）的值。因此，attn_outputs中每行代表每个token对全部token的关注权重及其值特征向量的加权和。但在CausalLM任务中实际上只有第三行有用。

换用上面的例子，attn_outputs中最后一行代表现任社长在综合考虑历任社长特质后对继任特质的期望，这个期望是由现任社长对包括自己在内的所有人的评价和所有人作为社长的内敛的特质决定的。现任社长对哪一任的评价越高，那么这位社长的内在特质将对继任社长的期望标准施加最大的影响（对结果值施加最大的权重）。

问题4：为什么两个token的状态向量点积就能表示互相打分？

答：说白了，每个token的hidden_state其实上功能等价于经典机器学习中的特征向量，只不过人类不可读。里面包含了关于自身的信息，而点积则是两个token的特征向量逐位相乘后相加的过程，可以视作简化版的余弦相似性（cos(θ) = (A · B) / (||A|| * ||B||)）计算，两个token如果点积后获得的数值大，就说明两者在这个hidden_size大小的语义空间内越相似，越有可能有关系，也越有可能是值得关注的对象。

问题5：Attention机制在当时的创新点（解决的痛点）是什么？

答：input_embed作为一个已经准备好了的输入数据的状态张量，可以理解为有3行，每行对应一个token且有一个长度为8的向量来表示这个token，语义信息就在这个状态张量里。在Transformer之前，主流采用LSTM这类RNN网络来处理文本序列，即将序列长度3视作时间步，每个时间步的长度为8的状态向量按照时间步进入LSTM网络运算，前一个时间步的运算结果将作为后一个时间步的隐状态。这种方式主要有以下缺点

串行结构：虽然直观地考虑了一句话的语义时序，但前一个运算完毕后才轮得到后一个，就算有空闲资源也不能同时计算。
长距离依赖问题：串行结构带来的另一个问题是，两个时间步相隔越远，前者给后者带来的影响越小。

RNN和LSTM细节在这里留个坑，总而言之传统的LSTM方法由于串行结构会出现训练开销大、梯度问题等情况，而Transformer就很好地解决了这个问题。简单而言，Attention机制让整个序列中所有token（在这里长度为3）的状态张量两两点积（给互相打分）形成一个方阵评分表，这个两两相互打分就是注意力核心优势的来源：

每次只涉及两个token/时间步之间进行点积，谁先谁后无所谓，因此可以并行化进行；
打分的依据只有两个token自身的状态向量，不涉及在序列中的前后位置，因此没有长距离依赖衰退的问题，对梯度更友好。

2 Key-Value Cache原理

Key-Value Cache是对Transformer架构在工程上的优化，核心是将Attention机制中k_proj和v_proj两个投影矩阵对每个token的投影结果保存在内存中以待后续复用，这是一种典型的以空间换时间的工程思想。kv cache大大加快了大模型推理的速度，也带来了一些问题，比如自此大模型的生成过程就有了第一次迭代和后续迭代之分，kv_cache会占用大量显存等。

实现形式上，第一次迭代计算当前长度为3的所有token的query_states，key_states，value_states后以元组形式保存后两者。至于为什么不保存query_states，因为只有最后一个token对其他token的注意力评分才有用，所以本质上第一次迭代时也没必要生成所有token的query_states，最后的outputs也只需要最后一层就行了（实际上在最后进入分类器前也是这么操作的）。但当前大模型中都没有讲究这个，我猜一方面是因为保持三个投影矩阵操作一致会让代码更简洁可读，另一方面就是仅在第一轮出现这个情况无伤大雅。

下面是后续迭代的示意图，输入变成了seq_len为1的张量，经过三个投影矩阵后生成的张量第一个维度也为1。key和value的状态通过kv_cache直接补充，可以看到，这张图中更加清楚地突出了前述现任社长对继任社长的人选有决定权的意思了吧？后续迭代的重点就是获得新token对之前所有token的评价形成加权和中的权重，反之则不重要。

attention机制2

一个大模型有n个Transformer Blocks堆叠，每个Block都缓存各自的kv_cache。所以模型一旦变大、对话序列一旦边长，kv_cache的显存消耗将成为瓶颈。

3 不同情况下的Attention

3.1 Self-Attention

1中主要讲述的是自注意力self-attention的机制。所谓自注意力，实际指的是attention点积中使用的query_states, key_states, value_states来自同一个输入input_embed的投影变换。自注意力常用于Encoder-Decoder架构中的Encoder和Decoder中的Masked Self-Attention。

3.2 Cross-Attention

交叉注意力Cross-Attention是相对Self-Attention的，这种注意力机制只用于Encoder-Decoder架构中的Decoder中。交叉指的是attention点积中使用的QKV状态来自不同输入的投影变换，其中key_states, value_states来自于Encoder层的输出，query_states后者来自本Deocder层的输入。

在这里插入图片描述

3.3 Masked Self-Attention

Masked Self-Attention也是Encoder-Decoder架构中的Decoder专有的，主要用于在训练中评估Mask掉序列中的未来token。比如我们训练的语料是"你今天真好看"，在自监督的范式下通常迭代使用前文预测后文，比如使用"你今天"三个字预测"真"这个自。由于Encoder是双向注意力的而Decoder则是一个个生成next token的，所以对Decoder来说，在预测"真"的时候不能看到"真好看"这三个相对于"你今天"的未来token。我们知道当今的大模型都是Decoder-only架构的，Transformer Block中只有一个self-attention模块，那么这种机制还成立吗？

答：还成立，只不过这种机制只用于训练阶段，通过input_ids伴随的attention_mask张量实现。在推理过程中大模型只有Self-Attention。这种Mask掉所有之后token的单向注意力范式叫Causal Decoder，即因果解码。此外，还出现了一种新的Mask方法叫前缀解码Prefix Decoder，核心原理是在token序列的前一部分（Prefix部分）采用Encoder式的双向注意力，而在后面的非Prefix部分采用Causal Decoder的单向注意力，区别如下图。

在这里插入图片描述

4 点积和加性注意力

加性注意力Additive Attention和点积注意力Multiplicative Attention（缩放点积注意力，Scaled Dot-Product Attention, SDPA）是两种常见的注意力机制，前者最初用在seq2seq，后者最初用在Transformer。

4.1 加性注意力

通过一个前馈神经网络计算相似度，将相似度通过softmax函数进行归一化，利用注意力权重对值value进行加权求和，得到输出。
$e_{ij}=v^Ttanh(W_qQ_i+W_kK_j)$

4.2 点积注意力

通过计算查询Q和键K的点积来衡量相似度，并加上一个缩放因子，将相似度通过softmax函数进行归一化，利用注意力权重对值value进行加权求和，得到输出。
$e_{ij}=\frac{Q_i·K_j^T}{\sqrt{d_k}}$

在Transformer原始论文中就讨论了两种注意力实现的优劣，主要观点是

加性注意力由于涉及拼接向量和通过前馈神经网络进行计算，在处理较小的维度时表现良好，因为前馈神经网络能够捕捉更复杂的关系。
点积注意力主要操作是矩阵乘法和缩放，计算复杂度和加性差不多，但点积因为可以利用矩阵乘法并行工具而更加快速和节约内存。此外，在处理较高维度时表现更好，因为点积操作能够更有效地计算相似性，并且缩放可以防止数值不稳定。

因此当前大模型主流使用缩放点击注意力SDPA。

5 多头和多查询注意力：MHA、MQA、GQA

5.1 多头注意力MHA

Transformer的原始论文中对自注意力采用了多头机制，即Multi-Head Attention (MHA)。之所以要分多个头，是因为多头注意力机制将隐状态向量hidden_states分成多个头，每个头形成一个子语义空间，并能关注不同维度的信息。这样，模型可以同时捕捉到不同类型的特征，从而提高了模型的表达能力和注意力分配。

在具体实现上，多个头可并行地完成各自地注意力计算，拼接成一个完整的结果，然后经过一个全新的投影矩阵o_proj。这个投影矩阵的作用是将多头拼接的结果Mix混合打乱，让局部特质影响全局。如下：

在这里插入图片描述

5.2 多查询注意力MQA

但MHA要求每个Query都要对应一组Key-Value，可能有冗余，计算也慢。于是2019年提出了Multi-Query Attention（多查询注意力，简称MQA）。所有的Query头共享同1份Key和Value矩阵，从而实现参数量和计算量的减少。MQA的目的是加快计算，但由于参数量减少了，生成的质量不如MHA。

在这里插入图片描述

5.3 分组查询注意力GQA

进一步，在2023年，对MHA和MQA进行折中，诞生了分组查询注意力（Grouped Query Attention，GQA），目标是在保持 MQA 速度的同时实现 MHA 的质量。基本实现是m个Q，n组KV，n可以乘除m，即m//n个Query头共享一组kv，而Kv有多组。

6 Attention计算和存储优化

关于Attention的优化研究是大模型优化的主要方向之一，这里主要介绍两篇最具代表性的，还有其他的优化方向如稀疏注意力等。

6.1 Flash Attention

2022年6月，Stanford团队在论文FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness中提出，优化I/O，主攻训练，针对softmax计算需要所有要素参与的问题，核心是将输入分块并在每个块上执行attention从而减少对高带宽内存(HBM)的读写操作。

6.2 Page Attention

页注意力，2023年，Berkeley团队在论文Efficient Memory Management for Large Language Model Serving with PagedAttention中提出，优化显存，主攻推理服务，核心在非连续的内存空间中存储kv_cache，是vLLM框架的核心技术。具体而言，这是一种受操作系统中虚拟内存和分页经典思想启发的注意力算法。与传统的注意力算法不同，PagedAttention 将每个序列的 KV cache 划分为块，每个块包含固定数量 token 的键和值，并可在内存和外存进行替换。