10. 注意力机制

10.1. 注意力提示

在心理学中,人根据非自主性提示自主性提示有选择的引导注意力的焦点。非自主性提示是基于环境中物体的突出性,自主性提示则基于人的意志力。
卷积层、全连接层和池化层都只考虑偏向于感官输入的非自主性提示,注意力机制则要考虑依赖于任务的意志提示。可以设计注意力汇聚, 以便给定的查询(自主性提示)可以与键(非自主性提示)进行匹配, 引导得出最匹配的值。
qkv.svg

  1. #@save
  2. def show_heatmaps(matrices, xlabel, ylabel, titles=None, figsize=(2.5, 2.5),cmap='Reds'):
  3. """
  4. 显示矩阵热图,用于可视化矩阵权重。仅当查询和键相同时,注意力权重为1,否则为0。
  5. matrices的形状是 (要显示的图片行数,要显示的图片列数,查询的数目,键的数目)。
  6. """
  7. d2l.use_svg_display()
  8. num_rows, num_cols = matrices.shape[0], matrices.shape[1]
  9. fig, axes = d2l.plt.subplots(num_rows, num_cols, figsize=figsize,sharex=True, sharey=True, squeeze=False)
  10. for i, (row_axes, row_matrices) in enumerate(zip(axes, matrices)):
  11. for j, (ax, matrix) in enumerate(zip(row_axes, row_matrices)):
  12. pcm = ax.imshow(matrix.detach().numpy(), cmap=cmap)
  13. if i == num_rows - 1:
  14. ax.set_xlabel(xlabel)
  15. if j == 0:
  16. ax.set_ylabel(ylabel)
  17. if titles:
  18. ax.set_title(titles[j])
  19. fig.colorbar(pcm, ax=axes, shrink=0.6);
  20. attention_weights = torch.eye(10).reshape((1, 1, 10, 10)) # torch.eye(n):生成n行n列的单位矩阵
  21. show_heatmaps(attention_weights, xlabel='Keys', ylabel='Queries')

output_attention-cues_054b1a_30_0.svg

10.2. 注意力汇聚:Nadaraya-Watson 核回归

给定成对的“输入-输出”数据集 10. 注意力机制 - 图3%2C%E2%80%A6%2C(x%7Bn%7D%2Cy%7Bn%7D)%7D#card=math&code=%7B%28x%7B1%7D%2Cy%7B1%7D%29%2C%E2%80%A6%2C%28x%7Bn%7D%2Cy%7Bn%7D%29%7D&id=y0vmX), 如何学习10. 注意力机制 - 图4来预测任意新输入10. 注意力机制 - 图5的输出10. 注意力机制 - 图6#card=math&code=%5Chat%7By%7D%3Df%28x%29&id=GLnUs)?接下来探讨如何解决这个问题。
首先,根据非线性函数10. 注意力机制 - 图7%2Bx%5E%7B0.8%7D%7Bi%7D%2B%CF%B5%2C#card=math&code=y%7Bi%7D%3D2sin%28x%7Bi%7D%29%2Bx%5E%7B0.8%7D%7Bi%7D%2B%CF%B5%2C&id=EjOPK)生成一个人工数据集, 噪声项ϵ服从均值为0和标准差为0.5的正态分布。

  1. n_train = 50 # 训练样本数
  2. x_train, _ = torch.sort(torch.rand(n_train) * 5) # 排序后的训练样本。torch.rand()从均匀分布[0,1]中返回随机数。x_train存储排序后的列表,_存储每个元素在原列表中的索引位置
  3. def f(x):
  4. return 2 * torch.sin(x) + x**0.8
  5. y_train = f(x_train) + torch.normal(0.0, 0.5, (n_train,)) # 训练样本的输出(加了噪声)
  6. x_test = torch.arange(0, 5, 0.1) # 测试样本
  7. y_truth = f(x_test) # 测试样本的真实输出。(不加噪声)
  8. def plot_kernel_reg(y_hat):
  9. d2l.plot(x_test, [y_truth, y_hat], 'x', 'y', legend=['Truth', 'Pred'],xlim=[0, 5], ylim=[-1, 5]) # 生成数据和预测数据。线
  10. d2l.plt.plot(x_train, y_train, 'o', alpha=0.5); # 绘制训练样本。散点

10.2.1 非参数注意力汇聚

下面的方法根据输入的位置对输出10. 注意力机制 - 图8进行加权:
10. 注意力机制 - 图9%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%5Cfrac%7BK(x-x%7Bi%7D)%7D%7B%5Csum%7Bj%3D1%7D%5E%7Bn%7DK(x-x%7Bj%7D)%7Dy%7Bi%7D%0A#card=math&code=f%28x%29%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%5Cfrac%7BK%28x-x%7Bi%7D%29%7D%7B%5Csum%7Bj%3D1%7D%5E%7Bn%7DK%28x-x%7Bj%7D%29%7Dy%7Bi%7D%0A&id=aqgpL)
其中K是(kernel)。上式估计器被称为 Nadaraya-Watson核回归(Nadaraya-Watson kernel regression)。下面按注意力框架将其重写为通用的注意力汇聚公式:
10. 注意力机制 - 图10%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%CE%B1(x%2Cx%7Bi%7D)y%7Bi%7D%0A#card=math&code=f%28x%29%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%CE%B1%28x%2Cx%7Bi%7D%29y%7Bi%7D%0A&id=edb7o)
其中10. 注意力机制 - 图11是查询,10. 注意力机制 - 图12#card=math&code=%28x%7Bi%7D%2Cy%7Bi%7D%29&id=BmMhX)是键值对。将查询10. 注意力机制 - 图13和键10. 注意力机制 - 图14之间的关系建模为 注意力权重(attention weight)。我们考虑一个高斯核(Gaussian kernel),其定义为:
10. 注意力机制 - 图15%3D%5Cfrac%7B1%7D%7B%5Csqrt%7B2%CF%80%7D%7D%5Ctext%7Bexp%7D(-%5Cfrac%7Bu%5E%7B2%7D%7D%7B2%7D)%0A#card=math&code=K%28u%29%3D%5Cfrac%7B1%7D%7B%5Csqrt%7B2%CF%80%7D%7D%5Ctext%7Bexp%7D%28-%5Cfrac%7Bu%5E%7B2%7D%7D%7B2%7D%29%0A&id=ckxEG)
将高斯核带入(1)、(2)可得:
10. 注意力机制 - 图16%20%26%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%20%5Calpha%5Cleft(x%2C%20x%7Bi%7D%5Cright)%20y%7Bi%7D%20%5C%5C%0A%26%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%20%5Cfrac%7B%5Cexp%20%5Cleft(-%5Cfrac%7B1%7D%7B2%7D%5Cleft(x-x%7Bi%7D%5Cright)%5E%7B2%7D%5Cright)%7D%7B%5Csum%7Bj%3D1%7D%5E%7Bn%7D%20%5Cexp%20%5Cleft(-%5Cfrac%7B1%7D%7B2%7D%5Cleft(x-x%7Bj%7D%5Cright)%5E%7B2%7D%5Cright)%7D%20y%7Bi%7D%20%5C%5C%0A%26%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%20%5Coperatorname%7Bsoftmax%7D%5Cleft(-%5Cfrac%7B1%7D%7B2%7D%5Cleft(x-x%7Bi%7D%5Cright)%5E%7B2%7D%5Cright)%20y%7Bi%7D%0A%5Cend%7Baligned%7D%0A#card=math&code=%5Cbegin%7Baligned%7D%0Af%28x%29%20%26%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%20%5Calpha%5Cleft%28x%2C%20x%7Bi%7D%5Cright%29%20y%7Bi%7D%20%5C%5C%0A%26%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%20%5Cfrac%7B%5Cexp%20%5Cleft%28-%5Cfrac%7B1%7D%7B2%7D%5Cleft%28x-x%7Bi%7D%5Cright%29%5E%7B2%7D%5Cright%29%7D%7B%5Csum%7Bj%3D1%7D%5E%7Bn%7D%20%5Cexp%20%5Cleft%28-%5Cfrac%7B1%7D%7B2%7D%5Cleft%28x-x%7Bj%7D%5Cright%29%5E%7B2%7D%5Cright%29%7D%20y%7Bi%7D%20%5C%5C%0A%26%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%20%5Coperatorname%7Bsoftmax%7D%5Cleft%28-%5Cfrac%7B1%7D%7B2%7D%5Cleft%28x-x%7Bi%7D%5Cright%29%5E%7B2%7D%5Cright%29%20y%7Bi%7D%0A%5Cend%7Baligned%7D%0A&id=tbNg9)
如果一个键10. 注意力机制 - 图17越是接近给定的查询10. 注意力机制 - 图18, 那么分配给这个键对应值10. 注意力机制 - 图19的注意力权重就会越大, 也就“获得了更多的注意力”。值得注意的是,该回归是一个非参数模型。下面基于这个模型来绘制预测结果:

  1. # X_repeat的形状:(n_test,n_train),
  2. # 每一行都包含着相同的测试输入(例如:同样的查询),而每一列是一组
  3. X_repeat = x_test.repeat_interleave(n_train).reshape((-1, n_train))#重复n_train
  4. # x_train包含着键。attention_weights的形状:(n_test,n_train),
  5. # 每一行都包含着要在给定的每个查询的值(y_train)之间分配的注意力权重
  6. attention_weights = nn.functional.softmax(-(X_repeat - x_train)**2 / 2, dim=1)
  7. # y_hat的每个元素都是值的加权平均值,其中的权重是注意力权重
  8. y_hat = torch.matmul(attention_weights, y_train)
  9. plot_kernel_reg(y_hat)

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

  1. d2l.show_heatmaps(attention_weights.unsqueeze(0).unsqueeze(0),
  2. xlabel='Sorted training inputs',
  3. ylabel='Sorted testing inputs')
  4. # unsequeeze(i):在第i维增加一个维度。这里将2维变为了4维

output_nadaraya-waston_736177_56_0.svg
非参数的Nadaraya-Watson核回归具有一致性(consistency)的优点: 如果有足够的数据,此模型会收敛到最优结果。 尽管如此,我们还是可以轻松地将可学习的参数集成到注意力汇聚中。

10.2.2. 带参数注意力汇聚

唯一不同是在查询10. 注意力机制 - 图22和键10. 注意力机制 - 图23之间的距离乘以可学习参数10. 注意力机制 - 图24
10. 注意力机制 - 图25%20%26%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%20%5Calpha%5Cleft(x%2C%20x%7Bi%7D%5Cright)%20y%7Bi%7D%20%5C%5C%0A%26%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%20%5Cfrac%7B%5Cexp%20%5Cleft(-%5Cfrac%7B1%7D%7B2%7D%5Cleft(%20%5Cleft(x-x%7Bi%7D%5Cright)w%5Cright)%5E%7B2%7D%5Cright)%7D%7B%5Csum%7Bj%3D1%7D%5E%7Bn%7D%20%5Cexp%20%5Cleft(-%5Cfrac%7B1%7D%7B2%7D%5Cleft(%20%5Cleft(x-x%7Bj%7D%5Cright)w%5Cright)%5E%7B2%7D%5Cright)%7D%20y%7Bi%7D%20%5C%5C%0A%26%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%20%5Coperatorname%7Bsoftmax%7D%5Cleft(-%5Cfrac%7B1%7D%7B2%7D%5Cleft(%20%5Cleft(x-x%7Bi%7D%5Cright)w%5Cright)%5E%7B2%7D%5Cright)%20y%7Bi%7D%0A%5Cend%7Baligned%7D%0A#card=math&code=%5Cbegin%7Baligned%7D%0Af%28x%29%20%26%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%20%5Calpha%5Cleft%28x%2C%20x%7Bi%7D%5Cright%29%20y%7Bi%7D%20%5C%5C%0A%26%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%20%5Cfrac%7B%5Cexp%20%5Cleft%28-%5Cfrac%7B1%7D%7B2%7D%5Cleft%28%20%5Cleft%28x-x%7Bi%7D%5Cright%29w%5Cright%29%5E%7B2%7D%5Cright%29%7D%7B%5Csum%7Bj%3D1%7D%5E%7Bn%7D%20%5Cexp%20%5Cleft%28-%5Cfrac%7B1%7D%7B2%7D%5Cleft%28%20%5Cleft%28x-x%7Bj%7D%5Cright%29w%5Cright%29%5E%7B2%7D%5Cright%29%7D%20y%7Bi%7D%20%5C%5C%0A%26%3D%5Csum%7Bi%3D1%7D%5E%7Bn%7D%20%5Coperatorname%7Bsoftmax%7D%5Cleft%28-%5Cfrac%7B1%7D%7B2%7D%5Cleft%28%20%5Cleft%28x-x%7Bi%7D%5Cright%29w%5Cright%29%5E%7B2%7D%5Cright%29%20y%7Bi%7D%0A%5Cend%7Baligned%7D%0A&id=o7Tez)
为了更有效地计算小批量数据的注意力, 可以利用批量矩阵乘法。假设第一个小批量数据包含n个矩阵10. 注意力机制 - 图26, 每个矩阵形状为a×b, 第二个小批量包含n个矩阵10. 注意力机制 - 图27, 每个矩阵形状为b×c。 它们的批量矩阵乘法得到n个矩阵10. 注意力机制 - 图28, 每个矩阵形状为a×c。 因此,假定两个张量的形状分别是(n,a,b)和(n,b,c), 它们的批量矩阵乘法输出的形状为(n,a,c)。批量矩阵乘法使用torch.bmm(X, Y)

  1. weights = torch.ones((2, 10)) * 0.1
  2. values = torch.arange(20.0).reshape((2, 10))
  3. torch.bmm(weights.unsqueeze(1), values.unsqueeze(-1)) # [2, 10] -> ->[2, 1, 10]; [2, 10] -> [2, 10, 1]
  4. """
  5. tensor([[[ 4.5000]],
  6. [[14.5000]]])
  7. """
  1. class NWKernelRegression(nn.Module): # 定义模型
  2. def __init__(self, **kwargs):
  3. super().__init__(**kwargs)
  4. self.w = nn.Parameter(torch.rand((1,), requires_grad=True))
  5. def forward(self, queries, keys, values):
  6. # queries和attention_weights的形状为(查询个数,“键-值”对个数)
  7. queries = queries.repeat_interleave(keys.shape[1]).reshape((-1, keys.shape[1]))
  8. self.attention_weights = nn.functional.softmax(
  9. -((queries - keys) * self.w)**2 / 2, dim=1)
  10. # values的形状为(查询个数,“键-值”对个数)
  11. return torch.bmm(self.attention_weights.unsqueeze(1),
  12. values.unsqueeze(-1)).reshape(-1)

在带参数的注意力汇聚模型中, 任何一个训练样本的输入都会和除自己以外的所有训练样本的“键-值”对进行计算, 从而得到其对应的预测输出。

# 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()
    l = loss(net(x_train, keys, values), y_train)
    l.sum().backward()
    trainer.step()
    print(f'epoch {epoch + 1}, loss {float(l.sum()):.6f}')
    animator.add(epoch + 1, float(l.sum()))

output_nadaraya-waston_736177_116_0.svg
我们发现: 在尝试拟合带噪声的训练数据时, 预测结果绘制的线不如之前非参数模型的平滑。

# 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)

output_nadaraya-waston_736177_128_0.svg
参数W实际上是控制了高斯核的窗口大小,当窗口较大时,曲线较平滑;当窗口较小时,曲线会变的不平滑。不平滑的好处是曲线可以拟合的更好,坏处是可能产生一些抖动,存在部分的过拟合。
output_nadaraya-waston_736177_140_0-16520899035831.svg
与非参数的注意力汇聚模型相比, 带参数的模型加入可学习的参数后, 曲线在注意力权重较大的区域变得更不平滑。从上图可以看到,预测注意力权重与实际权重更接近了。

10.3. 注意力评分函数

10.2中的高斯核指数部分就是一种注意力评分函数(attention scoring function)。下图说明了 如何将注意力汇聚的输出计算成为值的加权和, 其中a表示注意力评分函数。 由于注意力权重是概率分布, 因此加权和其本质上是加权平均值。选择不同的注意力评分函数a会导致不同的注意力汇聚操作。
image-20220430172444169.png
假设有一个查询 10. 注意力机制 - 图33 和 m 个“键 一 值”对 10. 注意力机制 - 图34%2C%20%5Cldots%2C%5Cleft(%5Cmathbf%7Bk%7D%7Bm%7D%2C%20%5Cmathbf%7Bv%7D%7Bm%7D%5Cright)#card=math&code=%5Cleft%28%5Cmathbf%7Bk%7D%7B1%7D%2C%20%5Cmathbf%7Bv%7D%7B1%7D%5Cright%29%2C%20%5Cldots%2C%5Cleft%28%5Cmathbf%7Bk%7D%7Bm%7D%2C%20%5Cmathbf%7Bv%7D%7Bm%7D%5Cright%29&id=gt8rg) , 其中 10. 注意力机制 - 图35 。 注意力汇聚函数 10. 注意力机制 - 图36 就被表示成 值的加权和:
10. 注意力机制 - 图37%2C%20%5Cldots%2C%5Cleft(%5Cmathbf%7Bk%7D%7Bm%7D%2C%20%5Cmathbf%7Bv%7D%7Bm%7D%5Cright)%5Cright)%3D%5Csum%7Bi%3D1%7D%5E%7Bm%7D%20%5Calpha%5Cleft(%5Cmathbf%7Bq%7D%2C%20%5Cmathbf%7Bk%7D%7Bi%7D%5Cright)%20%5Cmathbf%7Bv%7D%7Bi%7D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bv%7D%5C%5C%0A%5Calpha%5Cleft(%5Cmathbf%7Bq%7D%2C%20%5Cmathbf%7Bk%7D%7Bi%7D%5Cright)%3D%5Coperatorname%7Bsoftmax%7D%5Cleft(a%5Cleft(%5Cmathbf%7Bq%7D%2C%20%5Cmathbf%7Bk%7D%7Bi%7D%5Cright)%5Cright)%3D%5Cfrac%7B%5Cexp%20%5Cleft(a%5Cleft(%5Cmathbf%7Bq%7D%2C%20%5Cmathbf%7Bk%7D%7Bi%7D%5Cright)%5Cright)%7D%7B%5Csum%7Bj%3D1%7D%5E%7Bm%7D%20%5Cexp%20%5Cleft(a%5Cleft(%5Cmathbf%7Bq%7D%2C%20%5Cmathbf%7Bk%7D%7Bj%7D%5Cright)%5Cright)%7D%20%5Cin%20%5Cmathbb%7BR%7D%0A#card=math&code=f%5Cleft%28%5Cmathbf%7Bq%7D%2C%5Cleft%28%5Cmathbf%7Bk%7D%7B1%7D%2C%20%5Cmathbf%7Bv%7D%7B1%7D%5Cright%29%2C%20%5Cldots%2C%5Cleft%28%5Cmathbf%7Bk%7D%7Bm%7D%2C%20%5Cmathbf%7Bv%7D%7Bm%7D%5Cright%29%5Cright%29%3D%5Csum%7Bi%3D1%7D%5E%7Bm%7D%20%5Calpha%5Cleft%28%5Cmathbf%7Bq%7D%2C%20%5Cmathbf%7Bk%7D%7Bi%7D%5Cright%29%20%5Cmathbf%7Bv%7D%7Bi%7D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bv%7D%5C%5C%0A%5Calpha%5Cleft%28%5Cmathbf%7Bq%7D%2C%20%5Cmathbf%7Bk%7D%7Bi%7D%5Cright%29%3D%5Coperatorname%7Bsoftmax%7D%5Cleft%28a%5Cleft%28%5Cmathbf%7Bq%7D%2C%20%5Cmathbf%7Bk%7D%7Bi%7D%5Cright%29%5Cright%29%3D%5Cfrac%7B%5Cexp%20%5Cleft%28a%5Cleft%28%5Cmathbf%7Bq%7D%2C%20%5Cmathbf%7Bk%7D%7Bi%7D%5Cright%29%5Cright%29%7D%7B%5Csum%7Bj%3D1%7D%5E%7Bm%7D%20%5Cexp%20%5Cleft%28a%5Cleft%28%5Cmathbf%7Bq%7D%2C%20%5Cmathbf%7Bk%7D%7Bj%7D%5Cright%29%5Cright%29%7D%20%5Cin%20%5Cmathbb%7BR%7D%0A&id=aYrrf)

10.3.1. 掩蔽softmax操作

softmax用于输出一个概率分布作为注意力权重。但在计算时,应当规避填充词元等无意义的特殊词元, 可以指定一个有效序列长度(即词元的个数), 以便在计算softmax时过滤掉超出指定范围的位置。下面的操作方式是将任何超出有效长度的位置都置为一个非常小的值。

#@save
def masked_softmax(X, valid_lens):
    """通过在最后一个轴上掩蔽元素来执行softmax操作"""
    # X:3D张量,valid_lens:1D或2D张量
    if valid_lens is None:
        return nn.functional.softmax(X, dim=-1)    # 对每一行执行softmax计算
    else:
        shape = X.shape
        if valid_lens.dim() == 1: 
            valid_lens = torch.repeat_interleave(valid_lens, shape[1])
        else:
            valid_lens = valid_lens.reshape(-1)    # 转化为一维向量
        # 最后一轴上被掩蔽的元素使用一个非常大的负值替换,从而其softmax输出为0
        X = d2l.sequence_mask(X.reshape(-1, shape[-1]), valid_lens,value=-1e6)
        return nn.functional.softmax(X.reshape(shape), dim=-1)

masked_softmax(torch.rand(2, 2, 4), torch.tensor([2, 3]))
"""
tensor([[[0.5301, 0.4699, 0.0000, 0.0000],
         [0.5574, 0.4426, 0.0000, 0.0000]],

        [[0.3333, 0.2810, 0.3857, 0.0000],
         [0.3744, 0.2307, 0.3949, 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.3439, 0.4590, 0.1970, 0.0000]],

        [[0.5145, 0.4855, 0.0000, 0.0000],
         [0.1810, 0.2662, 0.2728, 0.2800]]])
"""

10.3.2. 加性注意力

当查询和键是不同长度的矢量时可以使用加性注意力作为评分函数。给定查询 10. 注意力机制 - 图38 和 键 10. 注意力机制 - 图39 , 加性注意力(additive attention)的评分函数 为:
10. 注意力机制 - 图40%3D%5Cmathbf%7Bw%7D%7Bv%7D%5E%7B%5Ctop%7D%20%5Ctanh%20%5Cleft(%5Cmathbf%7BW%7D%7Bq%7D%20%5Cmathbf%7Bq%7D%2B%5Cmathbf%7BW%7D%7Bk%7D%20%5Cmathbf%7Bk%7D%5Cright)%20%5Cin%20%5Cmathbb%7BR%7D%0A#card=math&code=a%28%5Cmathbf%7Bq%7D%2C%20%5Cmathbf%7Bk%7D%29%3D%5Cmathbf%7Bw%7D%7Bv%7D%5E%7B%5Ctop%7D%20%5Ctanh%20%5Cleft%28%5Cmathbf%7BW%7D%7Bq%7D%20%5Cmathbf%7Bq%7D%2B%5Cmathbf%7BW%7D%7Bk%7D%20%5Cmathbf%7Bk%7D%5Cright%29%20%5Cin%20%5Cmathbb%7BR%7D%0A&id=nOBBc)
其中可学习的参数是 10. 注意力机制 - 图4110. 注意力机制 - 图4210. 注意力机制 - 图43

#@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, keys = self.W_q(queries), self.W_k(keys)
        # 由于queries和keys中间维度不同,因此需要扩展维度以进行叠加
        # queries的形状:(batch_size,查询的个数,1,num_hidden)
        # 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)

# 演示: 查询、键和值的形状为(批量大小,步数或词元序列长度,特征大小)。
# normal中前两个参数是均值与方差
queries, keys = torch.normal(0, 1, (2, 1, 20)), torch.ones((2, 10, 2))
# values的小批量,两个值矩阵是相同的
values = torch.arange(40, dtype=torch.float32).reshape(1, 10, 4).repeat(2, 1, 1)    # 维度为(2,10,4)
valid_lens = torch.tensor([2, 6])

attention=AdditiveAttention(key_size=2,query_size=20,num_hiddens=8,dropout=0.1)
attention.eval()
attention(queries, keys, values, valid_lens)
"""
注意力汇聚输出的形状为(批量大小,查询的步数,值的维度)
tensor([[[ 2.0000,  3.0000,  4.0000,  5.0000]],
        [[10.0000, 11.0000, 12.0000, 13.0000]]], grad_fn=<BmmBackward0>)
"""

尽管加性注意力包含了可学习的参数,但由于本例子中每个键都是相同的, 所以注意力权重是均匀的,由指定的有效长度决定。

d2l.show_heatmaps(attention.attention_weights.reshape((1, 1, 2, 10)),xlabel='Keys', ylabel='Queries')

output_attention-scoring-functions_2a8fdc_78_0.svg

10.3.3. 缩放点积注意力

点积可以得到计算效率更高的评分函数, 但是要求查询和键具有相同的长度d。 假设查询和键的所有元素都是独立的随机变量, 并且都满足零均值和单位方差, 那么两个向量的点积的均值为0,方差为d。 为确保无论向量长度如何, 点积的方差在不考虑向量长度的情况下仍然是1, 我们将点积除以√d,基于n个查询和m个键-值对计算注意力, 其中查询和键的长度为d,值的长度为v。 查询10. 注意力机制 - 图45、 键10. 注意力机制 - 图46和 值10. 注意力机制 - 图47的缩放点积注意力是:
10. 注意力机制 - 图48%20%5Cmathbf%7BV%7D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bn%20%5Ctimes%20v%7D%0A#card=math&code=%5Coperatorname%7Bsoftmax%7D%5Cleft%28%5Cfrac%7B%5Cmathbf%7BQ%7D%20%5Cmathbf%7BK%7D%5E%7B%5Ctop%7D%7D%7B%5Csqrt%7Bd%7D%7D%5Cright%29%20%5Cmathbf%7BV%7D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bn%20%5Ctimes%20v%7D%0A&id=NGrOb)

#@save
class DotProductAttention(nn.Module):
    """缩放点积注意力"""
    def __init__(self, dropout, **kwargs):
        super(DotProductAttention, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)    # 使用暂退法正则化

    # 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_b=True为了交换keys的最后两个维度
        scores = torch.bmm(queries, keys.transpose(1,2)) / math.sqrt(d)
        self.attention_weights = masked_softmax(scores, valid_lens)
        return torch.bmm(self.dropout(self.attention_weights), values)

使用与先前加性注意力例子中相同的键、值和有效长度。 对于点积操作,我们令查询的特征维度与键的特征维度大小相同。

queries = torch.normal(0, 1, (2, 1, 2))
attention = DotProductAttention(dropout=0.5)
attention.eval()
attention(queries, keys, values, valid_lens)
"""
tensor([[[ 2.0000,  3.0000,  4.0000,  5.0000]],
        [[10.0000, 11.0000, 12.0000, 13.0000]]])
"""

与加性注意力演示相同,由于键包含的是相同的元素, 而这些元素无法通过任何查询进行区分,因此获得了均匀的注意力权重。

d2l.show_heatmaps(attention.attention_weights.reshape((1, 1, 2, 10)),
                  xlabel='Keys', ylabel='Queries')

output_attention-scoring-functions_2a8fdc_114_0.svg

10.4. Bahdanau 注意力

QQ图片20220510145220-16521656239311.png
先前设计的seq2seq模型在每个解码步骤中对上下文变量的使用都是无差别的,其仅输出一个上下文变量,不能对此直接建模。Bahdanau等人提出了一种可微注意力模型,在预测词元时,如果不是所有输入词元都相关,模型将仅对齐输入序列中与当前预测相关的部分。

10.4.1. 模型

下面描述的Bahdanau 注意力模型将遵循9.7节的符号表达。新的注意力模型与9.7中的模型几乎相同,只不过上下文变量10. 注意力机制 - 图51在任何解码时间步10. 注意力机制 - 图52都会被10. 注意力机制 - 图53替换(T为输入序列中词元的个数):
10. 注意力机制 - 图54%5Cmathbf%7Bh%7D%7Bt%7D%0A#card=math&code=%5Cmathbf%7Bc%7D%7Bt%27%7D%3D%5Csum%7Bt%3D1%7D%5E%7BT%7D%CE%B1%28%5Cmathbf%7Bs%7D%7Bt%27-1%7D%2C%5Cmathbf%7Bh%7D%7Bt%7D%29%5Cmathbf%7Bh%7D%7Bt%7D%0A&id=cf19W)
时间步t′−1时的解码器隐状态10. 注意力机制 - 图55是查询, 编码器隐状态10. 注意力机制 - 图56既是键,也是值, 注意力权重α是使用 (10.3)节中所定义的加性注意力打分函数计算的。下面是带有Bahdanau 注意力的 编码器-解码器 架构 。
seq2seq-attention-details.svg

10.4.2. 定义注意力解码器

上面结构相比之前仅需更改解码器就够了。

#@save
class AttentionDecoder(d2l.Decoder):
    """带有注意力机制解码器的基本接口"""
    def __init__(self, **kwargs):
        super(AttentionDecoder, self).__init__(**kwargs)

    @property
    def attention_weights(self):
        raise NotImplementedError

接下来实现带有Bahdanau注意力的循环神经网络解码器。需要下面的输入:

  1. 编码器在所有时间步的最终层隐状态,将作为注意力的键和值;
  2. 上一时间步的编码器全层隐状态,将作为初始化解码器的隐状态;
  3. 编码器有效长度(排除在注意力池中填充词元)。

在每个解码时间步骤中,解码器上一个时间步的最终层隐状态将用作查询。 因此,注意力输出和输入嵌入都连结为循环神经网络解码器的输入。

class Seq2SeqAttentionDecoder(AttentionDecoder):
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,dropout=0, **kwargs):
        super(Seq2SeqAttentionDecoder, self).__init__(**kwargs)
        self.attention = d2l.AdditiveAttention(
            num_hiddens, num_hiddens, num_hiddens, dropout)
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.rnn = nn.GRU(
            embed_size + num_hiddens, num_hiddens, num_layers,
            dropout=dropout)
        self.dense = nn.Linear(num_hiddens, vocab_size)

    def init_state(self, enc_outputs, enc_valid_lens, *args):
        # outputs的形状为(batch_size,num_steps,num_hiddens).
        # hidden_state的形状为(num_layers,batch_size,num_hiddens)
        outputs, hidden_state = enc_outputs
        return (outputs.permute(1, 0, 2), hidden_state, enc_valid_lens)

    def forward(self, X, state):
        # enc_outputs的形状为(batch_size,num_steps,num_hiddens).
        # hidden_state的形状为(num_layers,batch_size,
        # num_hiddens)
        enc_outputs, hidden_state, enc_valid_lens = state
        # 输出X的形状为(num_steps,batch_size,embed_size)
        X = self.embedding(X).permute(1, 0, 2)
        outputs, self._attention_weights = [], []
        for x in X:
            # query的形状为(batch_size,1,num_hiddens)
            query = torch.unsqueeze(hidden_state[-1], dim=1)
            # context的形状为(batch_size,1,num_hiddens)
            context = self.attention(
                query, enc_outputs, enc_outputs, enc_valid_lens)
            # 在特征维度上连结
            x = torch.cat((context, torch.unsqueeze(x, dim=1)), dim=-1)
            # 将x变形为(1,batch_size,embed_size+num_hiddens)
            out, hidden_state = self.rnn(x.permute(1, 0, 2), hidden_state)
            outputs.append(out)
            self._attention_weights.append(self.attention.attention_weights)
        # 全连接层变换后,outputs的形状为
        # (num_steps,batch_size,vocab_size)
        outputs = self.dense(torch.cat(outputs, dim=0))
        return outputs.permute(1, 0, 2), [enc_outputs, hidden_state,
                                          enc_valid_lens]

    @property
    def attention_weights(self):
        return self._attention_weights

# 使用包含7个时间步的4个序列输入的小批量测试Bahdanau注意力解码器。
encoder = d2l.Seq2SeqEncoder(vocab_size=10, embed_size=8, num_hiddens=16,num_layers=2)
encoder.eval()
decoder = Seq2SeqAttentionDecoder(vocab_size=10, embed_size=8, num_hiddens=16,num_layers=2)
decoder.eval()
X = torch.zeros((4, 7), dtype=torch.long)  # (batch_size,num_steps)
state = decoder.init_state(encoder(X), None)
output, state = decoder(X, state)
output.shape, len(state), state[0].shape, len(state[1]), state[1][0].shape
# (torch.Size([4, 7, 10]), 3, torch.Size([4, 7, 16]), 2, torch.Size([4, 16]))

10.4.3. 训练

# 对模型进行机器翻译训练
embed_size, num_hiddens, num_layers, dropout = 32, 32, 2, 0.1
batch_size, num_steps = 64, 10
lr, num_epochs, device = 0.005, 250, d2l.try_gpu()

train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)
encoder = d2l.Seq2SeqEncoder(
    len(src_vocab), embed_size, num_hiddens, num_layers, dropout)
decoder = Seq2SeqAttentionDecoder(
    len(tgt_vocab), embed_size, num_hiddens, num_layers, dropout)
net = d2l.EncoderDecoder(encoder, decoder)
d2l.train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)
# loss 0.020, 4754.2 tokens/sec on cuda:0

output_bahdanau-attention_7f08d9_44_1.svg
训练结束后通过可视化注意力权重发现,每个查询都会在键值对上分配不同的权重,这说明在每个解码步中,输入序列的不同部分被选择性地聚集在注意力池中。

attention_weights = torch.cat([step[0][0][0] for step in dec_attention_weight_seq], 0).reshape((
    1, 1, -1, num_steps))
# 加上一个包含序列结束词元
d2l.show_heatmaps(
    attention_weights[:, :, :, :len(engs[-1].split()) + 1].cpu(),
    xlabel='Key positions', ylabel='Query positions')

output_bahdanau-attention_7f08d9_71_0.svg

10.5. 多头注意力

在实践中,当给定相同的查询、键和值的集合时, 我们希望模型可以学习到不同的行为, 然后将不同的行为组合起来, 捕获序列内各种范围的依赖关系 (例如,短距离依赖和长距离依赖关系)。
可以用独立学习得到的h组不同的线性投影来变换查询、键和值。 然后,这h组变换后的查询、键和值将并行地送到注意力汇聚中。 最后,将这h个注意力汇聚的输出拼接在一起, 并且通过另一个可以学习的线性投影进行变换。 这种设计被称为多头注意力(multihead attention) [Vaswani et al., 2017]。 下图展示了使用全连接层来实现可学习的线性变换的多头注意力。
multi-head-attention.svg
给定查询 10. 注意力机制 - 图61 、键 10. 注意力机制 - 图62 和 值 10. 注意力机制 - 图63 , 每个注意力头 10. 注意力机制 - 图64#card=math&code=%5Cmathbf%7Bh%7D%7Bi%7D%28i%3D1%2C%20%5Cldots%2C%20h%29&id=kstmG) 的计算方法为:
![](https://g.yuque.com/gr/latex?%5Cmathbf%7Bh%7D
%7Bi%7D%3Df%5Cleft(%5Cmathbf%7BW%7D%7Bi%7D%5E%7B(q)%7D%20%5Cmathbf%7Bq%7D%2C%20%5Cmathbf%7BW%7D%7Bi%7D%5E%7B(k)%7D%20%5Cmathbf%7Bk%7D%2C%20%5Cmathbf%7BW%7D%7Bi%7D%5E%7B(v)%7D%20%5Cmathbf%7Bv%7D%5Cright)%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bp%7Bv%7D%7D%0A#card=math&code=%5Cmathbf%7Bh%7D%7Bi%7D%3Df%5Cleft%28%5Cmathbf%7BW%7D%7Bi%7D%5E%7B%28q%29%7D%20%5Cmathbf%7Bq%7D%2C%20%5Cmathbf%7BW%7D%7Bi%7D%5E%7B%28k%29%7D%20%5Cmathbf%7Bk%7D%2C%20%5Cmathbf%7BW%7D%7Bi%7D%5E%7B%28v%29%7D%20%5Cmathbf%7Bv%7D%5Cright%29%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bp%7Bv%7D%7D%0A&id=SiMzq)
可学习的参数包括 ![](https://g.yuque.com/gr/latex?%5Cmathbf%7BW%7D
%7Bi%7D%5E%7B(q)%7D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bp%7Bq%7D%20%5Ctimes%20d%7Bq%7D%7D%20%E3%80%81%20%5Cmathbf%7BW%7D%7Bi%7D%5E%7B(k)%7D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bp%7Bk%7D%20%5Ctimes%20d%7Bk%7D%7D%20%E3%80%81%20%5Cmathbf%7BW%7D%7Bi%7D%5E%7B(v)%7D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bp%7Bv%7D%20%5Ctimes%20d%7Bv%7D%7D#card=math&code=%5Cmathbf%7BW%7D%7Bi%7D%5E%7B%28q%29%7D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bp%7Bq%7D%20%5Ctimes%20d%7Bq%7D%7D%20%E3%80%81%20%5Cmathbf%7BW%7D%7Bi%7D%5E%7B%28k%29%7D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bp%7Bk%7D%20%5Ctimes%20d%7Bk%7D%7D%20%E3%80%81%20%5Cmathbf%7BW%7D%7Bi%7D%5E%7B%28v%29%7D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bp%7Bv%7D%20%5Ctimes%20d%7Bv%7D%7D&id=eliam) , 以及注意力汇聚函数 10. 注意力机制 - 图65 (加性注意力或缩放点积注意力)。多头注意力的输出需要经过另一个线性转换, 它对应着 h 个头连结后的结果, 因此其可学习参数是 ![](https://g.yuque.com/gr/latex?%5Cmathbf%7BW%7D%7Bo%7D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bp%7Bo%7D%20%5Ctimes%20h%20p%7Bv%7D%7D#card=math&code=%5Cmathbf%7BW%7D%7Bo%7D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bp%7Bo%7D%20%5Ctimes%20h%20p%7Bv%7D%7D&id=plnbi) :
![](https://g.yuque.com/gr/latex?%5Cmathbf%7BW%7D
%7Bo%7D%5Cleft%5B%5Cbegin%7Barray%7D%7Bc%7D%0A%5Cmathbf%7Bh%7D%7B1%7D%20%5C%5C%0A%5Cvdots%20%5C%5C%0A%5Cmathbf%7Bh%7D%7Bh%7D%0A%5Cend%7Barray%7D%5Cright%5D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bp%7Bo%7D%7D%0A#card=math&code=%5Cmathbf%7BW%7D%7Bo%7D%5Cleft%5B%5Cbegin%7Barray%7D%7Bc%7D%0A%5Cmathbf%7Bh%7D%7B1%7D%20%5C%5C%0A%5Cvdots%20%5C%5C%0A%5Cmathbf%7Bh%7D%7Bh%7D%0A%5Cend%7Barray%7D%5Cright%5D%20%5Cin%20%5Cmathbb%7BR%7D%5E%7Bp%7Bo%7D%7D%0A&id=caGYG)
基于这种设计, 每个头都可能会关注输入的不同部分,可以表示比简单加权平均值更复杂的函数
我们选择缩放点积注意力作为每一个注意力头。 值得注意的是,如果我们将查询、键和值的线性变换的输出数量设置为 ![](https://g.yuque.com/gr/latex?p
%7Bq%7Dh%3Dp%7Bk%7Dh%3Dp%7Bv%7Dh%3Dp%7Bo%7D#card=math&code=p%7Bq%7Dh%3Dp%7Bk%7Dh%3Dp%7Bv%7Dh%3Dp%7Bo%7D&id=krIBN), 则可以并行计算h个头。 在下面的实现中,![](https://g.yuque.com/gr/latex?p%7Bo%7D#card=math&code=p_%7Bo%7D&id=i36L3)是通过参数num_hiddens指定的。

#@save
class MultiHeadAttention(nn.Module):
    """多头注意力"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 num_heads, dropout, bias=False, **kwargs):
        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)
        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)

下面我们使用键和值相同的小例子来测试我们编写的MultiHeadAttention类。 多头注意力输出的形状是(batch_sizenum_queriesnum_hiddens)。

num_hiddens, num_heads = 100, 5
attention = MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,num_hiddens, num_heads, 0.5)
attention.eval()
"""
MultiHeadAttention(
  (attention): DotProductAttention(
    (dropout): Dropout(p=0.5, inplace=False)
  )
  (W_q): Linear(in_features=100, out_features=100, bias=False)
  (W_k): Linear(in_features=100, out_features=100, bias=False)
  (W_v): Linear(in_features=100, out_features=100, bias=False)
  (W_o): Linear(in_features=100, out_features=100, bias=False)
)
"""
batch_size, num_queries = 2, 4
num_kvpairs, valid_lens =  6, torch.tensor([3, 2])
X = torch.ones((batch_size, num_queries, num_hiddens))
Y = torch.ones((batch_size, num_kvpairs, num_hiddens))
attention(X, Y, Y, valid_lens).shape
# torch.Size([2, 4, 100])

10.6. 自注意力和位置编码

在深度学习中,经常使用卷积神经网络(CNN)或循环神经网络(RNN)对序列进行编码。 有了注意力机制之后,我们将词元序列输入注意力池化中,每个查询都会关注所有的键-值对并生成一个注意力输出。 由于查询、键和值来自同一组输入,因此被称为 自注意力(self-attention) [Lin et al., 2017b][Vaswani et al., 2017]。 在本节中,我们将使用自注意力进行序列编码。

10.6.1. 自注意力

给定一个由词元组成的输入序列10. 注意力机制 - 图66,其中任意10. 注意力机制 - 图67(1≤i≤n)。该序列的自注意力输出为一个长度相同的序列10. 注意力机制 - 图68,其中:
10. 注意力机制 - 图69%2C%E2%80%A6%2C(%5Cmathbf%7Bx%7D%7Bn%7D%2C%5Cmathbf%7Bx%7D%7Bn%7D))%E2%88%88%5Cmathbb%7BR%7D%5E%7Bd%7D%0A#card=math&code=%5Cmathbf%7By%7D%7Bi%7D%3Df%28%5Cmathbf%7Bx%7D%7Bi%7D%2C%28%5Cmathbf%7Bx%7D%7B1%7D%2C%5Cmathbf%7Bx%7D%7B1%7D%29%2C%E2%80%A6%2C%28%5Cmathbf%7Bx%7D%7Bn%7D%2C%5Cmathbf%7Bx%7D%7Bn%7D%29%29%E2%88%88%5Cmathbb%7BR%7D%5E%7Bd%7D%0A&id=paJuf)
根据 (10.2)节中定义的注意力池化函数10. 注意力机制 - 图70。下面的代码片段是基于多头注意力对一个张量完成自注意力的计算,张量的形状为(批量大小,时间步的数目或词元序列的长度,d)。输出与输入的张量形状相同。

num_hiddens, num_heads = 100, 5
attention = d2l.MultiHeadAttention(num_hiddens, num_hiddens, num_hiddens,num_hiddens, num_heads, 0.5)
attention.eval()
"""
MultiHeadAttention(
  (attention): DotProductAttention(
    (dropout): Dropout(p=0.5, inplace=False)
  )
  (W_q): Linear(in_features=100, out_features=100, bias=False)
  (W_k): Linear(in_features=100, out_features=100, bias=False)
  (W_v): Linear(in_features=100, out_features=100, bias=False)
  (W_o): Linear(in_features=100, out_features=100, bias=False)
)
"""
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    # torch.Size([2, 4, 100])

QQ图片20220510145842.jpg
RNN必须顺序计算,无法并行处理,且由于序列的增长会逐渐忘记前面的信息。但其强时序的特性,相比于其他模型对时序信息的记忆是更好的。
自注意力可以与每个“键值对”计算,因此其可以处理非常长的序列文本,所以使用它的Transformer、BERT、GPT等输入序列通常都非常大。但此时其计算复杂度可能变得非常大,这些模型都需要非常大的算力进行计算,BERT、GPT等模型甚至需要成百上千的GPU并行计算。

10.6.2. 位置编码

在处理词元序列时,循环神经网络是逐个的重复地处理词元的, 而自注意力则因为并行计算而放弃了顺序操作。 为了使用序列的顺序信息,在输入表示中添加 位置编码(positional encoding)来注入绝对的或相对的位置信息。 位置编码可以通过学习得到也可以直接固定得到。 接下来,我们描述的是基于正弦函数和余弦函数的固定位置编码 [Vaswani et al., 2017]
假设输入 10. 注意力机制 - 图72 包含一个序列中 n 个词元的 d 维嵌入表示。位置编码使用相同形状的位置嵌入矩阵 10. 注意力机制 - 图73 输出 10. 注意力机制 - 图74 , 矩阵第 i 行、第 2 j 列和 2 j+1 列上 的元素为:
10. 注意力机制 - 图75%2C%20%5C%5C%0Ap%7Bi%2C%202%20j%2B1%7D%20%26%3D%5Ccos%20%5Cleft(%5Cfrac%7Bi%7D%7B10000%5E%7B2%20j%20%2F%20d%7D%7D%5Cright)%20.%0A%5Cend%7Baligned%7D%0A#card=math&code=%5Cbegin%7Baligned%7D%0Ap%7Bi%2C%202%20j%7D%20%26%3D%5Csin%20%5Cleft%28%5Cfrac%7Bi%7D%7B10000%5E%7B2%20j%20%2F%20d%7D%7D%5Cright%29%2C%20%5C%5C%0Ap_%7Bi%2C%202%20j%2B1%7D%20%26%3D%5Ccos%20%5Cleft%28%5Cfrac%7Bi%7D%7B10000%5E%7B2%20j%20%2F%20d%7D%7D%5Cright%29%20.%0A%5Cend%7Baligned%7D%0A&id=ugIEB)
在解释这个设计之前, 让我们先在 下面的PositionalEncoding类中实现它。

#@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)

在位置嵌入矩阵10. 注意力机制 - 图76中, 行代表词元在序列中的位置,列代表位置编码的不同维度。 在下面的例子中,我们可以看到每个位置的所有值相加,其所得结果都是不同的。

encoding_dim, num_steps = 32, 60
pos_encoding = PositionalEncoding(encoding_dim, 0)
pos_encoding.eval()
X = pos_encoding(torch.zeros((1, num_steps, encoding_dim)))
P = pos_encoding.P[:, :X.shape[1], :]
d2l.plot(torch.arange(num_steps), P[0, :, 6:10].T, xlabel='Row (position)',figsize=(6, 2.5), legend=["Col %d" % d for d in torch.arange(6, 10)])

output_self-attention-and-positional-encoding_d76d5a_43_0.svg

# 类似于计算机的二进制编码
P = P[0, :, :].unsqueeze(0).unsqueeze(0)
d2l.show_heatmaps(P, xlabel='Column (encoding dimension)',
                  ylabel='Row (position)', figsize=(3.5, 4), cmap='Blues')

output_self-attention-and-positional-encoding_d76d5a_67_0.svg
使用上述位置编码的优势是除了捕获绝对位置信息之外,上述的位置编码还允许模型学习得到输入序列中相对位置信息。 这是因为对于任何确定的位置偏移δ,位置i+δ处 的位置编码可以线性投影位置i处的位置编码来表示。这种投影的数学解释是, 令 10. 注意力机制 - 图79 , 对于任何确定的位置偏移 10. 注意力机制 - 图80 , (10.6.2).中的任何一对 10. 注意力机制 - 图81#card=math&code=%5Cleft%28p%7Bi%2C%202%20j%7D%2C%20p%7Bi%2C%202%20j%2B1%7D%5Cright%29&id=WZ2vB) 都可以线性投影到 10. 注意力机制 - 图82#card=math&code=%5Cleft%28p%7Bi%2B%5Cdelta%2C%202%20j%7D%2C%20p%7Bi%2B%5Cdelta%2C%202%20j%2B1%7D%5Cright%29&id=C1TYX) :

10. 注意力机制 - 图83%20%26%20%5Csin%20%5Cleft(%5Cdelta%20%5Comega%7Bj%7D%5Cright)%20%5C%5C%0A-%5Csin%20%5Cleft(%5Cdelta%20%5Comega%7Bj%7D%5Cright)%20%26%20%5Ccos%20%5Cleft(%5Cdelta%20%5Comega%7Bj%7D%5Cright)%0A%5Cend%7Barray%7D%5Cright%5D%5Cleft%5B%5Cbegin%7Barray%7D%7Bc%7D%0Ap%7Bi%2C%202%20j%7D%20%5C%5C%0Ap%7Bi%2C%202%20j%2B1%7D%0A%5Cend%7Barray%7D%5Cright%5D%20%7D%20%5C%5C%0A%3D%26%20%7B%5Cleft%5B%5Cbegin%7Barray%7D%7Bc%7D%0A%5Ccos%20%5Cleft(%5Cdelta%20%5Comega%7Bj%7D%5Cright)%20%5Csin%20%5Cleft(i%20%5Comega%7Bj%7D%5Cright)%2B%5Csin%20%5Cleft(%5Cdelta%20%5Comega%7Bj%7D%5Cright)%20%5Ccos%20%5Cleft(i%20%5Comega%7Bj%7D%5Cright)%20%5C%5C%0A-%5Csin%20%5Cleft(%5Cdelta%20%5Comega%7Bj%7D%5Cright)%20%5Csin%20%5Cleft(i%20%5Comega%7Bj%7D%5Cright)%2B%5Ccos%20%5Cleft(%5Cdelta%20%5Comega%7Bj%7D%5Cright)%20%5Ccos%20%5Cleft(i%20%5Comega%7Bj%7D%5Cright)%0A%5Cend%7Barray%7D%5Cright%5D%20%7D%20%5C%5C%0A%3D%26%20%7B%5Cleft%5B%5Cbegin%7Barray%7D%7Bc%7D%0A%5Csin%20%5Cleft((i%2B%5Cdelta)%20%5Comega%7Bj%7D%5Cright)%20%5C%5C%0A%5Ccos%20%5Cleft((i%2B%5Cdelta)%20%5Comega%7Bj%7D%5Cright)%0A%5Cend%7Barray%7D%5Cright%5D%20%7D%20%5C%5C%0A%3D%26%20%7B%5Cleft%5B%5Cbegin%7Barray%7D%7Bc%7D%0Ap%7Bi%2B%5Cdelta%2C%202%20j%7D%20%5C%5C%0Ap%7Bi%2B%5Cdelta%2C%202%20j%2B1%7D%0A%5Cend%7Barray%7D%5Cright%5D%20%7D%0A%5Cend%7Baligned%7D%0A#card=math&code=%5Cbegin%7Baligned%7D%0A%26%20%7B%5Cleft%5B%5Cbegin%7Barray%7D%7Bcc%7D%0A%5Ccos%20%5Cleft%28%5Cdelta%20%5Comega%7Bj%7D%5Cright%29%20%26%20%5Csin%20%5Cleft%28%5Cdelta%20%5Comega%7Bj%7D%5Cright%29%20%5C%5C%0A-%5Csin%20%5Cleft%28%5Cdelta%20%5Comega%7Bj%7D%5Cright%29%20%26%20%5Ccos%20%5Cleft%28%5Cdelta%20%5Comega%7Bj%7D%5Cright%29%0A%5Cend%7Barray%7D%5Cright%5D%5Cleft%5B%5Cbegin%7Barray%7D%7Bc%7D%0Ap%7Bi%2C%202%20j%7D%20%5C%5C%0Ap%7Bi%2C%202%20j%2B1%7D%0A%5Cend%7Barray%7D%5Cright%5D%20%7D%20%5C%5C%0A%3D%26%20%7B%5Cleft%5B%5Cbegin%7Barray%7D%7Bc%7D%0A%5Ccos%20%5Cleft%28%5Cdelta%20%5Comega%7Bj%7D%5Cright%29%20%5Csin%20%5Cleft%28i%20%5Comega%7Bj%7D%5Cright%29%2B%5Csin%20%5Cleft%28%5Cdelta%20%5Comega%7Bj%7D%5Cright%29%20%5Ccos%20%5Cleft%28i%20%5Comega%7Bj%7D%5Cright%29%20%5C%5C%0A-%5Csin%20%5Cleft%28%5Cdelta%20%5Comega%7Bj%7D%5Cright%29%20%5Csin%20%5Cleft%28i%20%5Comega%7Bj%7D%5Cright%29%2B%5Ccos%20%5Cleft%28%5Cdelta%20%5Comega%7Bj%7D%5Cright%29%20%5Ccos%20%5Cleft%28i%20%5Comega%7Bj%7D%5Cright%29%0A%5Cend%7Barray%7D%5Cright%5D%20%7D%20%5C%5C%0A%3D%26%20%7B%5Cleft%5B%5Cbegin%7Barray%7D%7Bc%7D%0A%5Csin%20%5Cleft%28%28i%2B%5Cdelta%29%20%5Comega%7Bj%7D%5Cright%29%20%5C%5C%0A%5Ccos%20%5Cleft%28%28i%2B%5Cdelta%29%20%5Comega%7Bj%7D%5Cright%29%0A%5Cend%7Barray%7D%5Cright%5D%20%7D%20%5C%5C%0A%3D%26%20%7B%5Cleft%5B%5Cbegin%7Barray%7D%7Bc%7D%0Ap%7Bi%2B%5Cdelta%2C%202%20j%7D%20%5C%5C%0Ap_%7Bi%2B%5Cdelta%2C%202%20j%2B1%7D%0A%5Cend%7Barray%7D%5Cright%5D%20%7D%0A%5Cend%7Baligned%7D%0A&id=mHNVb)
10. 注意力机制 - 图84 投影矩阵不依赖于任何位置的索引10. 注意力机制 - 图85

10.7 Transformer

Transformer模型完全基于注意力机制,没有任何卷积层或循环神经网络层 [Vaswani et al., 2017]

10.7.1. 模型

与10.4节中基于Bahdanau注意力实现的序列到序列的学习相比,transformer的编码器和解码器是基于自注意力的模块叠加而成的,源(输入)序列和目标(输出)序列的嵌入(embedding)表示将加上位置编码(positional encoding),再分别输入到编码器和解码器中。
transformer.svg
编码器是由多个相同的层叠加而成的,每个层都有两个子层:多头自注意力汇聚基于位置的前馈网络。在计算编码器的自注意力时,查询、键和值都来自前一个编码器层的输出。在transformer中,对于序列中任何位置的任何输入10. 注意力机制 - 图87,都要求满足10. 注意力机制 - 图88%E2%88%88%5Cmathbb%7BR%7D%5E%7Bd%7D#card=math&code=%5Ctext%7Bsublayer%7D%28%5Cmathbf%7Bx%7D%29%E2%88%88%5Cmathbb%7BR%7D%5E%7Bd%7D&id=qxhfO),以便残差连接满足10. 注意力机制 - 图89%E2%88%88%5Cmathbb%7BR%7D%5E%7Bd%7D#card=math&code=%5Cmathbf%7Bx%7D%2B%5Ctext%7Bsublayer%7D%28%5Cmathbf%7Bx%7D%29%E2%88%88%5Cmathbb%7BR%7D%5E%7Bd%7D&id=c3a9u)。在残差连接的加法计算之后,紧接着应用层规范化(layer normalization) [Ba et al., 2016]
在解码器多头注意力中,查询来自前一个解码器层的输出,而键和值来自整个编码器的输出。在解码器掩蔽多头注意力中,查询、键和值都来自上一个解码器层的输出。Attention在中间的任意一步都可看到后面的所有信息,这在编码器中是可以的;但在解码器中不应当考虑其本身或之后的元素,这可以通过掩码来实现,即计算10. 注意力机制 - 图90输出时,假设当前序列为10. 注意力机制 - 图91
我们已经描述并实现了基于缩放点积多头注意力 [10.5节]和位置编码 [10.6.3节]。接下来,我们将实现transformer模型的剩余部分。

10.7.2. 基于位置的前馈网络

基于位置的前馈网络由两个全连接层构成。第一层将输入状态由(b, n, d)转化为(bn, d),第二层再将其变回(b, n, d),其中,n为可变的序列长度。理论上来说,模型应当可以处理任意长度的序列,因此从矩阵乘法角度上,10. 注意力机制 - 图92是不可取的,(bn, d)可表达为将每个序列中的元素取出,做全连接。在下面的实现中,输入X的形状(批量大小,时间步数或序列长度,隐单元数或特征维度)将被一个两层的感知机转换成形状为(批量大小,时间步数,ffn_num_outputs)的输出张量。

#@save
class PositionWiseFFN(nn.Module):
    """基于位置的前馈网络"""
    def __init__(self, ffn_num_input, ffn_num_hiddens, ffn_num_outputs,**kwargs):
        super(PositionWiseFFN, self).__init__(**kwargs)
        self.dense1 = nn.Linear(ffn_num_input, ffn_num_hiddens)
        self.relu = nn.ReLU()
        self.dense2 = nn.Linear(ffn_num_hiddens, ffn_num_outputs)

    def forward(self, X):
        return self.dense2(self.relu(self.dense1(X)))

10.7.3. 残差连接和层规范化

在 [7.5节]中,我们解释了在一个小批量的样本内基于批量规范化对数据进行重新中心化和重新缩放的调整。层规范化和批量规范化的目标相同,但层规范化是基于特征维度进行规范化。尽管批量规范化在计算机视觉中被广泛应用,但在自然语言处理任务中(输入通常是变长序列)批量规范化通常不如层规范化的效果好。

  • 批量归一化是对每个特征/通道里元素做归一化,不适合序列长度会变的NLP应用。
  • 层归一化是对每个样本里的元素归一化。

D851B4FB19D31498F5BD535E8FC0B91C.png

#@save
class AddNorm(nn.Module):
    """残差连接后进行层规范化"""
    def __init__(self, normalized_shape, dropout, **kwargs):
        super(AddNorm, self).__init__(**kwargs)
        self.dropout = nn.Dropout(dropout)
        self.ln = nn.LayerNorm(normalized_shape)

    def forward(self, X, Y):
        return self.ln(self.dropout(Y) + X)

# 残差连接要求两个输入的形状相同,以便加法操作后输出张量的形状相同。
add_norm = AddNorm([3, 4], 0.5)
add_norm.eval()
add_norm(torch.ones((2, 3, 4)), torch.ones((2, 3, 4))).shape    # torch.Size([2, 3, 4])

10.7.4. 编码器

下面实现了一个Transformer block,下面的EncoderBlock类包含两个子层:多头自注意力和基于位置的前馈网络,这两个子层都使用了残差连接和紧随的层规范化。编码器中的任何层都不会改变输入的形状。

#@save
class EncoderBlock(nn.Module):
    """transformer编码器块"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,dropout, use_bias=False, **kwargs):
        super(EncoderBlock, self).__init__(**kwargs)
        self.attention = d2l.MultiHeadAttention(key_size, query_size, value_size, num_hiddens, num_heads, dropout,use_bias)
        self.addnorm1 = AddNorm(norm_shape, dropout)
        self.ffn = PositionWiseFFN(
            ffn_num_input, ffn_num_hiddens, num_hiddens)
        self.addnorm2 = AddNorm(norm_shape, dropout)

    def forward(self, X, valid_lens):
        Y = self.addnorm1(X, self.attention(X, X, X, valid_lens))
        return self.addnorm2(Y, self.ffn(Y))

下面的transformer编码器堆叠了num_layersEncoderBlock类的实例。由于我们使用的是值范围在−1和1之间的固定位置编码,因此通过学习得到的输入的嵌入表示的值需要先乘以嵌入维度的平方根进行重新缩放,然后再与位置编码相加。

#@save
class TransformerEncoder(d2l.Encoder):
    """transformer编码器"""
    def __init__(self, vocab_size, key_size, query_size, value_size,num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,num_heads, num_layers, dropout, use_bias=False, **kwargs):
        super(TransformerEncoder, self).__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
        self.blks = nn.Sequential()
        for i in range(num_layers):
            self.blks.add_module("block"+str(i),
                EncoderBlock(key_size, query_size, value_size, num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,num_heads, dropout, use_bias))

    def forward(self, X, valid_lens, *args):
        # 位置编码在(-1,1)之间,而Embedding的输出不在此区间内,因此对其缩放保证其在(-1,1)
        # 因此嵌入值乘以嵌入维度的平方根进行缩放,
        # 然后再与位置编码相加。
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        self.attention_weights = [None] * len(self.blks)
        for i, blk in enumerate(self.blks):
            X = blk(X, valid_lens)
            self.attention_weights[i] = blk.attention.attention.attention_weights
        return X

10.7.5. 解码器

在训练阶段,输出序列的所有位置(时间步)的词元都是已知的;然而,在预测阶段,其输出序列的词元是逐个生成的。因此,在任何解码器时间步中,只有生成的词元才能用于解码器的自注意力计算中,因此掩蔽自注意力设定了参数dec_valid_lens,以便任何查询都只会与解码器中所有已经生成词元的位置(即直到该查询位置为止)进行注意力计算。

class DecoderBlock(nn.Module):
    """解码器中第i个块"""
    def __init__(self, key_size, query_size, value_size, num_hiddens,
                 norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
                 dropout, i, **kwargs):
        super(DecoderBlock, self).__init__(**kwargs)
        self.i = i
        self.attention1 = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        self.addnorm1 = AddNorm(norm_shape, dropout)
        self.attention2 = d2l.MultiHeadAttention(
            key_size, query_size, value_size, num_hiddens, num_heads, dropout)
        self.addnorm2 = AddNorm(norm_shape, dropout)
        self.ffn = PositionWiseFFN(ffn_num_input, ffn_num_hiddens,
                                   num_hiddens)
        self.addnorm3 = AddNorm(norm_shape, dropout)

    def forward(self, X, state):# state包含3个元素
        enc_outputs, enc_valid_lens = state[0], state[1]
        # 训练阶段,输出序列的所有词元都在同一时间处理,
        # 因此state[2][self.i]初始化为None。
        # 预测阶段,输出序列是通过词元一个接着一个解码的,
        # 因此state[2][self.i]包含着直到当前时间步第i个块解码的输出表示
        if state[2][self.i] is None:
            key_values = X
        else:
            key_values = torch.cat((state[2][self.i], X), axis=1)# 上一个解码器的隐状态与当前的输入
        state[2][self.i] = key_values
        if self.training:# 训练时,由于可以看到后面的数据,因此需要掩蔽
            batch_size, num_steps, _ = X.shape
            # dec_valid_lens:(batch_size,num_steps),
            # 其中每一行是[1,2,...,num_steps]
            dec_valid_lens = torch.arange(
                1, num_steps + 1, device=X.device).repeat(batch_size, 1)
        else:# 预测时,由于本来就是按顺序进行,因此无需掩蔽
            dec_valid_lens = None

        # 自注意力
        X2 = self.attention1(X, key_values, key_values, dec_valid_lens)
        Y = self.addnorm1(X, X2)
        # 编码器-解码器注意力。
        # enc_outputs的开头:(batch_size,num_steps,num_hiddens)
        Y2 = self.attention2(Y, enc_outputs, enc_outputs, enc_valid_lens)
        Z = self.addnorm2(Y, Y2)
        return self.addnorm3(Z, self.ffn(Z)), state

为了便于在“编码器-解码器”注意力中进行缩放点积计算和残差连接中进行加法计算,编码器和解码器的特征维度都是num_hiddens
现在构建由num_layersDecoderBlock实例组成的transformer解码器。最后,通过一个全连接层计算所有vocab_size个可能的输出词元的预测值。解码器的自注意力权重和编码器解码器注意力权重都被存储下来,方便日后可视化的需要。编码器的执行方式与seq2seq是相同的,每次前进一个词,并将输出作为输入

class TransformerDecoder(d2l.AttentionDecoder):
    def __init__(self, vocab_size, key_size, query_size, value_size,
                 num_hiddens, norm_shape, ffn_num_input, ffn_num_hiddens,
                 num_heads, num_layers, dropout, **kwargs):
        super(TransformerDecoder, self).__init__(**kwargs)
        self.num_hiddens = num_hiddens
        self.num_layers = num_layers
        self.embedding = nn.Embedding(vocab_size, num_hiddens)
        self.pos_encoding = d2l.PositionalEncoding(num_hiddens, dropout)
        self.blks = nn.Sequential()
        for i in range(num_layers):
            self.blks.add_module("block"+str(i),
                DecoderBlock(key_size, query_size, value_size, num_hiddens,
                             norm_shape, ffn_num_input, ffn_num_hiddens,
                             num_heads, dropout, i))
        self.dense = nn.Linear(num_hiddens, vocab_size)

    def init_state(self, enc_outputs, enc_valid_lens, *args):
        return [enc_outputs, enc_valid_lens, [None] * self.num_layers]

    def forward(self, X, state):
        X = self.pos_encoding(self.embedding(X) * math.sqrt(self.num_hiddens))
        self._attention_weights = [[None] * len(self.blks) for _ in range (2)]
        for i, blk in enumerate(self.blks):
            X, state = blk(X, state)
            # 解码器自注意力权重
            self._attention_weights[0][
                i] = blk.attention1.attention.attention_weights
            # “编码器-解码器”自注意力权重
            self._attention_weights[1][
                i] = blk.attention2.attention.attention_weights
        return self.dense(X), state

    @property
    def attention_weights(self):
        return self._attention_weights

10.7.6. 信息传递与预测

编码器中的输出10. 注意力机制 - 图94,将其作为解码中第i个Transformer块中多头注意力的key和value,它的query来自目标序列。这意味着编码器和解码器中块的个数和输出维度都是一样的。
预测第10. 注意力机制 - 图95个输出时,在自注意力中,前t个预测值作为key和value,第t个预测值还作为query。
image-20220511150828695.png

10.7.7. 训练

在这里,指定transformer的编码器和解码器都是2层,都使用4头注意力。我们在“英语-法语”机器翻译数据集上训练transformer模型。

# 对于每个block来说,主要参数是num_hiddens与num_heads
num_hiddens, num_layers, dropout, batch_size, num_steps = 32, 2, 0.1, 64, 10
lr, num_epochs, device = 0.005, 200, d2l.try_gpu()
ffn_num_input, ffn_num_hiddens, num_heads = 32, 64, 4
# q,k,v大小的取法与num_hidden相同,为2倍数
key_size, query_size, value_size = 32, 32, 32
norm_shape = [32]

train_iter, src_vocab, tgt_vocab = d2l.load_data_nmt(batch_size, num_steps)

encoder = TransformerEncoder(
    len(src_vocab), key_size, query_size, value_size, num_hiddens,
    norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
    num_layers, dropout)
decoder = TransformerDecoder(
    len(tgt_vocab), key_size, query_size, value_size, num_hiddens,
    norm_shape, ffn_num_input, ffn_num_hiddens, num_heads,
    num_layers, dropout)
net = d2l.EncoderDecoder(encoder, decoder)
d2l.train_seq2seq(net, train_iter, lr, num_epochs, tgt_vocab, device)
# loss 0.030, 5244.3 tokens/sec on cuda:0

output_transformer_5722f1_162_1.svg
训练结束后,使用transformer模型将一些英语句子翻译成法语,并且计算它们的BLEU分数。

engs = ['go .', "i lost .", 'he\'s calm .', 'i\'m home .']
fras = ['va !', 'j\'ai perdu .', 'il est calme .', 'je suis chez moi .']
for eng, fra in zip(engs, fras):
    translation, dec_attention_weight_seq = d2l.predict_seq2seq(
        net, eng, src_vocab, tgt_vocab, num_steps, device, True)
    print(f'{eng} => {translation}, ',
          f'bleu {d2l.bleu(translation, fra, k=2):.3f}')
"""
go . => va !,  bleu 1.000
i lost . => j'ai perdu .,  bleu 1.000
he's calm . => il est calme .,  bleu 1.000
i'm home . => je suis chez moi .,  bleu 1.000
"""

当进行最后一个英语到法语的句子翻译工作时,让我们可视化transformer的注意力权重。编码器自注意力权重的形状为(编码器层数,注意力头数,num_steps或查询的数目,num_steps或“键-值”对的数目)。

enc_attention_weights = torch.cat(net.encoder.attention_weights, 0).reshape((num_layers, num_heads,
    -1, num_steps))
enc_attention_weights.shape    # torch.Size([2, 4, 10, 10])

在编码器的自注意力中,查询和键都来自相同的输入序列。因为填充词元是不携带信息的,因此通过指定输入序列的有效长度可以避免查询与使用填充词元的位置计算注意力。接下来,将逐行呈现两层多头注意力的权重。每个注意力头都根据查询、键和值的不同的表示子空间来表示不同的注意力。

d2l.show_heatmaps(
    enc_attention_weights.cpu(), xlabel='Key positions',
    ylabel='Query positions', titles=['Head %d' % i for i in range(1, 5)],
    figsize=(7, 3.5))

output_transformer_5722f1_198_0.svg
为了可视化解码器的自注意力权重和“编码器-解码器”的注意力权重,我们需要完成更多的数据操作工作。例如,我们用零填充被掩蔽住的注意力权重。值得注意的是,解码器的自注意力权重和“编码器-解码器”的注意力权重都有相同的查询:即以序列开始词元(beginning-of-sequence,BOS)打头,再与后续输出的词元共同组成序列。

dec_attention_weights_2d = [head[0].tolist()
                            for step in dec_attention_weight_seq
                            for attn in step for blk in attn for head in blk]
dec_attention_weights_filled = torch.tensor(
    pd.DataFrame(dec_attention_weights_2d).fillna(0.0).values)
dec_attention_weights = dec_attention_weights_filled.reshape((-1, 2, num_layers, num_heads, num_steps))
dec_self_attention_weights, dec_inter_attention_weights = \
    dec_attention_weights.permute(1, 2, 3, 0, 4)
dec_self_attention_weights.shape, dec_inter_attention_weights.shape
# (torch.Size([2, 4, 6, 10]), torch.Size([2, 4, 6, 10]))

由于解码器自注意力的自回归属性,查询不会对当前位置之后的“键-值”对进行注意力计算。

# Plusonetoincludethebeginning-of-sequencetoken
d2l.show_heatmaps(
    dec_self_attention_weights[:, :, :, :len(translation.split()) + 1],
    xlabel='Key positions', ylabel='Query positions',
    titles=['Head %d' % i for i in range(1, 5)], figsize=(7, 3.5))

output_transformer_5722f1_222_0.svg
与编码器的自注意力的情况类似,通过指定输入序列的有效长度,输出序列的查询不会与输入序列中填充位置的词元进行注意力计算。

d2l.show_heatmaps(
    dec_inter_attention_weights, xlabel='Key positions',
    ylabel='Query positions', titles=['Head %d' % i for i in range(1, 5)],
    figsize=(7, 3.5))

output_transformer_5722f1_234_0.svg

10.7.8. 总结

  • Transformer是一个纯使用注意力的编码-解码器
  • 编码器和解码器都有个n个transformer块
  • 每个块里使用多头(自)注意力,基于位置的前馈网络,和层归一化