注意力机制

注意力机制

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

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

查询、键和值

类似数据库中的元组\((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\)的宽度也会影响结果，例如下面的核：

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

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

注意力评分函数

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

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

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

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

便捷函数

下面介绍几个便捷函数：

1. 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]]])

2. 批量矩阵乘法

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

那么批量矩阵乘法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} \]

其中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\)为注意力汇聚。例如加性注意力和缩放点积注意力。多头注意力的输出是另一个线性变化：

\[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 (\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 \]

其中\(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)

自注意力的另一种理解方式

自注意力机制实际上是想让机器注意到整个输入中不同部分之间的相关性。可参考第四周(2)自注意力机制(Self-Attention) - 知乎 (zhihu.com)

4个输入项，经过自注意力机制，可以输出4个输出项：

如何实现的呢？首先经过线性变换得出查询、键和值：

最关键的是注意力权重矩阵的求法，实际上\(\alpha _{i,j}\)就是表示输入项i和输入项j之间的相关关系，

输出向量为，其中\(\alpha'\)是经过激活函数处理之后的\(\alpha_{i,j}\)：