《动手学深度学习》学习笔记第10章注意力机制

本系列为《动手学深度学习》学习笔记

书籍链接：动手学深度学习

笔记是从第四章开始，前面三章为基础知识，有需要的可以自己去看看

关于本系列笔记： 书里为了让读者更好的理解，有大篇幅的描述性的文字，内容很多，笔记只保留主要内容，同时也是对之前知识的查漏补缺

《动手学深度学习》学习笔记第4章多层感知机
《动手学深度学习》学习笔记第5章深度学习计算
《动手学深度学习》学习笔记第6章卷积神经网络
《动手学深度学习》学习笔记第7章现代卷积神经网络
《动手学深度学习》学习笔记第8章循环神经网络
《动手学深度学习》学习笔记第9章现代循环神经网络

10.1 注意力提示

环境中的干扰注意力的信息却并不少。比如人类的视觉神经系统大约每秒收到108位的信息，这远远超过了大脑能够完全处理的水平。幸运的是，人类的祖先已经从经验（也称为数据）中认识到“并非感官的所有输入都是一样的”。在整个人类历史中，这种只将注意力引向感兴趣的一小部分信息的能力。

10.1.1 生物学中的注意力提示

注意力是如何应用于视觉世界中的呢？这要从当今十分普及的双组件（two‐component）的框架开始讲起：这个框架的出现可以追溯到19世纪90年代的威廉·詹姆斯，他被认为是“美国心理学之父”(James, 2007)。在这个框架中，受试者基于非自主性提示和自主性提示有选择地引导注意力的焦点。

非自主性提示是基于环境中物体的突出性和易见性。想象一下，假如面前有五个物品：一份报纸、一篇研究论文、一杯咖啡、一本笔记本和一本书，就像图10.1.1。所有纸制品都是黑白印刷的，但咖啡杯是红色的。我们会把视力最敏锐的地方放到咖啡上
在这里插入图片描述
喝咖啡后，我们会变得兴奋并想读书，所以转过头，重新聚焦眼睛，然后看看书，就像图10.1.2中描述那样。与图10.1.1中由于突出性导致的选择不同，此时选择书是受到了认知和意识的控制，因此注意力在基于自主性提示去辅助选择时将更为谨慎。受试者的主观意愿推动，选择的力量也就更强大。
在这里插入图片描述

10.1.2 查询、键和值

用神经网络来设计注意力机制的框架，首先，考虑一个相对简单的状况，即只使用非自主性提示。要想将选择偏向于感官输入，则可以简单地使用参数化的全连接层，甚至是非参数化的最大汇聚层或平均汇聚层。因此，“是否包含自主性提示”将注意力机制与全连接层或汇聚层区别开来。

在注意力机制的背景下，自主性提示被称为查询（query）。给定任何查询，注意力机制通过注意力汇聚（attention pooling）将选择引导至感官输入（sensory inputs，例如中间特征表示）。
在注意力机制中，这些感官输入被称为值（value）。更通俗的解释，每个值都与一个键（key）配对，这可以想象为感官输入的非自主提示。

如图10.1.3所示，可以通过设计注意力汇聚的方式，便于给定的查询（自主性提示）与键（非自主性提示）进行匹配，这将引导得出最匹配的值（感官输入）。
在这里插入图片描述
图10.1.3: 注意力机制通过注意力汇聚将查询（自主性提示）和键（非自主性提示）结合在一起，实现对值（感官输入）的选择倾向

鉴于上面所提框架在图10.1.3中的主导地位，因此这个框架下的模型将成为本章的中心。然而，注意力机制的设计有许多替代方案。例如可以设计一个不可微的注意力模型，该模型可以使用强化学习方法 (Mnih et al., 2014)进行训练。

10.1.3 注意力的可视化

平均汇聚层可以被视为输入的加权平均值，其中各输入的权重是一样的。实际上，注意力汇聚得到的是加权平均的总和值，其中权重是在给定的查询和不同的键之间计算得出的。

import torch
from d2l import torch as d2l

为了可视化注意力权重，需要定义一个show_heatmaps函数。其输入matrices的形状是（要显示的行数，要显示的列数，查询的数目，键的数目）。

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

下面使用一个简单的例子进行演示。在本例子中，仅当查询和键相同时，注意力权重为1，否则为0。

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

在这里插入图片描述

后面的章节内容将经常调用show_heatmaps函数来显示注意力权重。

小结

人类的注意力是有限的、有价值和稀缺的资源。
受试者使用非自主性和自主性提示有选择性地引导注意力。前者基于突出性，后者则依赖于意识。
注意力机制与全连接层或者汇聚层的区别源于增加的自主提示。
由于包含了自主性提示，注意力机制与全连接的层或汇聚层不同。
注意力机制通过注意力汇聚使选择偏向于值（感官输入），其中包含查询（自主性提示）和键（非自主性提示）。键和值是成对的。可视化查询和键之间的注意力权重是可行的。

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

上节介绍了框架下的注意力机制的主要成分图10.1.3：查询（自主提示）和键（非自主提示） 之间的交互形成了注意力汇聚；注意力汇聚有选择地聚合了值（感官输入）以生成最终的输出。

1964年提出的Nadaraya‐Watson核回归模型是一个简单但完整的例子，可以用于演示具有注意力机制的机器学习。

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

10.2.1 生成数据集

给定的成对的“输入－输出”数据集 ${(x_1, y_1), . . . ,(x_n, y_n)}$ ，如何学习f来预测任意新输入x的输出 $\hat{y}= f(x)$ ？
根据下面的非线性函数生成一个人工数据集，其中加入的噪声项为 $ϵ$ ：
$y_i = 2 sin(x_i) + x^{0.8}_i + ϵ,$
其中 $ϵ$ 服从均值为0和标准差为0.5的正态分布。
（在这里生成了50个训练样本和50个测试样本。为了更好地可视化之后的注意力模式，需要将训练样本进行排序。）

n_train = 50 # 训练样本数
x_train, _ = torch.sort(torch.rand(n_train) * 5) # 排序后的训练样本


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) # 测试样本的真实输出
n_test = len(x_test) # 测试样本数
n_test

下面的函数将绘制所有的训练样本（样本由圆圈表示），不带噪声项的真实数据生成函数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);

10.2.2 平均汇聚

先使用最简单的估计器来解决回归问题。基于平均汇聚来计算所有训练样本输出值的平均值：
$\frac{1}{n} \sum^{n}_{i=1}{y_i}(10.2.2)$
如下图所示，这个估计器确实不够聪明。真实函数 $f$ （“Truth”）和预测函数（“Pred”）相差很大。

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

在这里插入图片描述

10.2.3 非参数注意力汇聚

显然，平均汇聚忽略了输入 $x_i$ 。于是Nadaraya (Nadaraya, 1964)和 Watson (Watson, 1964)提出了一个更好的想法，根据输入的位置对输出 $y_i$ 进行加权：
在这里插入图片描述
其中 $K$ 是核（kernel）。公式 (10.2.3)所描述的估计器被称为 Nadaraya-Watson核回归（Nadaraya‐Watson kernel regression）。受此启发，可以从图10.1.3中的注意力机制框架的角度重写 (10.2.3)，成为一个更加通用的注意力汇聚（attention pooling）公式：

其中 $x$ 是查询， $x_i, y_i)$ 是键值对。

比较 (10.2.4)和 (10.2.2)，注意力汇聚是 $y_i$ 的加权平均。将查询 $x$ 和键 $x_i$ 之间的关系建模为 注意力权重（attention weight） $α(x, x_i)$ ，如 (10.2.4)所示，这个权重将被分配给每一个对应值 $y_i$ 。

对于任何查询，模型在所有键值对注意力权重都是一个有效的概率分布：它们是非负的，并且总和为1。为了更好地理解注意力汇聚，下面考虑一个高斯核（Gaussian kernel），其定义为：
在这里插入图片描述
将高斯核代入 (10.2.4)和 (10.2.3)可以得到：

在 (10.2.6)中，如果一个键 $x_i$ 越是接近给定的查询 $x$ ，那么分配给这个键对应值 $y_i$ 的注意力权重就会越大，也就“获得了更多的注意力”。

值得注意的是，Nadaraya‐Watson核回归是一个非参数模型。因此，(10.2.6)是非参数的注意力汇聚（non‐parametric 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的每个元素都是值的加权平均值，其中的权重是注意力权重
y_hat = torch.matmul(attention_weights, y_train)
plot_kernel_reg(y_ha

在这里插入图片描述

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

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

在这里插入图片描述

10.2.4 带参数注意力汇聚

非参数的Nadaraya‐Watson核回归具有一致性（consistency）的优点：如果有足够的数据，此模型会收敛到最优结果。

例如，与 (10.2.6)略有不同，在下面的查询 $x$ 和键 $x_i$ 之间的距离乘以可学习参数 $w$ ：
在这里插入图片描述
批量矩阵乘法

为了更有效地计算小批量数据的注意力，可以利用深度学习开发框架中提供的批量矩阵乘法。

假设第一个小批量数据包含 $n$ 个矩阵 $X_1, . . . , X_n$ ，形状为 $a \times b$ ，
第二个小批量包含 $n$ 个矩阵 $Y 1, ..., Yn$ ，形状为 $b \times c$ 。

它们的批量矩阵乘法得到 $n$ 个矩阵 $X_1Y_1, . . . , X_nY_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
values = torch.arange(20.0).reshape((2, 10))
torch.bmm(weights.unsqueeze(1), values.unsqueeze(-1))
==============================
tensor([[[ 4.5000]],
		[[14.5000]]])

定义模型

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

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)

训练

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

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

在这里插入图片描述

如下所示，训练完带参数的注意力汇聚模型后可以发现：在尝试拟合带噪声的训练数据时，预测结果绘制的线不如之前非参数模型的平滑。

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

在这里插入图片描述
为什么新的模型更不平滑了呢？下面看一下输出结果的绘制图：与非参数的注意力汇聚模型相比，带参数的模型加入可学习的参数后，曲线在注意力权重较大的区域变得更不平滑。

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

在这里插入图片描述
小结

Nadaraya‐Watson核回归是具有注意力机制的机器学习范例。
Nadaraya‐Watson核回归的注意力汇聚是对训练数据中输出的加权平均。从注意力的角度来看，分配给每个值的注意力权重取决于将值所对应的键和查询作为输入的函数。
注意力汇聚可以分为非参数型和带参数型。

10.3 注意力评分函数

10.2节使用了高斯核来对查询和键之间的关系建模。(10.2.6)中的高斯核指数部分可以视为注意力评分函数（attention scoring function），简称评分函数（scoring function），然后把这个函数的输出结果输入到softmax函数中进行运算。通过上述步骤，将得到与键对应的值的概率分布（即注意力权重）。最后，注意力汇聚的输出就是基于这些注意力权重的值的加权和。

从宏观来看，上述算法可以用来实现图10.1.3中的注意力机制框架。图10.3.1说明了如何将注意力汇聚的输出计算成为值的加权和，其中a表示注意力评分函数。由于注意力权重是概率分布，因此加权和其本质上是加权平均值。
在这里插入图片描述
图10.3.1: 计算注意力汇聚的输出为值的加权和

用数学语言描述，假设有一个查询 $q ∈ R^q$ 和 $m$ 个“键－值”对 $k_1, v_1), . . . ,(k_m, v_m)$ ，其中 $k_i ∈ R^k$ ， $v_i ∈ R^v$ 。

注意力汇聚函数 $f$ 就被表示成值的加权和：
$f(q,(k_1, v_1), . . . ,(k_m, v_m))= \sum^{m}_{i=1}{α(q, k_i)v_i} ∈ R^v, (10.3.1)$
其中查询 $q$ 和键 $k_i$ 的注意力权重（标量）是通过注意力评分函数 $a$ 将两个向量映射成标量，再经过softmax运算得到的：
在这里插入图片描述
正如上图所示，选择不同的注意力评分函数a会导致不同的注意力汇聚操作。本节将介绍两个流行的评分函数。

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

10.3.1 掩蔽softmax操作

正如上面提到的，softmax操作用于输出一个概率分布作为注意力权重。在某些情况下，并非所有的值都应该被纳入到注意力汇聚中。

例如，为了在 9.5节中高效处理小批量数据集，某些文本序列被填充了没有意义的特殊词元。为了仅将有意义的词元作为值来获取注意力汇聚，可以指定一个有效序列长度（即词元的个数），以便在计算softmax时过滤掉超出指定范围的位置。下面的masked_softmax函数实现了这样的掩蔽softmax操作（masked softmax operation），其中任何超出有效长度的位置都被掩蔽并置为0。

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

为了演示此函数是如何工作的，考虑由两个2 × 4矩阵表示的样本，这两个样本的有效长度分别为2和3。经过掩蔽softmax操作，超出有效长度的值都被掩蔽为0。

masked_softmax(torch.rand(2, 2, 4), torch.tensor([2, 3]))
====================================================
tensor([[[0.5980, 0.4020, 0.0000, 0.0000],
		[0.5548, 0.4452, 0.0000, 0.0000]],
		[[0.3716, 0.3926, 0.2358, 0.0000],
		[0.3455, 0.3337, 0.3208, 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.4125, 0.3273, 0.2602, 0.0000]],
		[[0.5254, 0.4746, 0.0000, 0.0000],
		[0.3117, 0.2130, 0.1801, 0.2952]]])

10.3.2 加性注意力

一般来说，当查询和键是不同长度的矢量时，可以使用加性注意力作为评分函数。给定查询 $q ∈ R^q$ 和键 $k ∈ R^k$ ，加性注意力（additive attention）的评分函数为：
在这里插入图片描述
其中可学习的参数是 $W_q ∈ R^{h×q}$ 、 $W_k ∈ R^{h×k}$ 和 $W_v ∈ R^{h}$ 。如 (10.3.3)所示，将查询和键连结起来后输入到一个多层感知机（MLP）中，感知机包含一个隐藏层，其隐藏单元数是一个超参数h。通过使用tanh作为激活函数，并且禁用偏置项。
下面来实现加性注意力。

#@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的形状：(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)

用一个小例子来演示上面的AdditiveAttention类，其中查询、键和值的形状为（批量大小，步数或词元序列长度，特征大小），实际输出为(2, 1, 20)、(2, 10, 2)和(2, 10, 4)。注意力汇聚输出的形状为（批量大小，查询的步数，值的维度）。

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

在这里插入图片描述

10.3.3 缩放点积注意力

使用点积可以得到计算效率更高的评分函数，但是点积操作要求查询和键具有相同的长度d。假设查询和键的所有元素都是独立的随机变量，并且都满足零均值和单位方差，那么两个向量的点积的均值为0，方差为d。为确保无论向量长度如何，点积的方差在不考虑向量长度的情况下仍然是1，我们再将点积除以√d，则缩放点积注意力（scaled dot‐product attention）评分函数为：
在这里插入图片描述
在实践中，通常从小批量的角度来考虑提高效率，例如基于 $n$ 个查询和 $m$ 个键－值对计算注意力，其中查询和键的长度为 $d$ ，值的长为v。查询 $Q ∈ R^{n×d}$ 、键 $K ∈ R^{m×d}$ 和值 $V ∈ R^{m×v}$ 的缩放点积注意力是：
在这里插入图片描述
下面的缩放点积注意力的实现使用了暂退法进行模型正则化。

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

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

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

小结

将注意力汇聚的输出计算可以作为值的加权平均，选择不同的注意力评分函数会带来不同的注意力汇聚操作。
当查询和键是不同长度的矢量时，可以使用可加性注意力评分函数。当它们的长度相同时，使用缩放的“点－积”注意力评分函数的计算效率更高。

10.4 Bahdanau 注意力

9.7节中探讨了机器翻译问题：通过设计一个基于两个循环神经网络的编码器‐解码器架构，用于序列到序列学习。具体来说，

循环神经网络编码器将长度可变的序列转换为固定形状的上下文变量，
然后循环神经网络解码器根据生成的词元和上下文变量按词元生成输出（目标）序列词元。

然而，即使并非所有输入（源）词元都对解码某个词元都有用，在每个解码步骤中仍使用编码相同的上下文变量。有什么方法能改变上下文变量呢？

Bahdanau等人提出了一个没有严格单向对齐限制的可微注意力模型 (Bahdanau et al., 2014)。在预测词元时，如果不是所有输入词元都相关，模型将仅对齐（或参与）输入序列中与当前预测相关的部分。这是通过将上下文变量视为注意力集中的输出来实现的。

10.4.1 模型

假设输入序列中有 $T$ 个词元，解码时间步 $t^′$ 的上下文变量是注意力集中的输出：
在这里插入图片描述
其中，时间步 $t^′ − 1$ 时的解码器隐状态 $s_{t^′−1}$ 是查询，编码器隐状态 $h_t$ 既是键，也是值，注意力权重 $α$ 是使用(10.3.2) 所定义的加性注意力打分函数计算的。与图9.7.2中的循环神经网络编码器‐解码器架构略有不同，图10.4.1描述了Bahdanau注意力的架构。
在这里插入图片描述
图10.4.1: 一个带有Bahdanau注意力的循环神经网络编码器‐解码器模型

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

10.4.2 定义注意力解码器

下面看看如何定义Bahdanau注意力，实现循环神经网络编码器‐解码器。为了更方便地显示学习的注意力权重，以下AttentionDecoder类定义了带有注意力机制解码器的基本接口。

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

在接下来的Seq2SeqAttentionDecoder类中实现带有Bahdanau注意力的循环神经网络解码器。

首先，初始化解码器的状态，需要下面的输入：

编码器在所有时间步的最终层隐状态，将作为注意力的键和值；
上一时间步的编码器全层隐状态，将作为初始化解码器的隐状态；
编码器有效长度（排除在注意力池中填充词元）。

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

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 训练

与 9.7.4节类似，这里指定超参数，实例化一个带有Bahdanau注意力的编码器和解码器，并对这个模型进行机器翻译训练。由于新增的注意力机制，训练要比没有注意力机制的 9.7.4节慢得多。

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.021, 4948.7 tokens/sec on cuda:0

模型训练后，我们用它将几个英语句子翻译成法语并计算它们的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 paresseux ., bleu 0.658
i'm home . => je suis chez moi ., bleu 1.000

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

在这里插入图片描述
小结

在预测词元时，如果不是所有输入词元都是相关的，那么具有Bahdanau注意力的循环神经网络编码器‐解码器会有选择地统计输入序列的不同部分。这是通过将上下文变量视为加性注意力池化的输出来实现的。
在循环神经网络编码器‐解码器中，Bahdanau注意力将上一时间步的解码器隐状态视为查询，在所有时间步的编码器隐状态同时视为键和值。

10.5 多头注意力

在实践中，当给定相同的查询、键和值的集合时，希望模型可以基于相同的注意力机制学习到不同的行为，然后将不同的行为作为知识组合起来，捕获序列内各种范围的依赖关系（例如，短距离依赖和长距离依赖关系）。因此，允许注意力机制组合使用查询、键和值的不同子空间表示（representation subspaces）可能是有益的。

为此，与其只使用单独一个注意力汇聚，可以用独立学习得到的 $h$ 组不同的线性投影（linear projections）来变换查询、键和值。然后，这 $h$ 组变换后的查询、键和值将并行地送到注意力汇聚中。最后，将这 $h$ 个注意力汇聚的输出拼接在一起，并且通过另一个可以学习的线性投影进行变换，以产生最终输出。这种设计被称为多头注意力（multihead attention）(Vaswani et al., 2017)。对于 $h$ 个注意力汇聚输出，每一个注意力汇聚都被称作一个头（head）。图10.5.1 展示了使用全连接层来实现可学习的线性变换的多头注意力。
在这里插入图片描述图10.5.1: 多头注意力：多个头连结然后线性变换

10.5.1 模型

用数学语言将这个模型形式化地描述出来。给定查询 $q ∈ R^{d_q}$ 、键 $k ∈ R^{d_k}$ 和值 $v ∈ R^{d_v}$ ，每个注意力头 $h_i（i = 1, . . . , h）$ 的计算方法为：
在这里插入图片描述

其中， $f$ 代表注意力汇聚的函数。 $f$ 可以是 10.3节中的加性注意力和缩放点积注意力。多头注意力的输出需要经过另一个线性转换，它对应着 $h$ 个头连结后的结果，因此其可学习参数是 $W_o ∈ R^{p_o}×h^{p_v}$ ：
在这里插入图片描述
基于这种设计，每个头都可能会关注输入的不同部分，可以表示比简单加权平均值更复杂的函数。

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

10.5.2 实现

在实现过程中通常选择缩放点积注意力作为每一个注意力头。为了避免计算代价和参数代价的大幅增长，我们设定 $p_q = p_k = p_v = p_o/h$ 。值得注意的是，如果将查询、键和值的线性变换的输出数量设置为 $p_qh = p_kh =p_vh = p_o$ ，则可以并行计算 $h$ 个头。在下面的实现中， $p_o$ 是通过参数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)

为了能够使多个头并行计算，上面的MultiHeadAttention类将使用下面定义的两个转置函数。具体来说，transpose_output函数反转了transpose_qkv函数的操作。

#@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_size，num_queries，num_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.7 Transformer

自注意力同时具有并行计算和最短的最大路径长度这两个优势。因此，使用自注意力来设计深度架构是很有吸引力的。对比之前仍然依赖循环神经网络实现输入表示的自注意力模型 (Cheng et al., 2016, Lin et al., 2017,Paulus et al., 2017)，Transformer模型完全基于注意力机制，没有任何卷积层或循环神经网络层 (Vaswani etal., 2017)。

10.7.1 模型

Transformer作为编码器－解码器架构的一个实例，其整体架构图在图10.7.1中展示。正如所见到的，Transformer是由编码器和解码器组成的。Transformer的编码器和解码器是基于自注意力的模块叠加而成的，源（输入）序列和目标（输出）序列的嵌入（embedding）表示将加上位置编码（positional encoding），再分别输入到编码器和解码器中。
在这里插入图片描述
图10.7.1: transformer架构

图10.7.1中概述了Transformer的架构。

从宏观角度来看，Transformer的编码器是由多个相同的层叠加而成的，每个层都有两个子层（子层表示为sublayer）。

第一个子层是多头自注意力（multi‐head self‐attention）汇聚；
第二个子层是基于位置的前馈网络（positionwise feed‐forward network）。

具体来说，在计算编码器的自注意力时，查询、键和值都来自前一个编码器层的输出。每个子层都采用了残差连接（residual connection）。
在Transformer中，对于序列中任何位置的任何输入 $x ∈ R^d$ ，都要求满足 $sublayer(x) ∈ R^d$ ，以便残差连接满足 $x + sublayer(x) ∈ R^d$ 。
在残差连接的加法计算之后，紧接着应用层规范化（layer normalization）(Ba et al., 2016)。因此，输入序列对应的每个位置，Transformer编码器都将输出一个d维表示向量。

Transformer解码器也是由多个相同的层叠加而成的，并且层中使用了残差连接和层规范化。除了编码器中描述的两个子层之外，解码器还在这两个子层之间插入了第三个子层，称为编码器－解码器注意力（encoder‐decoder attention）层。在编码器－解码器注意力中，查询来自前一个解码器层的输出，而键和值来自整个编码器的输出。在解码器自注意力中，查询、键和值都来自上一个解码器层的输出。但是，解码器中的每个位置只能考虑该位置之前的所有位置。这种掩蔽（masked）注意力保留了自回归（auto‐regressive）属性，确保预测仅依赖于已生成的输出词元。

import math
import pandas as pd
import torch
from torch import nn
from d2l import torch as d2l

10.7.2 基于位置的前馈网络

基于位置的前馈网络对序列中的所有位置的表示进行变换时使用的是同一个多层感知机（MLP），这就是称前馈网络是基于位置的（positionwise）的原因。在下面的实现中，输入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)))

下面的例子显示，改变张量的最里层维度的尺寸，会改变成基于位置的前馈网络的输出尺寸。因为用同一个多层感知机对所有位置上的输入进行变换，所以当所有这些位置的输入相同时，它们的输出也是相同的。

ffn = PositionWiseFFN(4, 4, 8)
ffn.eval()
ffn(torch.ones((2, 3, 4)))[0]
==================================
tensor([[-0.8290, 1.0067, 0.3619, 0.3594, -0.5328, 0.2712, 0.7394, 0.0747],
	[-0.8290, 1.0067, 0.3619, 0.3594, -0.5328, 0.2712, 0.7394, 0.0747],
	[-0.8290, 1.0067, 0.3619, 0.3594, -0.5328, 0.2712, 0.7394, 0.0747]],
grad_fn=<SelectBackward0>)

10.7.3 残差连接和层规范化

现在让我们关注图10.7.1中的加法和规范化（add&norm）组件。正如在本节开头所述，这是由残差连接和紧随其后的层规范化组成的。两者都是构建有效的深度架构的关键。7.5节中解释了在一个小批量的样本内基于批量规范化对数据进行重新中心化和重新缩放的调整。层规范化和批量规范化的目标相同，但层规范化是基于特征维度进行规范化。尽管批量规范化在计算机视觉中被广泛应用，但在自然语言处理任务中（输入通常是变长序列）批量规范化通常不如层规范化的效果好。以下代码对比不同维度的层规范化和批量规范化的效果。

ln = nn.LayerNorm(2)
bn = nn.BatchNorm1d(2)
X = torch.tensor([[1, 2], [2, 3]], dtype=torch.float32)
# 在训练模式下计算X的均值和方差
print('layer norm:', ln(X), '\nbatch norm:', bn(X))

layer norm: tensor([[-1.0000, 1.0000],[-1.0000, 1.0000]], grad_fn=<NativeLayerNormBackward0>)
batch norm: tensor([[-1.0000, -1.0000],[ 1.0000, 1.0000]], grad_fn=<NativeBatchNormBackward0>)

现在可以使用残差连接和层规范化来实现AddNorm类。暂退法也被作为正则化方法使用。

#@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编码器的基础组件，现在可以先实现编码器中的一个层。下面的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编码器中的任何层都不会改变其输入的形状。

X = torch.ones((2, 100, 24))
valid_lens = torch.tensor([3, 2])
encoder_blk = EncoderBlock(24, 24, 24, 24, [100, 24], 24, 48, 8, 0.5)
encoder_blk.eval()
encoder_blk(X, valid_lens).shape
==========================================
torch.Size([2, 100, 24])

下面实现的Transformer编码器的代码中，堆叠了num_layers个EncoderBlock类的实例。由于这里使用的是值范围在−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之间，
		# 因此嵌入值乘以嵌入维度的平方根进行缩放，
		# 然后再与位置编码相加。
		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

下面指定了超参数来创建一个两层的Transformer编码器。Transformer编码器输出的形状是（批量大小，时间步数目，num_hiddens）。

encoder = TransformerEncoder(200, 24, 24, 24, 24, [100, 24], 24, 48, 8, 2, 0.5)
encoder.eval()
encoder(torch.ones((2, 100), dtype=torch.long), valid_lens).shape
=======================
torch.Size([2, 100, 24])

10.7.5 解码器

如图10.7.1所示，Transformer解码器也是由多个相同的层组成。在DecoderBlock类中实现的每个层包含了三个子层：解码器自注意力、“编码器‐解码器”注意力和基于位置的前馈网络。这些子层也都被残差连接和紧随的层规范化围绕。正如在本节前面所述，在掩蔽多头解码器自注意力层（第一个子层）中，查询、键和值都来自上一个解码器层的输出。关于序列到序列模型（sequence‐to‐sequence model），在训练阶段，其输出序列的所有位置（时间步）的词元都是已知的；然而，在预测阶段，其输出序列的词元是逐个生成的。因此，在任何解码器时间步中，只有生成的词元才能用于解码器的自注意力计算中。为了在解码器中保留自回归的属性，其掩蔽自注意力设定了参数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):
		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。


decoder_blk = DecoderBlock(24, 24, 24, 24, [100, 24], 24, 48, 8, 0.5, 0)
decoder_blk.eval()
X = torch.ones((2, 100, 24))
state = [encoder_blk(X, valid_lens), valid_lens, [None]]
decoder_blk(X, state)[0].shape
=========================
torch.Size([2, 100, 24])

现在我们构建了由num_layers个DecoderBlock实例组成的完整的Transformer解码器。最后，通过一个全连接层计算所有vocab_size个可能的输出词元的预测值。解码器的自注意力权重和编码器解码器注意力权重都被存储下来，方便日后可视化的需要。

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 训练

依照Transformer架构来实例化编码器－解码器模型。在这里，指定Transformer的编码器和解码器都是2层，都使用4头注意力。与 9.7.4节类似，为了进行序列到序列的学习，下面在“英语－法语”机器翻译数据集上训练Transformer模型。

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
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, 5202.9 tokens/sec on cuda:0

在这里插入图片描述
训练结束后，使用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))

在这里插入图片描述

为了可视化解码器的自注意力权重和“编码器－解码器”的注意力权重，我们需要完成更多的数据操作工作。例如用零填充被掩蔽住的注意力权重。值得注意的是，解码器的自注意力权重和“编码器－解码器”的注意力权重都有相同的查询：即以序列开始词元（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))

在这里插入图片描述
与编码器的自注意力的情况类似，通过指定输入序列的有效长度，输出序列的查询不会与输入序列中填充位置的词元进行注意力计算。

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