动手学深度学习（四十六）——注意力机制

一、注意力机制

  灵长类动物的视觉系统中的视神经接受了大量的感官输入，其内容远远超过了大脑能够完全处理的程度。幸运的是，并非所有刺激的影响都是相等的。意识的聚集和专注使得灵长类动物能够在复杂的视觉环境中将注意力引向感兴趣的物体，例如猎物和天敌。只关注一小部分信息的能力具有进化意义，使人类得以生存和成功。

  自 19 世纪以来，科学家们一直在研究认知神经科学领域的注意力。在本章中，我们将首先回顾一个热门框架，解释如何在视觉场景中展开注意力。受此框架中的 注意力提示（attention cues）的启发，我们将设计能够利用这些注意力线索的模型。特别是 1964 年的 Nadaraya-Waston 核回归（kernel regression）正是具有 注意力机制（attention mechanisms）的机器学习的简单演示。

  然后，我们继续介绍的是注意力函数，它们在深度学习的注意力模型设计中被广泛使用。具体来说，我们将展示如何使用这些函数来设计 Bahdanau 注意力。Bahdanau 注意力是深度学习中的具有突破性价值的注意力模型，它是双向对齐并且可以微分。

  最后，我们将描述仅仅基于了注意力机制的 Transformer 架构，架构中使用的是最新的 多头注意力（multi-head attention）和 自注意力（self-attention）设计。自 2017 年被构想出来，Transformer 一直都普遍存在于现代的深度学习应用中，例如语言、视觉、语音和强化学习领域。

二、注意力提示

  注意力是一种稀缺的资源：此刻你正在阅读这篇blog而忽略了其他的blog。因此，你的注意力是用机会成本（与金钱类似）来支付的。注意力在我们的环境中是稀缺的，而信息不是。在检查视觉场景时，我们的视神经系统收到的信息大约为每秒 $10^8$ 位，远远超过了大脑能够完全处理的水平。幸运的是，我们的祖先已经从经验（也称为数据）中学到 并非所有的感官输入都是一样的。在整个人类历史中，只将注意力引向感兴趣的一小部分信息的能力使我们的大脑能够更明智地分配资源来生存、成长和社交，例如检测天敌、食物和伴侣。

2.1 生物学中的注意力提示

  为了解释我们的注意力是如何在视觉世界中展开的，一个双组件（two-component）的框架已经出现并流行开来。这个框架的出现可以追溯到 19 世纪 90 年代的威廉·詹姆斯，他被认为是 “美国心理学之父” :cite:James.2007。在这个框架中，受试者基于 非自主提示 和 自主提示 有选择地引导注意力的焦点。

  非自主性提示是基于环境中物体的突出性和易见性。想象一下，你面前有五个物品：一份报纸、一篇研究论文、一杯咖啡、一本笔记本和一本下图展示中的书。虽然所有纸制品都是黑白印刷的，但咖啡杯是红色的。换句话说，这种咖啡在这种视觉环境中本质上是突出和显眼的，自动而且非自愿地引起人们的注意。所以你把 fovea（视力最高的黄斑中心）带到咖啡上。

  喝咖啡后，你会变得兴奋并想读书。所以你转过头，重新聚焦你的眼睛，然后看看书。与 之前突出性导致选择会偏向于咖啡不同，在任务依赖案例中选择书本是受到了认知和意识的控制，因此注意力在基于变量选择准则的自主提示去辅助选择时将更为谨慎。受到主体的主观意愿推动，选择的力量也就更强大。

2.2 查询、键和值

  自主的与非自主的提示解释了注意力展开的的方式，受这种提示的启发我们将在下文中描述用于设计注意力机制时的框架，框架中纳入了这两个注意力提示。

  首先，考虑一个相对简单的状况，即只使用非自主提示。要想将选择偏向于感官输入，我们可以简单地使用参数化的全连接层，甚至是非参数化的最大池化层或平均池化层。

  因此，通过包含自主提示将注意力机制与那些全连接层或池化层区别开来。在注意力机制的背景下，我们将自主提示称为 查询（Queries）。给定任何查询，注意力机制通过 注意力池化（attention pooling）将选择偏向于 感官输入（sensory inputs）（例如中间特征表示）。在注意力机制的背景下，这些感官输入被称为 值（Values）。更通俗的解释，每个值都与一个 键（Keys） 配对，这可以想象为该感官输入的非自主提示。我们可以设计注意力池，以便给定的查询（自主提示）可以与键（非自主提示）进行交互，这将指导选择偏向于值（感官输入）。

  请注意，注意力机制的设计有许多替代方案。例如，我们可以设计一个不可微分的注意力模型，该模型可以使用强化学习方法Mnih.Heess.Graves.ea.2014 进行训练。鉴于已经给出的框架在上图中占据主导地位，因此这个框架下的模型将成为本章我们关注的中心。

总结：

• 卷积、全连接、池化层都只考虑不随意线索
• 注意力机制显示的考虑随意线索
  • 每个线索被称之为查询（query）
  • 每个输入是一个值（value）和不随意线索（key）的对
  • 通过注意力池化层来有偏向地选择某些输入

2.3 注意力的可视化

  平均池化层可以被视为输入的加权平均值，其权重是均匀分布的。实际上，注意力池化得到的是加权平均的合计值，其中权重是在给定的查询和不同的键之间计算得出的。

import torch
from d2l import torch as d2l

# 可视化权重：输入为matrices
# 形状（要显示的行数、要显示的列数、查询的数量、键的数量）
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(title(titles[j]))
    fig.colorbar(pcm,ax=axes,shrink=0.6)

# 使用一个简单的例子用于演示，当查询和键相同时注意力权重为1，
attention_weights = torch.eye(10).reshape((1, 1, 10, 10))
show_heatmaps(attention_weights, xlabel='Keys', ylabel='Queries')

三、 注意力池化：Nadaraya-Watson 核回归

  在知道了Query-key-value框架下的注意力机制的主要成分。回顾一下，查询（自主提示）和键（非自主提示）之间的交互形成了 注意力池化（attention pooling）。注意力池化有选择地聚合了值（感官输入）以生成最终的输出1964 年提出的 Nadaraya-Watson 核回归模型是一个简单而完整的示例，可以用于演示具有注意力机制的机器学习。

import torch
from torch import nn
from d2l import torch as d2l

3.1 生成数据集

  简单起见，考虑下面这个回归问题：对于给定的成对的“输入－输出”数据集 $\{(x_1,y_1),\dots ,(x_n,y_n)\}$，如何学习 $f$ 来预测任意新的输入 $x$ 的输出 $\hat{y}=f(x)$？

  根据下面的非线性函数生成一个人工数据集，其中加入的噪声项为 $ϵ$：

$$y_i=2\sin (x_i)+x_i^{0.8}+ϵ,$$

  其中 $ϵ$ 服从均值为 $0$ 和标准差为 $0.5$ 的正态分布。同时生成了 $50$ 个训练样本和 $50$ 个测试样本。为了更好地可视化注意力模式，输入的训练样本将进行排序。

n_train = 50 # the number of train example
x_train,_ = torch.sort(torch.rand(n_train)*5) # the inputs of train example

def f(x):
    return 2*torch.sin(x)+x**0.8

y_train = f(x_train)+torch.normal(0.0,0.5,(n_train,))
x_test = torch.arange(0,5,0.1) 
y_truth = f(x_test) #the real outputs of train exaple
n_test = len(x_test)
n_test

50

3.2 平均池化

  先使用可能是这个世界上“最愚蠢”的估算器来解决回归问题：基于平均池化来计算所有训练样本输出值的平均值：

$$f(x)=\frac{1}{n}\sum _{i=1}^n{y}_i,$$

如下图所示，这个估算器确实不够聪明。

# 绘制全部的训练样本（圆形），不带噪声项的真实数据生成函数f（标记为“truth”）;学习到的预测函数（“Pred”）
def plot_kernel_reg(y_hat):
    d2l.plot(x_test, [y_truth, y_hat], 'x', 'y', legend=['Truth', 'Pred'],
             xlim=[0, 5], ylim=[-1, 5])
    d2l.plt.plot(x_train, y_train, 'o', alpha=0.5);

y_hat = torch.repeat_interleave(y_train.mean(), n_test)
plot_kernel_reg(y_hat)

3.3 非参数的注意力池化

  显然，平均池化忽略了输入 $x_i$。于是 Nadaraya Nadaraya.1964 和 WastonWatson.1964 提出了一个更好的想法，根据输入的位置对输出 $y_i$ 进行加权：

$$f(x)=\sum _{i=1}^n\frac{K(x-x_i)}{{\sum }_{j=1}^{n}K(x-x_j)}{y}_i,$$
其中 $K$ 是 核函数*（kernel）。公式所描述的估计器被称为 Nadaraya-Watson 核回归（Nadaraya-Watson kernel regression）(这家伙不就是加权平均吗！！，衡量与新添加内容距离更近的内容)。在这里我们不会深入讨论核函数的细节。回想一下注意力机制框架，我们可以从注意力机制的角度重写该方程成为一个更加通用的 注意力池化（attention pooling）公式：

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

  其中 $x$ 是查询，$(x_i,y_i)$ 是“键－值”对。比较两个公式可以发现注意力池化是 $y_i$ 的加权平均。将查询 $x$ 和键 $x_i$ 之间的关系建模为 注意力权重（attetnion weight） $\alpha (x,x_i)$，这个权重将被分配给每一个对应值 $y_i$。对于任何查询，模型在所有“键－值”对上的注意力权重都是一个有效的概率分布：它们是非负数的，并且总和为一。

为了更好地理解注意力池化，仅需要考虑一个 高斯核（Gaussian kernel），其定义为（u时距离$x-x_j$）：

$$K(u)=\frac{1}{\sqrt{2\pi }}\exp (-\frac{u^2}{2}).$$

将高斯核代入上面两个公式中就会得出

$$\begin{array}{rl}f(x)& =\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\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}\left(-\frac{1}{2}(x-x_i{)}^2\right)y_i.\end{array}$$

  如果一个键 $x_i$ 越是接近给定的查询 $x$, 那么分配给这个键对应的值 $y_i$ 的 注意力权重就会越大, 也就是 获得了更多的注意力。值得注意的是，Nadaraya-Watson 核回归是一个非参数模型；因此，由其推导出的注意力池化就是 非参数的注意力池化（nonparametric attention pooling）。接下来，我们将基于这个非参数的注意力池化模型来绘制预测结果。结果是预测线是平滑的，并且比平均池化产生的线更接近真实。

# `X_repeat` 的形状: (`n_test`, `n_train`),
# 每一行都包含着相同的测试输入（例如：同样的查询）
X_repeat = x_test.repeat_interleave(n_train).reshape((-1,n_train))
# `x_train` 包含着键。`attention_weights` 的形状：(`n_test`, `n_train`),
# 每一行都包含着要在给定的每个查询的值（`y_train`）之间分配的注意力权重
attention_weights = nn.functional.softmax(-(X_repeat-x_train)**2/2,dim=1)
y_hat = torch.matmul(attention_weights,y_train)
plot_kernel_reg(y_hat)

show_heatmaps(
    attention_weights.unsqueeze(0).unsqueeze(0),
    xlabel='Sorted training inputs', ylabel='Sorted testing inputs')

3.4 带参数的注意力池化

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

$$\begin{array}{rl}f(x)& =\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_i)w{)}^2\right)}{y}_i\\ & =\sum _{i=1}^n\mathrm{s}\mathrm{o}\mathrm{f}\mathrm{t}\mathrm{m}\mathrm{a}\mathrm{x}\left(-\frac{1}{2}((x-x_i)w{)}^2\right)y_i.\end{array}$$

  下面，我们将通过训练这个模型来学习注意力池化的参数。

3.4.1 批量矩阵乘法

  为了更有效地计算小批量数据的注意力，我们可以利用深度学习开发框架中提供的批量矩阵乘法。假设第一个小批量数据包含 $n$ 个矩阵 ${\mathbf{X}}_1,\dots ,{\mathbf{X}}_n$，形状为 $a\times b$，第二个小批量包含 $n$ 个矩阵 ${\mathbf{Y}}_1,\dots ,{\mathbf{Y}}_n$，形状为 $b\times c$。它们的批量矩阵乘法得到 $n$ 个矩阵 ${\mathbf{X}}_1{\mathbf{Y}}_1,\dots ,{\mathbf{X}}_n{\mathbf{Y}}_n$，形状为 $a\times c$。因此，假定两个张量的形状分别是 $(n,a,b)$ 和 $(n,b,c)$ ，它们的批量矩阵乘法输出的形状为 $(n,a,c)$。

X = torch.ones((2,1,4))
Y = torch.ones((2,4,6))
torch.bmm(X,Y).shape

torch.Size([2, 1, 6])

# 在注意力机制中，我们可以使用小批量矩阵乘法来计算小批量数据中值的加权平均
weights = torch.ones((2,10))*0.1
value = torch.arange(20.0).reshape((2,10))
torch.bmm(weights.unsqueeze(1),value.unsqueeze(-1))

tensor([[[ 4.5000]],
        [[14.5000]]])

3.4.2 模型定义

w控制的时高斯核的平滑程度

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)

3.4.2 模型训练

# `X_tile` 的形状: (`n_train`, `n_train`), 每一行都包含着相同的训练输入
X_tile = x_train.repeat((n_train, 1))
# `Y_tile` 的形状: (`n_train`, `n_train`), 每一行都包含着相同的训练输出
Y_tile = y_train.repeat((n_train, 1))
# `keys` 的形状: ('n_train', 'n_train' - 1)
keys = X_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape(
    (n_train, -1))
# `values` 的形状: ('n_train', 'n_train' - 1)
values = Y_tile[(1 - torch.eye(n_train)).type(torch.bool)].reshape(
    (n_train, -1))

net = NWKernelRegression()
loss = nn.MSELoss(reduction='none')
trainer = torch.optim.SGD(net.parameters(), lr=0.5)
animator = d2l.Animator(xlabel='epoch', ylabel='loss', xlim=[1, 5])

for epoch in range(5):
    trainer.zero_grad()
    # 注意：L2 Loss = 1/2 * MSE Loss。
    # PyTorch 的 MSE Loss 与 MXNet 的 L2Loss 差一个 2 的因子，因此被减半。
    l = loss(net(x_train, keys, values), y_train) / 2
    l.sum().backward()
    trainer.step()
    print(f'epoch {epoch + 1}, loss {float(l.sum()):.6f}')
    animator.add(epoch + 1, float(l.sum()))

# `keys` 的形状: (`n_test`, `n_train`), 每一行包含着相同的训练输入（例如：相同的键）
keys = x_train.repeat((n_test, 1))
# `value` 的形状: (`n_test`, `n_train`)
values = y_train.repeat((n_test, 1))
y_hat = net(x_test, keys, values).unsqueeze(1).detach()
plot_kernel_reg(y_hat)

show_heatmaps(
    net.attention_weights.unsqueeze(0).unsqueeze(0),
    xlabel='Sorted training inputs', ylabel='Sorted testing inputs')

四、总结

• 人类的注意力是有限的、宝贵的和稀缺的资源。
• 受试者使用非自主的和自主的提示有选择性地引导注意力。前者基于突出性，后者则依赖于任务。
• 注意力机制与全连接层或者池化层的区别源于增加的自主提示。
• 由于包含了自主提示，注意机制与全连接的层或池层不同。
• 注意力机制通过注意力池化使选择偏向于值（感官输入），其中包含查询（自主提示）和键（非自主提示）。键和值是成对的。
• 我们可以可视化查询和键之间的注意力权重。
• Nadaraya-Watson 核回归是具有注意力机制的机器学习的一个例子。
• Nadaraya-Watson 核回归的注意力池化是对训练数据中输出的加权平均。从注意力的角度来看，分配给每个值的注意力权重取决于将值所对应的键和查询作为输入的函数。
• 注意力池化可以分为非参数型和带参数型。