21 KiB
2.1 注意力机制
2.1.1 什么是注意力机制
随着 NLP 从统计机器学习向深度学习迈进，作为 NLP 核心问题的文本表示方法也逐渐从统计学习向深度学习迈进。 正如我们在第一章所介绍的，文本表示从最初的通过统计学习模型进行计算的向量空间模型、语言模型， 通过 Word2Vec 的单层神经网络进入到通过神经网络学习文本表示的时代。 但是，从 CV（Computer Vision，计算机视觉）为起源发展起来的神经网络，其核心架构有三种：
- 全连接网络（FNN），即每一层的神经元都和上下两层的每一个神经元完全连接:
- 卷积神经网络（CNN），即训练参数量远小于全连接层的卷积层来进行特征提取和学习:
- 循环神经网络（RNN），能够使用历史信息作为输入、包含环和自重复的网络:
由于 NLP 任务所需要处理的文本往往是序列，因此专用于处理序列、时序数据的 RNN 往往能够在 NLP 任务上取得最优的效果。 事实上，在注意力机制横空出世之前，RNN 以及 RNN 的衍生架构 LSTM 是 NLP 领域当之无愧的霸主。 例如，我们在第一章讲到过的开创了预训练思想的文本表示模型 ELMo，就是使用的双向 LSTM 作为网络架构。
但 RNN 及 LSTM 虽然具有捕捉时序信息、适合序列生成的优点，却有两个难以弥补的缺陷：
-
序列依序计算的模式能够很好地模拟时序信息，但限制了计算机并行计算的能力。 由于序列需要依次输入、依序计算，GPU 并行计算的能力受到了极大限制，导致 RNN 为基础架构的模型虽然参数量不算特别大， 但计算时间成本却很高；
-
RNN 难以捕捉长序列的相关关系。在 RNN 架构中，距离越远的输入之间的关系就越难被捕捉， 同时 RNN 需要将整个序列读入内存依次计算，也限制了序列的长度。虽然 LSTM 中通过门机制对此进行了一定优化， 但对于较远距离相关关系的捕捉，RNN 依旧是不如人意的。
针对这样的问题，Vaswani 等学者参考了在 CV 领域被提出、被经常融入到 RNN 中使用的注意力机制 （注意，虽然注意力机制在 NLP 被发扬光大，但其确实是在 CV 领域被提出的）， 创新性地搭建了完全由注意力机制构成的神经网络——Transformer，也就是 LLM 的鼻祖及核心架构， 从而让注意力机制一跃成为深度学习最核心的架构之一。
那么，究竟什么是注意力机制（Attention）？
Attention 机制最先源于计算机视觉领域，其核心思想为当我们关注一张图片，我们往往无需看清楚全部内容而仅将注意力集中在重点部分即可。 而在自然语言处理领域，我们往往也可以通过将重点注意力集中在一个或几个 token，从而取得更高效高质的计算效果。
Attention 机制有三个核心变量：Query（查询值）、Key（键值）和 Value（真值）。 我们可以通过一个案例来理解每一个变量所代表的含义。 例如，当我们有一篇新闻报道，我们想要找到这个报道的时间，那么，我们的 Query 可以是类似于“时间”、“日期”一类的向量 （为了便于理解，此处使用文本来表示，但其实际是稠密的向量）， Key 和 Value 会是整个文本。通过对 Query 和 Key 进行运算我们可以得到一个权重， 这个权重其实反映了从 Query 出发，对文本每一个 token 应该分布的注意力相对大小。 通过把权重和 Value 进行运算，得到的最后结果就是从 Query 出发计算整个文本注意力得到的结果。
具体而言，Attention 机制的特点是通过计算 Query 与Key的相关性为真值加权求和， 从而拟合序列中每个词同其他词的相关关系。
2.1.2 深入理解注意力机制
刚刚我们说到，注意力机制有三个核心变量：查询值 Query，键值 Key 和 真值 Value。 接下来我们以字典为例，逐步分析注意力机制的计算公式是如何得到的，从而帮助读者深入理解注意力机制。 首先，我们有这样一个字典：
{
"apple":10,
"banana":5,
"chair":2
}
此时，字典的键就是注意力机制中的键值 Key，而字典的值就是真值 Value。 字典支持我们进行精确的字符串匹配，例如，如果我们想要查找的值也就是查询值 Query 为“apple”， 那么我们可以直接通过将 Query 与 Key 做匹配来得到对应的 Value。
但是，如果我们想要匹配的 Query 是一个包含多个 Key 的概念呢？ 例如，我们想要查找“fruit”，此时，我们应该将 apple 和 banana 都匹配到，但不能匹配到 chair。 因此，我们往往会选择将 Key 对应的 Value 进行组合得到最终的 Value。
例如，当我们的 Query 为“fruit”，我们可以分别给三个 Key 赋予如下的权重：
{
"apple":0.6,
"banana":0.4,
"chair":0
}
那么，我们最终查询到的值应该是：
value = 0.6 * 10 + 0.4 * 5 + 0 * 2 = 8
给不同 Key 所赋予的不同权重，就是我们所说的注意力分数，也就是为了查询到 Query，我们应该赋予给每一个 Key 多少注意力。 但是，如何针对每一个 Query，计算出对应的注意力分数呢？ 从直观上讲，我们可以认为 Key 与 Query 相关性越高，则其所应该赋予的注意力权重就越大。 但是，我们如何能够找到一个合理的、能够计算出正确的注意力分数的方法呢？
在第一章中，我们有提到词向量的概念。通过合理的训练拟合，词向量能够表征语义信息， 从而让语义相近的词在向量空间中距离更近，语义较远的词在向量空间中距离更远。 我们往往用欧式距离来衡量词向量的相似性，但我们同样也可以用点积来进行度量：
v·w = \sum_{i}v_iw_i
根据词向量的定义，语义相似的两个词对应的词向量的点积应该大于0，而语义不相似的词向量点积应该小于0。
那么，我们就可以用点积来计算词之间的相似度。假设我们的 Query 为“fruit”，对应的词向量为 $q$；
我们的 Key 对应的词向量为 k = [v_{apple} v_{banana} v_{chair}],
则我们可以计算 Query 和每一个键的相似程度：
x = qK^T
此处的 K 即为将所有 Key 对应的词向量堆叠形成的矩阵。基于矩阵乘法的定义，x 即为 q 与每一个 k 值的点积。 现在我们得到的 x 即反映了 Query 和每一个 Key 的相似程度，我们再通过一个 Softmax 层将其转化为和为 1 的权重：
softmax(x)_i = \frac{e^{xi}}{\sum_{j}e^{x_j}}
这样，得到的向量就能够反映 Query 和每一个 Key 的相似程度，同时又相加权重为 1，也就是我们的注意力分数了。 最后，我们再将得到的注意力分数和值向量做对应乘积即可。根据上述过程，我们就可以得到注意力机制计算的基本公式：
attention(Q,K,V) = softmax(qK^T)v
不过，此时的值还是一个标量，同时，我们此次只查询了一个 Query。我们可以将值转化为维度为 d_v 的向量，
同时一次性查询多个 Query，同样将多个 Query 对应的词向量堆叠在一起形成矩阵 Q，得到公式：
attention(Q,K,V) = softmax(QK^T)V
目前，我们离标准的注意力公式还差最后一步。
在上一个公式中，如果 Q 和 K 对应的维度 d_k 比较大，
softmax 放缩时就非常容易受影响，使不同值之间的差异较大，从而影响梯度的稳定性。
因此，我们要将 Q 和 K 乘积的结果做一个放缩：
attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d_k}})V
这也就是注意力机制的核心计算公式了。
2.1.3 注意力机制的实现
基于上文，我们可以很简单地使用 Pytorch 来实现注意力机制的代码：
'''注意力计算函数'''
def attention(query, key, value, dropout=None):
'''
args:
query: 查询值矩阵
key: 键值矩阵
value: 真值矩阵
'''
# 获取键向量的维度，键向量的维度和值向量的维度相同
d_k = query.size(-1)
# 计算Q与K的内积并除以根号dk
# transpose——相当于转置
scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
# Softmax
p_attn = scores.softmax(dim=-1)
if dropout is not None:
p_attn = dropout(p_attn)
# 采样
# 根据计算结果对value进行加权求和
return torch.matmul(p_attn, value), p_attn
注意，在上文代码中，我们假设输入的 q、k、v 是已经经过转化的词向量矩阵，也就是公式中的 Q、K、V。 我们仅需要通过上述几行代码，就可以实现核心的 attention 计算。
2.1.4 自注意力
根据上文的分析，我们可以发现，Attention 机制的本质是对两段序列的元素依次进行相似度计算， 寻找出一个序列的每个元素对另一个序列的每个元素的相关度，然后基于相关度进行加权，即分配注意力。 而这两段序列即是我们计算过程中 Q、K、V 的来源。
但是，在我们的实际应用中，我们往往只需要计算 Query 和 Key 之间的注意力结果，很少存在额外的真值 Value。 也就是说，我们其实只需要拟合两个文本序列。 在经典的 Attention 机制中，Q 往往来自于一个序列，K 与 V 来自于另一个序列，都通过参数矩阵计算得到，从而可以拟合这两个序列之间的关系。 例如在 Transformer 的 Decoder 结构中，Q 来自于 Encoder 的输出，K 与 V 来自于 Decoder 的输入， 从而拟合了编码信息与历史信息之间的关系，便于综合这两种信息实现未来的预测。
但在 Transformer 的 Encoder 结构中，使用的是 Attention 机制的变种 —— self-attention （自注意力）机制。
所谓自注意力，即是计算本身序列中每个元素都其他元素的注意力分布，
即在计算过程中，Q、K、V 都由同一个输入通过不同的参数矩阵计算得到。
在 Encoder 中，Q、K、V 分别是输入对参数矩阵 $W_q$、$W_k$、W_v 做积得到，从而拟合输入语句中每一个 token 对其他所有 token 的关系。
通过自注意力机制，我们可以找到一段文本中每一个 token 与其他所有 token 的相关关系大小，从而建模文本之间的依赖关系。 在代码中的实现，self-attention 机制其实是通过给 Q、K、V 的输入传入同一个参数实现的：
# attention 为上文定义的注意力计算函数
attention(x, x, x)
2.1.5 掩码自注意力
掩码自注意力，即 Mask Self Attention，是指使用注意力掩码的自注意力机制。 掩码的作用是遮蔽一些特定位置的 token，模型在学习的过程中，会忽略掉被遮蔽的 token。
使用注意力掩码的核心动机是让模型只能使用历史信息进行预测而不能看到未来信息。 使用注意力机制的 Transformer 模型也是通过类似于 n-gram 的语言模型任务来学习的， 也就是对一个文本序列，不断根据之前的 token 来预测下一个 token，直到将整个文本序列补全。
例如，如果待学习的文本序列是 【BOS】I like you【EOS】，那么，模型会按如下顺序进行预测和学习：
Step 1：输入 【BOS】，输出 I
Step 2：输入 【BOS】I，输出 like
Step 3：输入 【BOS】I like，输出 you
Step 4：输入 【BOS】I like you，输出 【EOS】
理论上来说，只要学习的语料足够多，通过上述的过程，模型可以学会任意一种文本序列的建模方式，也就是可以对任意的文本进行补全。
但是，我们可以发现，上述过程是一个串行的过程，也就是需要先完成 Step 1，才能做 Step 2，接下来逐步完成整个序列的补全。 我们在一开始就说过，Transformer 相对于 RNN 的核心优势之一即在于其可以并行计算，具有更高的计算效率。 如果对于每一个训练语料，模型都需要串行完成上述过程才能完成学习，那么很明显没有做到并行计算，计算效率很低。
针对这个问题，Transformer 就提出了掩码自注意力的方法。掩码自注意力会生成一串掩码，来遮蔽未来信息。 例如，我们待学习的文本序列仍然是 【BOS】I like you【EOS】，我们使用的注意力掩码是【MASK】，那么模型的输入为：
<BOS> 【MASK】【MASK】【MASK】【MASK】
<BOS> I 【MASK】 【MASK】【MASK】
<BOS> I like 【MASK】【MASK】
<BOS> I like you 【MASK】
<BoS> I like you </EOS>
在每一行输入中，模型仍然是只看到前面的 token，预测下一个 token。但是注意，上述输入不再是串行的过程， 而可以一起并行地输入到模型中，模型只需要每一个样本根据未被遮蔽的 token 来预测下一个 token 即可，从而实现了并行的语言模型。
观察上述的掩码，我们可以发现其实则是一个和文本序列等长的上三角矩阵。 我们可以简单地通过创建一个和输入同等长度的上三角矩阵作为注意力掩码，再使用掩码来遮蔽掉输入即可。 也就是说，当输入维度为 （batch_size, seq_len, hidden_size）时，我们的 Mask 矩阵维度一般为 (1, seq_len, seq_len)（通过广播实现同一个 batch 中不同样本的计算）。
在具体实现中，我们通过以下代码生成 Mask 矩阵：
# 创建一个上三角矩阵，用于遮蔽未来信息。
# 先通过 full 函数创建一个 1 * seq_len * seq_len 的矩阵
mask = torch.full((1, args.max_seq_len, args.max_seq_len), float("-inf"))
# triu 函数的功能是创建一个上三角矩阵
mask = torch.triu(mask, diagonal=1)
生成的 Mask 矩阵会是一个上三角矩阵，上三角位置的元素均为 -inf，其他位置的元素置为0。
在注意力计算时，我们会将计算得到的注意力分数与这个掩码做和，再进行 Softmax 操作：
# 此处的 scores 为计算得到的注意力分数，mask 为上文生成的掩码矩阵
scores = scores + mask[:, :seqlen, :seqlen]
scores = F.softmax(scores.float(), dim=-1).type_as(xq)
通过做求和，上三角区域（也就是应该被遮蔽的 token 对应的位置）的注意力分数结果都变成了 -inf，而下三角区域的分数不变。 再做 Softmax 操作，-inf 的值在经过 Softmax 之后会被置为 0，从而忽略了上三角区域计算的注意力分数，从而实现了注意力遮蔽。
2.1.6 多头注意力
注意力机制可以实现并行化与长期依赖关系拟合，但一次注意力计算只能拟合一种相关关系， 单一的注意力机制很难全面拟合语句序列里的相关关系。 因此 Transformer 使用了 Multi-Head attention，也就是多头注意力机制， 即同时对一个语料进行多次注意力计算，每次注意力计算都能拟合不同的关系，将最后的多次结果拼接起来作为最后的输出，即可更全面深入地拟合语言信息。
在原论文中，作者也通过实验证实，多头注意力计算中，每个不同的注意力头能够拟合语句中的不同信息，如下图：
上层与下层分别是两个注意力头对同一段语句序列进行自注意力计算的结果，可以看到， 对于不同的注意力头，能够拟合不同层次的相关信息。 通过多个注意力头同时计算，能够更全面地拟合语句关系。
事实上，所谓的多头注意力机制其实就是将原始的输入序列进行多组的自注意力处理； 然后再将每一组得到的自注意力结果拼接起来，再通过一个线性层进行处理，得到最终的输出。 我们用公式可以表示为：
\mathrm{MultiHead}(Q, K, V) = \mathrm{Concat}(\mathrm{head_1}, ...,
\mathrm{head_h})W^O \\
\text{where}~\mathrm{head_i} = \mathrm{Attention}(QW^Q_i, KW^K_i, VW^V_i)
其最直观的代码实现并不复杂，即 n 个头就有 n 组3个参数矩阵，每一组进行同样的注意力计算， 但由于是不同的参数矩阵从而通过反向传播实现了不同的注意力结果，然后将 n 个结果拼接起来输出即可。
但上述实现时空复杂度均较高，我们可以通过矩阵运算巧妙地实现并行的多头计算， 其核心逻辑在于使用三个组合矩阵来代替了n个参数矩阵的组合，也就是矩阵内积再拼接其实等同于拼接矩阵再内积。 具体实现可以参考下列代码：
import torch.nn as nn
import torch
'''多头自注意力计算模块'''
class MultiHeadAttention(nn.Module):
def __init__(self, args: ModelArgs, is_causal=False):
# 构造函数
# args: 配置对象
super().__init__()
# 隐藏层维度必须是头数的整数倍，因为后面我们会将输入拆成头数个矩阵
assert args.n_embd % args.n_head == 0
# 模型并行处理大小，默认为1。
model_parallel_size = 1
# 本地计算头数，等于总头数除以模型并行处理大小。
self.n_local_heads = args.n_heads // model_parallel_size
# 每个头的维度，等于模型维度除以头的总数。
self.head_dim = args.dim // args.n_heads
# Wq, Wk, Wv 参数矩阵，每个参数矩阵为 n_embd x n_embd
# 这里通过三个组合矩阵来代替了n个参数矩阵的组合，其逻辑在于矩阵内积再拼接其实等同于拼接矩阵再内积，
# 不理解的读者可以自行模拟一下，每一个线性层其实相当于n个参数矩阵的拼接
self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
self.wk = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
self.wv = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
# 输出权重矩阵，维度为 n_embd x n_embd（head_dim = n_embeds / n_heads）
self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
# 注意力的 dropout
self.attn_dropout = nn.Dropout(args.dropout)
# 残差连接的 dropout
self.resid_dropout = nn.Dropout(args.dropout)
# 创建一个上三角矩阵，用于遮蔽未来信息
# 注意，因为是多头注意力，Mask 矩阵比之前我们定义的多一个维度
if is_causal:
mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
mask = torch.triu(mask, diagonal=1)
# 注册为模型的缓冲区
self.register_buffer("mask", mask)
def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
# 获取批次大小和序列长度，[batch_size, seq_len, dim]
bsz, seqlen, _ = q.shape
# 计算查询（Q）、键（K）、值（V）,输入通过参数矩阵层，维度为 (B, T, n_embed) x (n_embed, n_embed) -> (B, T, n_embed)
xq, xk, xv = self.wq(q), self.wk(k), self.wv(v)
# 将 Q、K、V 拆分成多头，维度为 (B, T, n_head, C // n_head)，然后交换维度，变成 (B, n_head, T, C // n_head)
# 因为在注意力计算中我们是取了后两个维度参与计算
# 为什么要先按B*T*n_head*C//n_head展开再互换1、2维度而不是直接按注意力输入展开，是因为view的展开方式是直接把输入全部排开，
# 然后按要求构造，可以发现只有上述操作能够实现我们将每个头对应部分取出来的目标
xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
xk = xk.view(bsz, seqlen, self.n_local_heads, self.head_dim)
xv = xv.view(bsz, seqlen, self.n_local_heads, self.head_dim)
xq = xq.transpose(1, 2)
xk = xk.transpose(1, 2)
xv = xv.transpose(1, 2)
# 注意力计算
# 计算 QK^T / sqrt(d_k)，维度为 (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
# 掩码自注意力必须有注意力掩码
if self.is_causal:
assert hasattr(self, 'mask')
# 这里截取到序列长度，因为有些序列可能比 max_seq_len 短
scores = scores + self.mask[:, :, :seqlen, :seqlen]
# 计算 softmax，维度为 (B, nh, T, T)
scores = F.softmax(scores.float(), dim=-1).type_as(xq)
# 做 Dropout
scores = self.attn_dropout(scores)
# V * Score，维度为(B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
output = torch.matmul(scores, xv)
# 恢复时间维度并合并头。
# 将多头的结果拼接起来, 先交换维度为 (B, T, n_head, C // n_head)，再拼接成 (B, T, n_head * C // n_head)
# contiguous 函数用于重新开辟一块新内存存储，因为Pytorch设置先transpose再view会报错，
# 因为view直接基于底层存储得到，然而transpose并不会改变底层存储，因此需要额外存储
output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
# 最终投影回残差流。
output = self.wo(output)
output = self.resid_dropout(output)
return output
参考文献