注意力机制，Transformer相关详解

本文遵循《动手学深度学习pytorch版》的内容组织，从注意力机制开始讲到Transfomer，偏重关键知识理解并附带图解和公式，未加说明时，插图均来自于该书，文本内容较长（9414字），建议收藏慢慢复习。

Transfomer论文原文
李沐：Transformer论文逐段精读【论文精读】

注意力机制，Transformer

注意力机制

注意力是一种稀缺资源。

在注意力机制的背景下，自主性提醒被称为查询（query）。给定任何查询，注意力机制通过注意力汇聚（attention pooling）将选择引导至感官输入（sensory inputs，例如中间特征表示）。在注意力机制中，这些感官输入被称为值（value）。更通俗的解释，每个值都与一个键（key）配对，这可以想象为感官输入的非自主提醒。可以通过设计注意力汇聚的方式，便于给定的查询（自主性提醒）与键（非自主性提醒）进行匹配，这将引导得出最匹配的值（感官输入）。

注意力机制与全连接层或汇聚层的主要区别在于它增加了自主提示。这意味着模型不仅仅是简单地处理所有的输入，而是根据其重要性赋予不同的权重。总的来说，注意力汇聚输出就是值的加权。
在这里插入图片描述

下面的例子以拟合 $y_i = 2 \sin(x_i) + x^{0.8}_i + \epsilon$ 这个非线性函数为例子，训练样本和测试样本分别由以下代码给出：

x_train, _ = torch.sort(torch.rand(50) * 5)  # 排序后的训练样本，x在[0,5)之间的50个随机数
x_test = torch.arange(0, 5, 0.1)  # 测试样本,为0开始，步长0.1的50步数据
def f(x):
    return 2 * torch.sin(x) + x ** 0.8
y_train = f(x_train) + torch.normal(0.0, 0.5, (50,))  # 训练样本的输出，加上了噪声
y_truth = f(x_test)  # 测试样本的真实输出

平均汇聚

在回归问题中，首先尝试使用最简单的估计器来解决问题。基于平均汇聚，可以计算所有训练样本输出值的平均值。显然，平均汇聚忽略了输入 $x_i$ 。导致真实函数f（“Truth”）和预测函数（“Pred”）相差很大。

$\frac{1}{n} \sum_{i=1}^{n} y_i$
在这里插入图片描述

非参数注意力汇聚

现在将输入的位置 $x$ 纳入考虑。在这里，我们有一组键值对 $x_i, y_i)$ ，其中 $x$ 是查询， $x_i, y_i)$ 是键值对。注意力汇聚是通过计算 $y_i$ 的加权平均来实现的。在这个过程中，我们计算查询 $x$ 和键 $x_i$ 之间的关系，并将其建模为注意力权重 $\alpha(x, x_i)$ 。这个权重将被分配给每一个对应的值 $y_i$ 。对于任何查询，注意力权重在所有键值对上构成一个有效的概率分布。这意味着注意力权重是非负的，并且它们的总和为1。

$\sum_{i=1}^{n} \alpha(x, x_i) y_i$

Nadaraya-Watson核回归：考虑下方的汇聚公式，如果一个键 $x_i$ 越是接近给定的查询 $x$ ，那么分配给这个键对应值 $y_i$ 的注意力权重就会越大，也就“获得了更多的注意力”。

Nadaraya-Watson核回归的注意力汇聚是对训练数据中输出的加权平均。从注意力的角度来看，分配给每个值的注意力权重取决于将值所对应的键和查询作为输入的函数。

$\sum_{i=1}^{n} \alpha(x, x_i) y_i = \sum_{i=1}^{n} \frac{\exp\left(-\frac{1}{2}(x - x_i)^2\right)}{\sum_{j=1}^{n} \exp\left(-\frac{1}{2}(x - x_j)^2\right)}y_i = \sum_{i=1}^{n} \text{softmax}\left(-\frac{1}{2}(x - x_i)^2\right) y_i$
在这里插入图片描述

带参数注意力汇聚

非参数的Nadaraya-Watson核回归具有一致性（consistency）的优点：如果有足够的数据，此模型会收敛到最优结果。尽管如此，我们还是可以轻松地将可学习的参数集成到注意力汇聚中。例如，在下面的查询 $x$ 和键 $x_i$ 之间的距离乘以可学习参数 $w$ ：

$\sum_{i=1}^{n} \alpha(x, x_i) y_i = \sum_{i=1}^{n} \frac{\exp\left(-\frac{1}{2}((x - x_i)w)^2\right)}{\sum_{j=1}^{n} \exp\left(-\frac{1}{2}((x - x_j)w)^2\right)} y_i = \sum_{i=1}^{n} \text{softmax}\left(-\frac{1}{2}((x - x_i)w)^2\right) y_i$

模型定义：

class NWKernelRegression(nn.Module):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.w = nn.Parameter(torch.rand((1,), requires_grad=True))

    def forward(self, queries, keys, values):
        # queries和attention_weights的形状为(查询个数，“键－值”对个数)
        queries = queries.repeat_interleave(keys.shape[1]).reshape((-1, keys.shape[1]))
        self.attention_weights = nn.functional.softmax(-((queries - keys) * self.w) ** 2 / 2, dim=1)
        # values的形状为(查询个数，“键－值”对个数)
        return torch.bmm(self.attention_weights.unsqueeze(1), values.unsqueeze(-1)).reshape(-1)

注意力评分函数

注意力评分函数是注意力机制中的一个重要组成部分，它的作用是计算查询和每个键之间的相似性。给定一个查询和一个键，注意力评分函数将它们作为输入，输出一个标量，表示查询和键之间的相似性。这个标量通常被解释为注意力权重，它决定了在注意力汇聚阶段，对应的值应该被赋予多大的权重。以下图中的a表示注意力评分函数。
在这里插入图片描述

计算注意力汇聚的输出为值的加权和

用数学语言描述，假设有一个查询 $\in R^q$ 和 $m$ 个“键－值”对 $k_1, v_1), . . . ,(k_m, v_m)$ ，其中 $k_i \in R^k，v_i \in R^v$ ，注意力汇聚函数 $f$ 就被表示成值的加权和：

$f(q,(k_1, v_1), . . . ,(k_m, v_m)) = \sum_{i=1}^{m} \alpha(q, k_i)v_i \in R^v$

其中查询 $q$ 和键 $ki$ 的注意力权重（标量）是通过注意力评分函数 $a$ 将两个向量映射成标量，再经过softmax运算得到的。选择不同的注意力评分函数 $a$ 会导致不同的注意力汇聚操作。

$\alpha(q, k_i) = \text{softmax}(a(q, k_i)) = \frac{\exp(a(q, k_i))}{\sum_{j=1}^{m} exp(a(q, k_j ))} \in R$

加性注意力：一般来说，当查询和键是不同长度的向量时，可以使用加性注意力作为评分函数。给定查询 $\in R^q$ 和键 $\in R^k$ ，加性注意力的评分函数为：

$w_v^\top \tanh(W_q q + W_k k) \in R$

其中可学习的参数是 $W_q \in R^{h \times q}$ 、 $W_k \in R^{h \times k}$ 和 $w_v \in R^h$ 。

缩放点积注意力：使用点积可以得到计算效率更高的评分函数，但是点积操作要求查询和键具有相同的长度 $d$ 。假设查询和键的所有元素都是独立的随机变量，并且都满足零均值和单位方差，那么两个向量的点积的均值为0，方差为 $d$ 。为确保无论向量长度如何，点积的方差在不考虑向量长度的情况下仍然是1，我们再将点积除以 $\sqrt{d}$ ，则缩放点积注意力评分函数为：

$\frac{q^\top k}{\sqrt{d}}$

使用注意力机制的seq2seq

动机：机器翻译中，每个生成的词可能相关于源句子中不同的词。seq2seq模型中不能对此直接建模，因为解码器默认只关注最后一个隐状态。
在这里插入图片描述

一个带有Bahdanau注意力的循环神经网络编码器-解码器模型

在预测词元时，如果并非所有输入词元都相关，模型将只对齐（或参与）输入序列中与当前预测相关的部分。这是通过将上下文变量视为注意力集中的输出来实现的。这个新的基于注意力的模型与前面的seq2seq模型相同，只不过在任何解码时间步 $t^{'}$ 中的上下文变量 $c$ 会被 $c_{t'}$ 替换。假设输入序列中有 $T$ 个词元，解码时间步 $t^{'}$ 的上下文变量是注意力集中的输出：

$c_{t'} = \sum_{t=1}^{T} \alpha(s_{t'-1}, h_t)h_t$

其中，时间步 $t^{'} - 1$ 时的解码器隐状态 $s_{t'-1}$ 是查询，编码器隐状态 $h_t$ 既是键，也是值，注意力权重是使用加性注意力打分函数计算的。

总的来说，编码器对每次词的输出作为key和value。解码器RNN对上一个词的输出是query，注意力的输出和下一个词的词嵌入合并进入解码器。

多头注意力

动机：对于同一个key，value，query，希望抽取到不同的信息，有一些注意力层关注到短距离关系，有一些关注到长距离关系。有一些类似于“通道”的概念。多头注意力融合了来自于多个注意力汇聚的不同知识，这些知识的不同来源于相同的查询、键和值的不同的子空间表示。
在这里插入图片描述

多头注意力：多个头连结然后线性变换

给定查询 $\in R^{dq}$ 、键 $\in R^{dk}$ 和值 $\in R^{dv}$ ，每个注意力头 $h_i（i = 1, . . . , h）$ 的计算方法为：

$h_i = f(W^{(q)}_i q, W^{(k)}_i k, W^{(v)}_i v) \in R^{pv}$

其中，可学习的参数包括 $W^{(q)}_i \in R^{pq \times dq}$ 、 $W^{(k)}_i \in R^{pk \times dk}$ 和 $W^{(v)}_i \in R^{pv \times dv}$ ，以及代表注意力汇聚的函数 $f$ 。函数 $f$ 可以是 10.3节中的加性注意力和缩放点积注意力。多头注意力的输出需要经过另一个线性转换，它对应着 $h$ 个头连结后的结果，因此其可学习参数是 $W_o \in R^{po \times hpv}$ ：

$W_o \begin{bmatrix} h_1 \\ . \\ . \\ h_h \end{bmatrix} \in R^{po}$

基于这种设计，每个头都可能会关注输入的不同部分，可以表示比简单加权平均值更复杂的函数。

自注意力和位置编码

自注意力：给定一个由词元组成的输入序列 $x_1, . . . , x_n$ ，其中任意 $x_i \in R^d（1 \leq i \leq n）$ 。该序列的自注意力输出为一个长度相同的序列 $y_1, . . . , y_n$ ，其中：

$y_i = f(x_i ,(x_1, x_1), . . . ,(x_n, x_n)) \in R^d$

在这里， $f$ 是一个函数，它接受一个查询和一组键值对，并返回一个输出。在自注意力机制中，查询、键和值都来自同一输入序列。在自注意力机制中，输入序列的每个元素都有机会“关注”序列中的其他元素，以便更好地表示和理解自己。

CNN，RNN，self-attention的比较：目标都是将由 $n$ 个词元组成的序列映射到另一个长度相等的序列，其中的每个输入词元或输出词元都由 $d$ 维向量表示。 $k$ 代表CNN的卷积核大小。注意到，若将CNN和自注意力相比，可以发现自注意力的“感受野”覆盖了整个序列。
在这里插入图片描述

比较卷积神经网络（填充词元被忽略）、循环神经网络和自注意力三种架构：

架构	CNN	RNN	自注意力
计算复杂度	$O(knd^2)$	$O(nd^2)$	$O(n^2d)$
并行度	$O (n)$	$O (1)$	$O (n)$
最大路径长度	$O (n / k)$	$O (n)$	$O (1)$

位置编码：在处理词元序列时，循环神经网络是逐个重复地处理词元的，而自注意力则因为并行计算而放弃了顺序操作。为了使用序列的顺序信息，通过在输入表示中添加位置编码来注入绝对的或相对的位置信息。

为什么是三角函数？一个好的位置编码方案应该满足：1.能够对每一词元（矩阵的行，序列的时间步）输出一个独一无二的编码。2.编码的值应该有界，使得模型能泛化到更长地句子。

假设输入表示 $\in R^{n \times d}$ 包含一个序列中 $n$ 个词元的 $d$ 维嵌入表示。位置编码使用相同形状的位置嵌入矩阵 $P $P \in R^{n \times d}$ 进行加法运算，输出序列 $X + P$ ，这样模型就能根据位置信息进行预测。位置嵌入矩阵 $P$ 的定义如下：

$P_{i,2j} = \sin(i / 10000^{2j/d})$
$P_{i,2j+1} = \cos(i / 10000^{2j/d})$

其中 $i$ 是位置索引， $j$ 是维度索引。该方案之所以有效，是因为对于任意固定的偏移量 $k$ ， $P_{i+k}$ 可以表示为 $P_i$ 的线性函数。因此，模型可以轻易地学习到相对位置，因为对于任意值 $k$ ， $P_{i+k}$ 和 $P_i$ 之间的转换是固定的。
在这里插入图片描述

Transformer

总览：

Transformer模型的主要特点是完全放弃了传统的RNN（循环神经网络）和CNN（卷积神经网络）结构，而是完全依赖于自注意力（Self-Attention）机制来并行地捕捉输入序列中的各个元素之间的依赖关系。这种设计使得Transformer模型在处理长距离依赖问题时具有优势。Transformer模型主要由编码器（Encoder）和解码器（Decoder）两部分组成：
在这里插入图片描述

编码器：

由N个相同的层堆叠而成，每一层包含两个主要部分，一个是多头自注意力（Multi-Head Attention）机制，另一个是位置全连接的前馈网络（Position-wise Feed-Forward Network）。每个部分都有残差连接和层归一化（Layer Normalization）。

解码器：

也是由N个相同的层堆叠而成，但在每一层中包含了一个掩码多头注意力机制，以及一个encoder-decoder attention层（第二层），这一层的K,V来自于编码器，Q来自于解码器。这个设计思想来自于传统RNN的encoder-decoder设计。

传统RNN的循环层所做的变换为

$h_t = f(x_t, h_{t−1})$

然后编码器通过一个函数q将所有隐状态转换为上下文变量

$c = q(h_1, . . . , h_T )$

例如取

$q(h_1, . . . , h_T ) = h_T$

上下文变量仅仅是输入序列在最后时间步的隐状态。实现解码器时，直接使用编码器最后一个时间步的隐状态来初始化解码器的隐状态。

按照这样理解，Transfomer解码器的encoder-decoder attention层实际上也是想让解码器的每一步都能看到编码器所编码的全局信息。在后来的纯解码器结构如GPT中，则丢弃了这一层，因为要求模型自回归地进行生成，要把未来信息“掩码到底”。

以下是作者原文，可以看到encoder-decoder attention的设计就是从 typical encoder-decoder
attention mechanisms启发得到的：

In “encoder-decoder attention” layers, the queries come from the
previous decoder layer, and the memory keys and values come from the
output of the encoder. This allows every position in the decoder to
attend over all positions in the input sequence. This mimics the
typical encoder-decoder attention mechanisms in sequence-to-sequence
models such as [38, 2, 9].

在这里插入图片描述

缩放点积注意力：

给定查询（Query）、键（Key）和值（Value）三个输入，首先计算查询和所有键的点积，然后对结果进行缩放（除以根号下键的维度大小），接着通过softmax函数得到权重，最后用这个权重对值进行加权求和。公式表示为：

$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$

其中， $d_k$ 是键的维度大小。

（掩码）多头注意力：

多头注意力融合了来自于多个注意力汇聚的不同知识，这些知识的不同来源于相同的查询、键和值的不同的子空间表示。其公式表示如下：

$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \text{head}_2, ..., \text{head}_h)W^O$

其中，每个"head"是一个独立的注意力层，计算方式如下：

$\text{head}_i = \text{Attention}(QW^Q_i, KW^K_i, VW^V_i)$

在这里， $Q, K, V$ 分别代表查询、键和值， $W^Q_i, W^K_i, W^V_i$ 和 $W^O$ 是模型需要学习的参数。

在掩码矩阵中，我们将希望模型关注的位置设置为0，将不希望模型关注的位置设置为一个非常大的负数（例如，-10^9）。这是因为，在计算注意力权重时，我们会将掩码矩阵加到 $QK^T$ 上，然后应用softmax函数。在这个过程中，非常大的负数会被映射到接近于0的值，从而使得对应位置的注意力权重接近于0。