注意力汇聚：Nadaraya-Watson 核回归
上节我们介绍了框架下的注意力机制的主要成分图10.1.3： 查询（自主提示）和键（非自主提示）之间的交互形成了注意力汇聚， 注意力汇聚有选择地聚合了值（感官输入）以生成最终的输出。 在本节中，我们将介绍注意力汇聚的更多细节， 以便从宏观上了解注意力机制在实践中的运作方式。 具体来说，1964年提出的Nadaraya-Watson核回归模型 是一个简单但完整的例子，可以用于演示具有注意力机制的机器学习。

In [ ]
import paddle
from paddle import nn
from d2l import paddle as d2l
[生成数据集]
简单起见，考虑下面这个回归问题： 给定的成对的“输入－输出”数据集 {(x1,y1),…,(xn,yn)}{(x_1, y_1), \ldots, (x_n, y_n)}{(x
1
​
,y
1
​
),…,(x
n
​
,y
n
​
)}， 如何学习fff来预测任意新输入xxx的输出y^=f(x)\hat{y} = f(x)
y
^
​
=f(x)？

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

yi=2sin⁡(xi)+xi0.8+ϵ, (10.2.1)y_i = 2\sin(x_i) + x_i^{0.8} + \epsilon,~~~~~~~~~~(10.2.1)
y
i
​
=2sin(x
i
​
)+x
i
0.8
​
+ϵ, (10.2.1)

其中ϵ\epsilonϵ服从均值为000和标准差为0.50.50.5的正态分布。 我们生成了505050个训练样本和505050个测试样本。 为了更好地可视化之后的注意力模式，我们将训练样本进行排序。

In [ ]
n_train = 50 # 训练样本数
x_train = paddle.sort(paddle.rand([n_train]) * 5) # 排序后的训练样本
In [ ]
def f(x):
return 2 * paddle.sin(x) + x**0.8

y_train = f(x_train) + paddle.normal(0.0, 0.5, (n_train,)) # 训练样本的输出
x_test = paddle.arange(0, 50, 1) * 0.1 # 测试样本
y_truth = f(x_test) # 测试样本的真实输出
n_test = len(x_test) # 测试样本数
n_test
下面的函数将绘制所有的训练样本（样本由圆圈表示）， 不带噪声项的真实数据生成函数fff（标记为“Truth”）， 以及学习得到的预测函数（标记为“Pred”）。

In [ ]
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);
平均汇聚
我们先使用最简单的估计器来解决回归问题： 基于平均汇聚来计算所有训练样本输出值的平均值：

f(x)=1n∑i=1nyi, (10.2.2)f(x) = \frac{1}{n}\sum_{i=1}^n y_i,~~~~~~~~~~(10.2.2)
f(x)=
n
1
​

i=1
∑
n
​
y
i
​
, (10.2.2)

如下图所示，这个估计器确实不够聪明： 真实函数fff（“Truth”）和预测函数（“Pred”）相差很大。

In [ ]
y_hat = paddle.tile(y_train.mean(), [n_test])
plot_kernel_reg(y_hat)
[非参数注意力汇聚]
显然，平均汇聚忽略了输入xix_ix
i
​
。 于是NadarayaNadaraya.1964和 WatsonWatson.1964提出了一个更好的想法， 根据输入的位置对输出yiy_iy
i
​
进行加权：

f(x)=∑i=1nK(x−xi)∑j=1nK(x−xj)yi, (10.2.3)f(x) = \sum_{i=1}^n \frac{K(x - x_i)}{\sum_{j=1}^n K(x - x_j)} y_i,~~~~~~~~~~(10.2.3)
f(x)=
i=1
∑
n
​

∑
j=1
n
​
K(x−x
j
​
)
K(x−x
i
​
)
​
y
i
​
, (10.2.3)

其中KKK是核（kernel）。 公式(10.2.3)所描述的估计器被称为 Nadaraya-Watson核回归（Nadaraya-Watson kernel regression）。 这里我们不会深入讨论核函数的细节， 但受此启发， 我们可以从图10.1.3中的注意力机制框架的角度 重写公式(10.2.3)， 成为一个更加通用的注意力汇聚（attention pooling）公式：

f(x)=∑i=1nα(x,xi)yi, (10.2.4)f(x) = \sum_{i=1}^n \alpha(x, x_i) y_i,~~~~~~~~~~(10.2.4)
f(x)=
i=1
∑
n
​
α(x,x
i
​
)y
i
​
, (10.2.4)

其中xxx是查询，(xi,yi)(x_i, y_i)(x
i
​
,y
i
​
)是键值对。 比较公式(10.2.4)和公式(10.2.2)， 注意力汇聚是yiy_iy
i
​
的加权平均。 将查询xxx和键xix_ix
i
​
之间的关系建模为 注意力权重（attention weight）α(x,xi)\alpha(x, x_i)α(x,x
i
​
)， 如公式(10.2.4)所示， 这个权重将被分配给每一个对应值yiy_iy
i
​
。 对于任何查询，模型在所有键值对注意力权重都是一个有效的概率分布： 它们是非负的，并且总和为1。

为了更好地理解注意力汇聚， 我们考虑一个高斯核（Gaussian kernel），其定义为：

K(u)=12πexp⁡(−u22). (10.2.5)K(u) = \frac{1}{\sqrt{2\pi}} \exp(-\frac{u^2}{2}).~~~~~~~~~~(10.2.5)
K(u)=
2π
​

1
​
exp(−
2
u
2

​
). (10.2.5)

将高斯核代入公式(10.2.4)和 公式(10.2.3)可以得到：

f(x)=∑i=1nα(x,xi)yi=∑i=1nexp⁡(−12(x−xi)2)∑j=1nexp⁡(−12(x−xj)2)yi=∑i=1nsoftmax(−12(x−xi)2)yi. (10.2.6)\begin{aligned} 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{aligned}~~~~~~~~~~(10.2.6)
f(x)
​

=
i=1
∑
n
​
α(x,x
i
​
)y
i
​

=
i=1
∑
n
​

∑
j=1
n
​
exp(−
2
1
​
(x−x
j
​
)
2
)
exp(−
2
1
​
(x−x
i
​
)
2
)
​
y
i
​

=
i=1
∑
n
​
softmax(−
2
1
​
(x−x
i
​
)
2
)y
i
​
.
​
(10.2.6)

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

值得注意的是，Nadaraya-Watson核回归是一个非参数模型。 因此公式(10.2.6)是 非参数的注意力汇聚（nonparametric attention pooling）模型。 接下来，我们将基于这个非参数的注意力汇聚模型来绘制预测结果。 你会发现新的模型预测线是平滑的，并且比平均汇聚的预测更接近真实。

In [ ]

X_repeat的形状:(n_test,n_train),

每一行都包含着相同的测试输入（例如：同样的查询）

X_repeat = x_test.reshape((50,1)).tile([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, axis=1)

y_hat的每个元素都是值的加权平均值，其中的权重是注意力权重

y_hat = paddle.matmul(attention_weights, y_train)
plot_kernel_reg(y_hat)
现在，我们来观察注意力的权重。 这里测试数据的输入相当于查询，而训练数据的输入相当于键。 因为两个输入都是经过排序的，因此由观察可知“查询-键”对越接近， 注意力汇聚的[注意力权重]就越高。

In [ ]
d2l.show_heatmaps(attention_weights.unsqueeze(0).unsqueeze(0),
xlabel=‘Sorted training inputs’,
ylabel=‘Sorted testing inputs’)
[带参数注意力汇聚]
非参数的Nadaraya-Watson核回归具有一致性（consistency）的优点： 如果有足够的数据，此模型会收敛到最优结果。 尽管如此，我们还是可以轻松地将可学习的参数集成到注意力汇聚中。

例如，与公式(10.2.6)略有不同， 在下面的查询xxx和键xix_ix
i
​
之间的距离乘以可学习参数www：

f(x)=∑i=1nα(x,xi)yi=∑i=1nexp⁡(−12((x−xi)w)2)∑j=1nexp⁡(−12((x−xj)w)2)yi=∑i=1nsoftmax(−12((x−xi)w)2)yi. (10.2.7)\begin{aligned}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{aligned}~~~~~~~~~~(10.2.7)
f(x)
​

=
i=1
∑
n
​
α(x,x
i
​
)y
i
​

=
i=1
∑
n
​

∑
j=1
n
​
exp(−
2
1
​
((x−x
j
​
)w)
2
)
exp(−
2
1
​
((x−x
i
​
)w)
2
)
​
y
i
​

=
i=1
∑
n
​
softmax(−
2
1
​
((x−x
i
​
)w)
2
)y
i
​
.
​
(10.2.7)

在本节的余下部分，我们将通过训练这个模型公式(10.2.7)来学习注意力汇聚的参数。

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

假设第一个小批量数据包含nnn个矩阵X1,…,Xn\mathbf{X}_1,\ldots, \mathbf{X}_nX
1
​
,…,X
n
​
， 形状为a×ba\times ba×b， 第二个小批量包含nnn个矩阵Y1,…,Yn\mathbf{Y}_1, \ldots, \mathbf{Y}_nY
1
​
,…,Y
n
​
， 形状为b×cb\times cb×c。 它们的批量矩阵乘法得到nnn个矩阵 X1Y1,…,XnYn\mathbf{X}_1\mathbf{Y}_1, \ldots, \mathbf{X}_n\mathbf{Y}_nX
1
​
Y
1
​
,…,X
n
​
Y
n
​
， 形状为a×ca\times ca×c。 因此，[假定两个张量的形状分别是(n,a,b)(n,a,b)(n,a,b)和(n,b,c)(n,b,c)(n,b,c)， 它们的批量矩阵乘法输出的形状为(n,a,c)(n,a,c)(n,a,c)]。

In [ ]
X = paddle.ones((2, 1, 4))
Y = paddle.ones((2, 4, 6))
paddle.bmm(X, Y).shape
在注意力机制的背景中，我们可以[使用小批量矩阵乘法来计算小批量数据中的加权平均值]。

In [ ]
weights = paddle.ones((2, 10)) * 0.1
values = paddle.arange(20, dtype=‘float32’).reshape((2,10))
paddle.bmm(weights.unsqueeze(1), values.unsqueeze(-1))
定义模型
基于公式(10.2.7)中的 [带参数的注意力汇聚]，使用小批量矩阵乘法， 定义Nadaraya-Watson核回归的带参数版本为：

In [ ]
class NWKernelRegression(nn.Layer):
def init(self, **kwargs):
super().init(**kwargs)
self.w = paddle.create_parameter((1,), dtype=‘float32’)

def forward(self, queries, keys, values):
    # queries和attention_weights的形状为(查询个数，“键－值”对个数)
    queries = queries.reshape((queries.shape[0], 1)).tile([keys.shape[1]]).reshape((-1, keys.shape[1]))
    self.attention_weight = nn.functional.softmax(
        -((queries - keys) * self.w)**2 / 2, axis=1)
    # values的形状为(查询个数，“键－值”对个数)
    return paddle.bmm(self.attention_weight.unsqueeze(1),
                      values.unsqueeze(-1)).reshape((-1, ))

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

In [ ]

X_tile的形状:(n_train，n_train)，每一行都包含着相同的训练输入

X_tile = x_train.tile([n_train, 1])

Y_tile的形状:(n_train，n_train)，每一行都包含着相同的训练输出

Y_tile = y_train.tile([n_train, 1])

keys的形状:(‘n_train’，‘n_train’-1)

keys = X_tile[(1 - paddle.eye(n_train)).astype(paddle.bool)].reshape((n_train, -1))

values的形状:(‘n_train’，‘n_train’-1)

values = Y_tile[(1 - paddle.eye(n_train)).astype(paddle.bool)].reshape((n_train, -1))
[训练带参数的注意力汇聚模型]时，使用平方损失函数和随机梯度下降。

In [ ]
net = NWKernelRegression()
loss = nn.MSELoss(reduction=‘none’)
trainer = paddle.optimizer.SGD(learning_rate=0.5, parameters=net.parameters())
animator = d2l.Animator(xlabel=‘epoch’, ylabel=‘loss’, xlim=[1,5])

for epoch in range(5):
trainer.clear_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()))
如下所示，训练完带参数的注意力汇聚模型后，我们发现： 在尝试拟合带噪声的训练数据时， [预测结果绘制]的线不如之前非参数模型的平滑。

In [ ]

keys的形状:(n_test，n_train)，每一行包含着相同的训练输入（例如，相同的键）

keys = x_train.tile([n_test, 1])

value的形状:(n_test，n_train)

values = y_train.tile([n_test, 1])
y_hat = net(x_test, keys, values).unsqueeze(1).detach()
plot_kernel_reg(y_hat)
为什么新的模型更不平滑了呢？ 我们看一下输出结果的绘制图： 与非参数的注意力汇聚模型相比， 带参数的模型加入可学习的参数后， [曲线在注意力权重较大的区域变得更不平滑]。

In [ ]
d2l.show_heatmaps(net.attention_weight.unsqueeze(0).unsqueeze(0),
xlabel=‘Sorted training inputs’,
ylabel=‘Sorter testing, inputs’)
小结
Nadaraya-Watson核回归是具有注意力机制的机器学习范例。
Nadaraya-Watson核回归的注意力汇聚是对训练数据中输出的加权平均。从注意力的角度来看，分配给每个值的注意力权重取决于将值所对应的键和查询作为输入的函数。
注意力汇聚可以分为非参数型和带参数型。
练习
增加训练数据的样本数量，你能否得到更好的非参数的Nadaraya-Watson核回归模型？
在带参数的注意力汇聚的实验中学习得到的参数www的价值是什么？为什么在可视化注意力权重时，它会使加权区域更加尖锐？
如何将超参数添加到非参数的Nadaraya-Watson核回归中以实现更好地预测结果？
为本节的核回归设计一个新的带参数的注意力汇聚模型。训练这个新模型并可视化其注意力权重。
如果想系统性学习该项目，可前往“动手学AI”课程（https://aistudio.baidu.com/aistudio/course/introduce/25851）查看完整章节

Discussions