注意力机制

在传统机器学习中，当数据集非常大时，可能导致性能的大幅下降。对于一张图片来说，我们人首先是看见它最显眼的部分，而背景有时我们会忽略。因此注意力机制就为每个输入项分配一个权重，使得模型去重点关注重要的输入。注意力机制可分为三类：

软注意力机制：对每个输入项分配权重在0-1之间，计算量也会比较大
硬注意力机制：对每个输入项分配权重要么是0，要么是1
自注意力机制：对每个输入项分配的权重取决于输入项之间的相互作用，即通过输入项内部的"表决"来决定应该关注哪些输入项。具有并行计算优势。我们主要研究这种机制，然而首先先介绍一般的注意力机制。

查询、键和值

类似数据库中的元组 $(k_i,v_i)$ ，分别对应着键和值。我们可以对数据库进行查询q，来获取结果。查询可能是精确查询，也可能是模糊查询。我们定义一个包含m个键值对的数据库 $D=\{(k_1,v_1),...,(k_m,v_m)\}$ ，并且定义一个查询q。然后我们可以定义在D上的注意力是：

其中 $\alpha(q,k_i)\in R$ ，是一个标量注意力权重。这个操作本身又叫做注意力汇聚（attention pooling）。“注意力”attention这个词来源于该操作会pay attention to 权重最大的那个 $\alpha$ 。

在深度学习中，我们要求 $\alpha$ 是非负的，并且求和等于1，来表示权重。因此可以这么做：

这样a可以是任意函数。下图是个示意图：

当应用各种不同的查询时，我们希望可视化它对给定的键集的影响。下面的函数可以实现这个功能。它的输入为一个四维张量，前两维分别规定了画布的行数和列数，后两维分别为查询的数量和键的数量。

热力图代码：

def show_heatmaps(matrices, xlabel, ylabel, titles=None, figsize=(2.5, 2.5),
                  cmap='Reds'):
    """显示矩阵热图"""
    d2l.use_svg_display()
    num_rows, num_cols = matrices.shape[0], matrices.shape[1]
    fig, axes = d2l.plt.subplots(num_rows, num_cols, figsize=figsize,
                                 sharex=True, sharey=True, squeeze=False)
    for i, (row_axes, row_matrices) in enumerate(zip(axes, matrices)):
        for j, (ax, matrix) in enumerate(zip(row_axes, row_matrices)):
            pcm = ax.imshow(matrix.detach().numpy(), cmap=cmap)
            if i == num_rows - 1:
                ax.set_xlabel(xlabel)
            if j == 0:
                ax.set_ylabel(ylabel)
            if titles:
                ax.set_title(titles[j])
    fig.colorbar(pcm, ax=axes, shrink=0.6);
    d2l.plt.show()

也就是说，热力图的横坐标是键，纵坐标是查询，显示了键以及查询这一对 $\alpha (k, q)$ 的注意力权重是多少。

注意力汇聚

我们通过回归问题来理解注意力汇聚。假设一个回归方程为：

对于特征和标签为 $x_i,y_i$ 的回归模型， $v_i=y_i$ 是标量， $k_i=x_i$ 是张量，查询 $q$ 是函数f应当评估的新位置。换句话说，我们给定一组点 $x_i,y_i$ ，其中加入了噪声，利用回归模型我们可以得到这组点的曲线。之后给定一个点q，要求求出这个点q对应的函数值f(q)。

由于q离哪个k最近，那么f(q)应该和f(k)最接近，因此那个k的权重应该就越大。因此对于核函数 $\alpha$ 的选择，有以下多种形式：

下面的代码以 $y=2*sin(x)+x$ 为例来说明注意力汇聚，为了简洁，k取成标量形式。

def f(x):
    return 2 * torch.sin(x) + x #回归方程

n = 40
x_train, _ = torch.sort(torch.rand(n) * 5)
y_train = f(x_train) + torch.randn(n) #加入噪声
x_val = torch.arange(0, 5, 0.1)
y_val = f(x_val)

我们将令验证特征x_val作为一个查询，将训练的特征-标签对(x_train, y_train)构成一个键值对，输入到下面函数中，就可以得到x_val在训练模型中的预测值以及注意力权重。

def nadaraya_watson(x_train, y_train, x_val, kernel):
    dists = x_train.reshape((-1, 1)) - x_val.reshape((1, -1)) #k-q
    # Each column/row corresponds to each query/key
    k = kernel(dists).type(torch.float32)
    # Normalization over keys for each query
    attention_w = k / k.sum(0)
    y_hat = y_train@attention_w
    return y_hat, attention_w

因此可以看到，注意力汇聚在本例中，就是回归模型求预测值。可以看到，注意力的权重是根据x_train和x_val之间的距离来确定的，这也符合我们的直觉。

另外在注意力汇聚的过程中，核函数 $\alpha$ 的宽度也会影响结果，例如下面的核：

α (q, k) = e^{- \frac{1}{2 σ^{2}} | | q - k | |^{2}}

$\alpha(q,k)=e^{-\frac{1}{2\sigma ^2}||q-k||^2}$

其中 $\sigma^2$ 就指定了核的大小。实验表明，核越窄，注意力权重范围越窄

注意力评分函数

前文中计算注意力权重是通过距离计算的，然而这计算量有些大，我们改用更简单的方式计算。

上文那个注意力评分函数为：

分析得出，只有第一项有必要保留，因此有：

之后通过softmax操作来获得注意力权重：

便捷函数

下面介绍几个便捷函数：

masked softmax operation

处理不同长度的序列时，我们会需要区分序列的有效长度。我们将超出有效长度的 $v_i$ 设置为0，并将注意力权重设置为一个较大的负数，经过softmax之后就会变为0，使其对梯度和值的贡献消失。

输入：注意力评分矩阵

输出：经过mask和softmax处理后的注意力权重矩阵

使用举例：

masked_softmax(torch.rand(2, 2, 4), torch.tensor([2, 3]))

tensor([[[0.5773, 0.4227, 0.0000, 0.0000],
         [0.5674, 0.4326, 0.0000, 0.0000]],

        [[0.5241, 0.2477, 0.2282, 0.0000],
         [0.3224, 0.2454, 0.4322, 0.0000]]])
         
masked_softmax(torch.rand(2, 2, 4), torch.tensor([[1, 3], [2, 4]]))

tensor([[[1.0000, 0.0000, 0.0000, 0.0000],
         [0.4743, 0.3170, 0.2087, 0.0000]],

        [[0.4712, 0.5288, 0.0000, 0.0000],
         [0.2280, 0.2086, 0.2058, 0.3576]]])

批量矩阵乘法

如果查询和键是以批量形式存在的，如：

那么批量矩阵乘法BMM计算如下：

torch.bmm(Q, K)

下面介绍两个流行的评分函数a：

加性注意力

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

其中 $W_q\in \R^{h\times q}, W_k\in \R^{h\times k}, w_v\in \R^{h}$ 均为可学习参数。h为隐藏单元数，为超参数。

#@save
class AdditiveAttention(nn.Module):
    """加性注意力"""
    def __init__(self, key_size, query_size, num_hiddens, dropout, **kwargs):
        super(AdditiveAttention, self).__init__(**kwargs)
        self.W_k = nn.Linear(key_size, num_hiddens, bias=False)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=False)
        self.w_v = nn.Linear(num_hiddens, 1, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, queries, keys, values, valid_lens):
        #输入queries形状：(一个小批量中的样本数量，一个样本中查询的数量即步数，查询向量的维数q)
        queries, keys = self.W_q(queries), self.W_k(keys)
        # 在维度扩展后，
        # queries的形状：(batch_size，查询的个数，1，num_hiddens)
        # key的形状：(batch_size，1，“键－值”对的个数，num_hiddens)
        # 使用广播方式进行求和
        features = queries.unsqueeze(2) + keys.unsqueeze(1)
        features = torch.tanh(features)
        # self.w_v仅有一个输出，因此从形状中移除最后那个维度。
        # scores的形状：(batch_size，查询的个数，“键-值”对的个数)
        scores = self.w_v(features).squeeze(-1)
        self.attention_weights = masked_softmax(scores, valid_lens)
        # values的形状：(batch_size，“键－值”对的个数，值的维度)
        return torch.bmm(self.dropout(self.attention_weights), values)

在实际输入中，queries和keys, values的形状为(批量大小，步数或词元序列长度，特征大小)。

将attention_weights矩阵输入到热力图函数中，得到：

该矩阵实际为：

tensor([[[0.5000, 0.5000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
          0.0000, 0.0000]],

        [[0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.1667, 0.0000, 0.0000,
          0.0000, 0.0000]]], grad_fn=<SoftmaxBackward0>)

从图上可以看出，这就显示了查询和键对应在注意力权重矩阵里相应位置的大小。颜色越深代表权重越大，模型越关注。而且可以看出，每行加和等于1。

缩放点积注意力

使用点积可以得到计算效率更高的评分函数，但是点积操作要求查询和键具有相同的长度d。假设我们有n个查询和m个键值对，查询和键的长度都为d，值长度为v，则缩放点积注意力对于 $Q\in \R^{n\times d}, K\in \R^{m\times d}, V\in \R^{m\times v}$ 为：

下面的函数以批量的查询和键值对作为输入，输出矩阵 $softmax(\frac{QK^T}{\sqrt{d}})\times V$ 。

class DotProductAttention(nn.Module):
    """缩放点积注意力"""
    def __init__(self, dropout, **kwargs):
        super(DotProductAttention, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout) #暂退层，将其所有元素按dropout的概率归0

    # queries的形状：(batch_size，查询的个数，d)
    # keys的形状：(batch_size，“键－值”对的个数，d)
    # values的形状：(batch_size，“键－值”对的个数，值的维度)
    # valid_lens的形状:(batch_size，)或者(batch_size，查询的个数)
    def forward(self, queries, keys, values, valid_lens=None):
        d = queries.shape[-1]
        # transpose将keys转置
        scores = torch.bmm(queries, keys.transpose(1,2)) / math.sqrt(d)
        #attention_weights形状：(batch_size, 查询的个数， 键值对的个数)
        self.attention_weights = masked_softmax(scores, valid_lens)
        #输出形状：(batch_size, 查询的个数， 值的维度)
        return torch.bmm(self.dropout(self.attention_weights), values)

总结：当查询和键是不同长度的矢量时，可以使用可加性注意力评分函数。当它们的长度相同时，使用缩放的“点－积”注意力评分函数的计算效率更高。

多头注意力机制

在实践中，当给定相同的查询、键和值的集合时，我们希望模型可以基于相同的注意力机制学习到不同的行为，然后将不同的行为作为知识组合起来，捕获序列内各种范围的依赖关系（例如，短距离依赖和长距离依赖关系）。

多头注意力可以在在同一个注意力汇聚模型中，结合不同的查询、键、值的子空间表示。具体而言，我们可以用独立学习得到的ℎ组不同的 线性投影（linear projections）来变换查询、键和值。然后，这ℎ组变换后的查询、键和值将并行地送到注意力汇聚中。最后，将这ℎ个注意力汇聚的输出拼接在一起，并且通过另一个可以学习的线性投影进行变换，以产生最终输出。对于ℎ个注意力汇聚输出，每一个注意力汇聚都被称作一个头（head）。

多头注意力模型的数学表示为：

给定一个查询 $q\in \R^{d_q}$ ，键 $k\in \R^{d_k}$ ，值 $v\in \R^{d_v}$ ，每个注意力头 $h_i, i\in [1..h]$ 为：

h_{i} = f (W_{i}^{(q)} q, W_{i}^{(k)} k, W_{i}^{(v)} v) \in R^{p_{v}}

$h_i=f(W_i^{(q)}q, W_i^{(k)}k, W_i^{(v)}v)\in \R^{p_v}$

其中W均为可学习参数， $W_i^{(q)}\in \R^{p_q\times d_q}$ ， $W_i^{(k)}\in \R^{p_k\times d_k}$ , $W_i^{(v)}\in \R^{p_v\times d_v}$ ， $f$ 为注意力汇聚。例如加性注意力和缩放点积注意力。多头注意力的输出是另一个线性变化：

\begin{matrix} (1) & W_{o} [\begin{array}{ccc} h_{1} \\ ⋮ \\ h_{h} \end{array}] \in R^{p_{o}} \end{matrix}

$W_o \begin{equation} \left[ \begin{array}{ccc} h_1\\ \vdots\\ h_h \end{array} \right] \end{equation} \in \R^{p_o}$

其中可学习参数 $W_o\in \R^{p_o\times hp_v}$ ，注意 $h_i$ 是一个列向量。

下面的示例代码，采用缩放点积注意力作为函数f，并令 $p_q=p_k=p_v=\frac{p_o}{h}$ ，且 $p_o=num\_hiddens$ 。计算 $h_i$ 时对各张量形状进行了改变，实现了并行一起计算。最终输出的形状为 (batch_size, num_queries, num_hiddens)

#@save
class MultiHeadAttention(nn.Module):
    """多头注意力"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs): #注意，英文版的写法是只有一个key_size，因为query_size和value_size的值一般都和key_size相同。key_size表示键特征向量的维度，也就是d_k
        super(MultiHeadAttention, self).__init__(**kwargs)
        self.num_heads = num_heads
        self.attention = d2l.DotProductAttention(dropout)
        self.W_q = nn.Linear(query_size, num_hiddens, bias=bias) #注意，这里原本输出维度是num_hiddens/h，但是由于我们要并行计算h个头，那么输出维度变为num_hiddens
        self.W_k = nn.Linear(key_size, num_hiddens, bias=bias)
        self.W_v = nn.Linear(value_size, num_hiddens, bias=bias)
        self.W_o = nn.Linear(num_hiddens, num_hiddens, bias=bias)

    def forward(self, queries, keys, values, valid_lens):
        # queries，keys，values的形状:
        # (batch_size，查询或者“键－值”对的个数，num_hiddens)
        # valid_lens　的形状:
        # (batch_size，)或(batch_size，查询的个数)
        # 经过变换后，输出的queries，keys，values　的形状:
        # (batch_size*num_heads，查询或者“键－值”对的个数，
        # num_hiddens/num_heads)
        queries = transpose_qkv(self.W_q(queries), self.num_heads)
        keys = transpose_qkv(self.W_k(keys), self.num_heads)
        values = transpose_qkv(self.W_v(values), self.num_heads)

        if valid_lens is not None:
            # 在轴0，将第一项（标量或者矢量）复制num_heads次，
            # 然后如此复制第二项，然后诸如此类。
            valid_lens = torch.repeat_interleave(
                valid_lens, repeats=self.num_heads, dim=0)

        # output的形状:(batch_size*num_heads，查询的个数，
        # num_hiddens/num_heads)
        output = self.attention(queries, keys, values, valid_lens)

        # output_concat的形状:(batch_size，查询的个数，num_hiddens)
        output_concat = transpose_output(output, self.num_heads)
        return self.W_o(output_concat)

为了变换，采取：

#@save
def transpose_qkv(X, num_heads):
    """为了多注意力头的并行计算而变换形状"""
    # 输入X的形状:(batch_size，查询或者“键－值”对的个数，num_hiddens)
    # 输出X的形状:(batch_size，查询或者“键－值”对的个数，num_heads，
    # num_hiddens/num_heads)
    X = X.reshape(X.shape[0], X.shape[1], num_heads, -1)

    # 输出X的形状:(batch_size，num_heads，查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    X = X.permute(0, 2, 1, 3)

    # 最终输出的形状:(batch_size*num_heads,查询或者“键－值”对的个数,
    # num_hiddens/num_heads)
    return X.reshape(-1, X.shape[2], X.shape[3])


#@save
def transpose_output(X, num_heads):
    """逆转transpose_qkv函数的操作"""
    X = X.reshape(-1, num_heads, X.shape[1], X.shape[2])
    X = X.permute(0, 2, 1, 3)
    return X.reshape(X.shape[0], X.shape[1], -1)

多头自注意力机制的详细阐述

参考Multi-headed Self-attention（多头自注意力）机制介绍 - 知乎 (zhihu.com)

如下图，底层有输入序列 $x_1,...,x_T$ ，其中 $x_1$ 可以表示某个句子的第一个单词对应的向量。通过嵌入层embedding，可以得到 $a_1,...,a_T$ 。经过图中右侧的计算，可以得到 $q_1,k_1,v_1$ ，利用 $q_1$ 和 $k_i$ 进行点积，就可以得到标量 $\alpha_{1,i}$ 。将经过softmax后的各个 $\alpha_{1,i}$ 分别与 $v_i$ 相乘，之后求和，就可以得到 $b_1$ 。即：

b_{1} = \sum_{i = 1}^{T} (α (q_{1}, k_{i}) \times v_{i}) = \sum_{i = 1}^{T} (\hat{α_{1, i}} \times v_{i})

$b_1 =\sum_{i=1}^T (\alpha(q_1, k_i)\times v_i) = \sum_{i=1}^T (\hat{\alpha_{1,i}}\times v_i)$

同理，对其他 $x_i$ 进行同样的操作即可，这样可以用并行计算来加快速度。

将输入用矩阵 $I$ 来表示，就可以得到三个矩阵 $Q,K,V$ ，之后经过计算就可以得到输出矩阵 $O$ 。

如果我们将一组 $q_i,k_i,v_i, i=1..T$ 看为一个”头“，那么多头的意思就是，要用多个 $W^Q,W^K,W^V$ 来和 $x_i$ 相乘，得到多组 $q_i,k_i,v_i$ 。

通过多头机制就可以得到多个输出 $b_{head}^i$ ，每个输出的形状和前面的 $b^i$ 是相同的，那么为了得到最终输出 $b^i$ ，我们将 $b^i_{head}$ 依次沿第一个维度进行拼接，然后乘以输出矩阵（线性变换），就可以得到最终输出 $b^i$ 。对于序列中其他输出也采用相同方法，他们都共享网络层的参数。

自注意力

我们可以让一组词元同时充当查询、键和值，这样查询、键和值都来自同一组输入，因此被称为自注意力。

给定一个由词元组成的输入序列 $x_1,...,x_n$ ，其中 $x_i\in \R^d$ ，则它的自注意力输出序列为 $y_1,...,y_n$ ，其中

y_{i} = f (x_{i}, (x_{1}, x_{1}), . . ., (x_{n}, x_{n})) \in R^{d}

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

其中 $x_i$ 表示查询， $(x_j, x_j)$ 表示一个键值对。f为注意力汇聚函数，下面代码基于多头注意力对X完成自注意力的计算，X的形状为（批量大小，时间步数量或词元序列的长度，d），d表示词元是由d维向量表示的。

num_hiddens, num_heads = 100, 5
attention = d2l.MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,
                                   num_hiddens, num_heads, 0.5) #特征维度为num_hiddens，输出维度也为num_hiddens

batch_size, num_queries, valid_lens = 2, 4, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
attention(X, X, X, valid_lens).shape

位置编码

在处理词元序列时，循环神经网络是逐个的重复地处理词元的，而自注意力则因为并行计算而放弃了顺序操作。为了使用序列的顺序信息，通过在输入表示中添加 位置编码（positional encoding）来注入绝对的或相对的位置信息。位置编码可以通过学习得到也可以直接固定得到。接下来描述的是基一种位置编码方法，它是基于正弦函数和余弦函数的固定位置编码：

假设输入表示 $X\in \R^{n\times d}$ 包含一个序列中n个词元的d维嵌入表示。位置编码使用相同形状的位置嵌入矩阵 $P\in \R^{n\times d}$ 输出 $X+P$ ，矩阵第i行、第2j列和2j+1列上的元素为：

P矩阵的行对应词元在序列中的位置，列表示位置编码的不同维度。

#@save
class PositionalEncoding(nn.Module):
    """位置编码"""
    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        # 创建一个足够长的P
        self.P = torch.zeros((1, max_len, num_hiddens))
        X = torch.arange(max_len, dtype=torch.float32).reshape(
            -1, 1) / torch.pow(10000, torch.arange(
            0, num_hiddens, 2, dtype=torch.float32) / num_hiddens)
        self.P[:, :, 0::2] = torch.sin(X)
        self.P[:, :, 1::2] = torch.cos(X)

    def forward(self, X):
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        return self.dropout(X)