李沐动手学深度学习V2-attention注意力机制

一. 注意力提示

1. 生物学中的注意力提示

非自主性提示是基于环境中物体的突出性和易见性。假如你面前有五个物品： 一份报纸、一篇研究论文、一杯咖啡、一本笔记本和一本书，如下图所示。 所有纸制品都是黑白印刷的，但咖啡杯是红色的。 也即是这个咖啡杯在这种视觉环境中是突出和显眼的， 不由自主地引起人们的注意。 所以你把视力最敏锐的地方放到咖啡上。

自主性提示：喝咖啡后，你会变得兴奋并想读书。 所以你转过头，重新聚焦你的眼睛，然后看看书，如下图所示。 与上面非自主性提示突出性导致的选择不同， 此时选择书是受到了认知和意识的控制， 因此注意力在基于自主性提示去辅助选择时将更为谨慎。 受试者的主观意愿推动，选择的力量也就更强大。

2. 查询、键和值

自主性的与非自主性的注意力提示解释了人类的注意力的方式， 下面通过这两种注意力提示用神经网络来设计注意力机制的框架。
首先考虑一个相对简单的状况， 即只使用非自主性提示。 要想将选择偏向于感官输入， 可以简单地使用参数化的全连接层， 甚至是非参数化的最大汇聚层或平均汇聚层。
因此，“是否包含自主性提示”将注意力机制与全连接层或汇聚层区别开来。 在注意力机制的背景下，将自主性提示称为查询（query）。 给定任何查询，注意力机制通过注意力汇聚（attention pooling） 将选择引导至感官输入（sensory inputs，例如中间特征表示）。 在注意力机制中，这些感官输入被称为值（value）。 更通俗的解释每个值都与一个键（key）配对， 可以想象为感官输入的非自主提示。如下图所示可以设计注意力汇聚， 以便给定的查询（自主性提示）可以与键（非自主性提示）进行匹配， 这将引导得出最匹配的值（感官输入）。

二. 注意力权重可视化

平均汇聚层可以被视为输入的加权平均值， 其中各输入的权重是一样的。 实际上注意力汇聚得到的是加权平均的总和值（经过softmax后加权平均的总和值）， 其中权重是在给定的查询和不同的键之间计算得出的。
show_heatmaps()函数作为可视化注意力权重，其输入matrices的形状是 （要显示的行数，要显示的列数，查询的数目，键的数目）。

import torch
from torch import nn
import d2l.torch
def show_heatmap(matrices,xlabel,ylabel,titles=None,figsize=(2.5,2.5),cmap='Reds'):
    """显示矩阵热图"""
    d2l.torch.use_svg_display()
    num_rows,num_colums = matrices.shape[0],matrices.shape[1]
    fig,axes = d2l.torch.plt.subplots(num_rows,num_colums,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)

下面使用一个简单的例子进行演示显示注意力权重的热力图，结果如下图所示：

attention_weights = torch.eye(10).reshape(1,1,10,10)
show_heatmap(attention_weights,xlabel='Keys',ylabel='Queries')

三. 注意力汇聚：Nadaraya-Watson 核回归

查询（自主提示）和键（非自主提示）之间的交互形成了注意力汇聚， 注意力汇聚有选择地聚合了值（感官输入）以生成最终的输出。下面使用Nadaraya-Watson核回归模型，一个简单完整的例子，用于演示具有注意力机制的机器学习。

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 # 训练样本数
x_train,_ = torch.sort(torch.rand(n_train)*5)  # 排序后的训练样本
x_train

输出结果如下所示：
tensor([0.1256, 0.1617, 0.3182, 0.3923, 0.4773, 0.6612, 0.8686, 0.9875, 1.0238,
        1.0555, 1.1690, 1.1852, 1.2633, 1.3428, 1.3704, 1.5106, 1.6027, 1.6256,
        1.6413, 1.8150, 1.9900, 2.1302, 2.1854, 2.2397, 2.3324, 2.3599, 2.5099,
        2.6554, 2.8174, 2.9159, 2.9763, 3.0784, 3.0822, 3.3583, 3.4745, 3.7680,
        3.7723, 3.7987, 3.8181, 3.8351, 4.0055, 4.0567, 4.0728, 4.2019, 4.3304,
        4.3628, 4.6436, 4.6437, 4.9091, 4.9439])

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.0,5.0,0.1)  # 测试样本
y_truth = f(x_test) # 测试样本的真实输出
n_test = len(x_test) # 测试样本数
y_truth

输出结果如下所示：
tensor([0.0000, 0.3582, 0.6733, 0.9727, 1.2593, 1.5332, 1.7938, 2.0402, 2.2712,
        2.4858, 2.6829, 2.8616, 3.0211, 3.1607, 3.2798, 3.3782, 3.4556, 3.5122,
        3.5481, 3.5637, 3.5597, 3.5368, 3.4960, 3.4385, 3.3654, 3.2783, 3.1787,
        3.0683, 2.9489, 2.8223, 2.6905, 2.5554, 2.4191, 2.2835, 2.1508, 2.0227,
        1.9013, 1.7885, 1.6858, 1.5951, 1.5178, 1.4554, 1.4089, 1.3797, 1.3684,
        1.3759, 1.4027, 1.4490, 1.5151, 1.6009])

plot_kernel_reg()函数将绘制所有的训练样本（样本由圆圈表示）， 不带噪声项的真实数据生成函数 （标记为“Truth”）， 以及学习得到的预测函数（标记为“Pred”）。

def plot_kernel_reg(y_hat):
    d2l.torch.plot(x_test,[y_truth,y_hat],'x','y',legend=['Truth','Pred'],xlim=[0,5],ylim=[-1,5])
    d2l.torch.plt.plot(x_train,y_train,'o',alpha=0.5)

2. 平均汇聚

基于平均汇聚来计算所有训练样本输出值的平均值：
$$f(x)=\frac{1}{n}\sum _{i=1}^n{y}_i,$$

如下图所示，真实函数$f$（“Truth”）和预测函数（“Pred”）相差很大。

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

3. 非参数注意力汇聚

平均汇聚忽略了输入$x_i$，Nadaraya 和Watson提出一个更好的想法，根据输入的位置对输出$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）。受此启发我们可以从注意力机制框架的角度重写Nadaraya-Watson核回归公式，成为一个更加通用的注意力汇聚（attention pooling）公式：
$$f(x)=\sum _{i=1}^n\alpha (x,x_i)y_i,$$

其中$x$是查询，$(x_i,y_i)$是键值对。注意力汇聚是$y_i$的加权平均。将查询$x$和键$x_i$之间的关系建模为注意力权重（attention weight）$\alpha (x,x_i)$，这个权重将被分配给每一个对应值$y_i$。对于任何查询，模型在所有键值对注意力权重都是一个有效的概率分布：它们是非负的，并且总和为1。
为了更好地理解注意力汇聚，考虑一个高斯核（Gaussian kernel），其定义为：
$$K(u)=\frac{1}{\sqrt{2\pi }}\exp (-\frac{u^2}{2}).$$

将高斯核代入Nadaraya-Watson核回归公式可以得到：

$$\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{softmax}\left(-\frac{1}{2}(x-x_i{)}^2\right)y_i.\end{array}$$

如果一个键$x_i$越是接近给定的查询$x$，那么分配给这个键对应值$y_i$的注意力权重就会越大，也就“获得了更多的注意力”。 注意：Nadaraya-Watson核回归是一个非参数的注意力汇聚。
下面将基于这个非参数的注意力汇聚模型来绘制预测结果，可以看出新的模型预测线是平滑的，并且比平均汇聚的预测更接近真实，如下图所示。

# X_repeat的形状:(n_test,n_train),
# 每一行都包含着相同的测试输入（例如：同样的查询）
X_repeat = torch.repeat_interleave(x_test,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的每个元素都是值的加权平均值，其中的权重是注意力权重
y_hat = torch.matmul(attention_weights,y_train)
plot_kernel_reg(y_hat)

下面观察注意力的权重，这里测试数据的输入相当于查询，而训练数据的输入相当于键。 因为两个输入都是经过排序的，因此由观察可知“查询-键”对越接近， 注意力汇聚的注意力权重就越高，如下图所示。

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

4.批量矩阵乘法

为了更有效地计算小批量数据的注意力，可以利用pytorch框架中提供的批量矩阵乘法。
假设第一个小批量数据包含$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(size=(2,1,4))
Y = torch.ones(size=(2,4,6))
torch.bmm(X,Y).shape
#输出结果如下：
#torch.Size([2, 1, 6])

在注意力机制的背景中，可以使用小批量矩阵乘法来计算小批量数据中的加权平均值（也即是将注意力attention权重和values值进行矩阵相乘）。

weights = torch.ones(size=(2,10))*0.1
values = torch.arange(20.0).reshape(2,10)
torch.bmm(weights.unsqueeze(1),values.unsqueeze(-1))
'''
输出结果如下：
tensor([[[ 4.5000]],
        [[14.5000]]])
'''

5. 带参数注意力汇聚

非参数的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_j)w{)}^2\right)}{y}_i\\ & =\sum _{i=1}^n\mathrm{softmax}\left(-\frac{1}{2}((x-x_i)w{)}^2\right)y_i.\end{array}$$

下面通过训练这个模型来学习注意力汇聚的参数。

6. 模型定义

基于带参数的注意力汇聚，使用小批量矩阵乘法， 定义Nadaraya-Watson核回归的带参数版本为：

class NWKernelRegresion(nn.Module):
    def __init__(self):
        super(NWKernelRegresion,self).__init__()
        self.w = nn.Parameter(torch.rand((1,),requires_grad=True))
    def forward(self,queries,keys,values):
        # queries和attention_weights的形状为(查询个数，“键－值”对个数)
        queries = torch.repeat_interleave(queries,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)

7. 模型训练

接下来，将训练数据集变换为键和值用于训练注意力模型。 在带参数的注意力汇聚模型中， 任何一个训练样本的输入都会和除自己以外的所有训练样本的“键－值”对进行计算， 从而得到其对应的预测输出。

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

训练带参数的注意力汇聚模型，使用平方损失函数和随机梯度下降，训练结果如下图所示。

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

for epoch in range(5):
    optim.zero_grad()
    y_hat = net(x_train,keys,values)
    l = loss(y_hat,y_train)
    l.sum().backward()
    optim.step()
    print('epoch ',epoch+1,'loss ',l.sum())
    animator.add(epoch+1,float(l.sum()))

训练完带参数的注意力汇聚模型后，把带噪声的训练数据作为(key,value)，把给定预测的x作为query与keys进行计算注意力权重经过softmax后然后再与对应的value相乘进行加权平均后得到给定查询x对应的预测结果，如下图所示预测结果绘制的线不如之前非参数模型的平滑。

# 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).detach()
plot_kernel_reg(y_hat)

查看带参数的模型加入可学习的参数后，每个query在keys中的注意力权重，如下图所示：

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

四. 小结

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

五. 全部代码

import torch
from torch import nn
import d2l.torch


def show_heatmap(matrices, xlabel, ylabel, titles=None, figsize=(2.5, 2.5), cmap='Reds'):
    """显示矩阵热图"""
    d2l.torch.use_svg_display()
    num_rows, num_colums = matrices.shape[0], matrices.shape[1]
    fig, axes = d2l.torch.plt.subplots(num_rows, num_colums, 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)


attention_weights = torch.eye(10).reshape(1, 1, 10, 10)
show_heatmap(attention_weights, xlabel='Keys', ylabel='Queries')
n_train = 50  # 训练样本数
x_train, _ = torch.sort(torch.rand(n_train) * 5)  # 排序后的训练样本
x_train


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.0, 5.0, 0.1)  # 测试样本
y_truth = f(x_test)  # 测试样本的真实输出
n_test = len(x_test)  # 测试样本数
y_truth


def plot_kernel_reg(y_hat):
    d2l.torch.plot(x_test, [y_truth, y_hat], 'x', 'y', legend=['Truth', 'Pred'], xlim=[0, 5], ylim=[-1, 5])
    d2l.torch.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)
# X_repeat的形状:(n_test,n_train),
# 每一行都包含着相同的测试输入（例如：同样的查询）
X_repeat = torch.repeat_interleave(x_test, 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的每个元素都是值的加权平均值，其中的权重是注意力权重
y_hat = torch.matmul(attention_weights, y_train)
plot_kernel_reg(y_hat)
show_heatmap(attention_weights.unsqueeze(0).unsqueeze(0), xlabel='Sorted training inputs',
             ylabel='Sorted testing inputs')
X = torch.ones(size=(2, 1, 4))
Y = torch.ones(size=(2, 4, 6))
torch.bmm(X, Y).shape
weights = torch.ones(size=(2, 10)) * 0.1
values = torch.arange(20.0).reshape(2, 10)
torch.bmm(weights.unsqueeze(1), values.unsqueeze(-1))


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

    def forward(self, queries, keys, values):
        # queries和attention_weights的形状为(查询个数，“键－值”对个数)
        queries = torch.repeat_interleave(queries, 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)


# X_tile的形状:(n_train，n_train)，每一行都包含着相同的训练输入
X_tiles = x_train.repeat((n_train, 1))
# Y_tile的形状:(n_train，n_train)，每一行都包含着相同的训练输出
Y_tiles = y_train.repeat((n_train, 1))
# keys的形状:('n_train'，'n_train'-1)
keys = X_tiles[(1 - torch.eye(n_train)).type(torch.bool)].reshape(n_train, -1)
# values的形状:('n_train'，'n_train'-1)
values = Y_tiles[(1 - torch.eye(n_train)).type(torch.bool)].reshape(n_train, -1)
net = NWKernelRegresion()
loss = nn.MSELoss(reduction='none')
optim = torch.optim.SGD(net.parameters(), lr=0.5)
animator = d2l.torch.Animator(xlabel='epoch', ylabel='loss', xlim=[1, 5])

for epoch in range(5):
    optim.zero_grad()
    y_hat = net(x_train, keys, values)
    l = loss(y_hat, y_train)
    l.sum().backward()
    optim.step()
    print('epoch ', epoch + 1, 'loss ', l.sum())
    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).detach()
plot_kernel_reg(y_hat)
show_heatmap(net.attention_weights.unsqueeze(0).unsqueeze(0), xlabel='Sorted training inputs',
             ylabel='Sorted testing inputs')

六. 相关链接

注意力机制第一篇：李沐动手学深度学习V2-注意力机制
注意力机制第二篇：李沐动手学深度学习V2-注意力评分函数
注意力机制第三篇：李沐动手学深度学习V2-基于注意力机制的seq2seq
注意力机制第四篇：李沐动手学深度学习V2-自注意力机制之位置编码
注意力机制第五篇：李沐动手学深度学习V2-自注意力机制
注意力机制第六篇：李沐动手学深度学习V2-多头注意力机制和代码实现
注意力机制第七篇：李沐动手学深度学习V2-transformer和代码实现