动手学深度学习 d2l-zh 注意力汇聚实战：用 Nadaraya-Watson 核回归读懂注意力机制的核心原理-平芜编程栈

人工智能
深度学习
机器学习
教程

【免费下载链接】d2l-zh

《动手学深度学习》：面向中文读者、能运行、可讨论。中英文版被70多个国家的500多所大学用于教学。

项目地址：https://gitcode.com/GitHub_Trending/d2/d2l-zh

注意力机制是当代深度学习（从机器翻译到 Transformer）的基石，而《动手学深度学习》（d2l-zh）用 1964 年提出的 Nadaraya-Watson 核回归作为"最小但完整"的教学案例，让读者在引入任何神经网络之前，就能直观看到"查询（query）—键（key）—值（value）"三者如何通过注意力权重交互。本篇以仓库中 nadaraya-waston_origin.md（中文版见 nadaraya-waston.md）为骨架，结合 d2l 工具库源码，完整还原从平均汇聚、非参数注意力汇聚到可学习参数的完整实验链路。读完本文，你将能亲手复现注意力热图、理解批量矩阵乘法在注意力计算中的角色，并掌握非参数与参数化注意力汇聚的差异。

从 QKV 框架到注意力汇聚：本节在全书中的定位

在进入代码之前，先回顾上一节（attention-cues.md）建立的注意力机制框架：查询（query，即自主提示）与键（key，即非自主提示）之间的交互形成注意力汇聚（attention pooling）；注意力汇聚有选择地聚合值（value，即感官输入）以生成输出。上图正是这一框架的示意图。

本节的目的不是引入新网络，而是把注意力汇聚讲深讲透：

用Nadaraya-Watson 核回归这一 1964 年的经典非参数方法，展示"加权平均"是如何自然涌现出注意力行为的；
给出非参数注意力汇聚（Gaussian 核 + softmax）的完整可运行实现；
再引入一个可学习参数，将其升级为参数化注意力汇聚，并用随机梯度下降训练；
借助热图可视化，直观比较两种模型注意力权重的分布差异。

d2l-zh 工具库为这一实验提供了全部底层支撑。以 PyTorch 后端为例，本文用到的d2l.rand、d2l.normal、d2l.sin、d2l.matmul、d2l.ones、d2l.arange等并不是独立实现的函数，而是 d2l/torch.py 中由config.ini驱动的别名映射（matmul = torch.matmul、rand = torch.rand、normal = torch.normal等），这样同一份章节代码可以在 MXNet、PyTorch、TensorFlow、PaddlePaddle 四个后端间无缝切换——这也是全书代码可通过#@tab标签多框架运行的基础。

实验环境只需要导入工具库与框架：

#@tab pytorch from d2l import torch as d2l import torch from torch import nn

（MXNet 后端还需from mxnet import autograd, gluon, np, npx并调用npx.set_np()；TensorFlow 后端建议tf.random.set_seed(seed=1322)固定随机种子以保证可复现，详见中文版文档。）

生成数据集：一个带噪声的非线性回归问题

为聚焦注意力本身，构造一个最简单的回归设定：给定输入-输出对 ${(x_1, y_1), \ldots, (x_n, y_n)}$，学习函数 $f$ 以预测任意新输入 $x$ 的输出 $\hat{y} = f(x)$。

数据由如下非线性函数叠加高斯噪声生成：

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

其中噪声 $\epsilon$ 服从均值为 0、标准差为 0.5 的正态分布。共生成 50 个训练样本和 50 个测试样本；为了让后续注意力模式可视化更清晰，训练输入被显式排序：

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

#@tab all def f(x): return 2 * d2l.sin(x) + x**0.8 y_train = f(x_train) + d2l.normal(0.0, 0.5, (n_train,)) # 训练输出 x_test = d2l.arange(0, 5, 0.1) # 测试样本 y_truth = f(x_test) # 测试样本的真实输出 n_test = len(x_test) # 测试样本数

这里测试输入从 0 以 0.1 为步长取到 5，共 50 个。下面的plot_kernel_reg辅助函数将同时绘制：所有训练样本（圆圈标记）、不含噪声的真实函数f（图例 "Truth"）以及学习得到的预测函数（图例 "Pred"），xlim/ylim固定坐标范围便于不同模型间对比：

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

d2l.plot的完整实现位于 d2l/torch.py，它内部复用set_figsize、set_axes等绘图基础设施，并约定 SVG 格式输出，确保 Jupyter 中跨平台一致渲染。

平均汇聚：全世界"最笨"的估计器

在引入任何智能机制前，先看一个极端朴素的基线：无视所有输入 $x_i$，直接用训练输出的均值预测任意 $x$：

$$f(x) = \frac{1}{n}\sum_{i=1}^n y_i.$$

#@tab pytorch y_hat = torch.repeat_interleave(y_train.mean(), n_test) plot_kernel_reg(y_hat)

绘制结果是一条水平直线，与真实函数 $f$ 的曲线形态相差甚远。它的数学意义在于：平均汇聚可以看作注意力汇聚的特例——所有 $y_i$ 的权重恒等于 $1/n$，与查询 $x$ 无关。这也为后文的加权平均埋下伏笔。

非参数注意力汇聚：高斯核与 softmax 的天然结合

平均汇聚的致命缺陷是忽略了输入 $x_i$。Nadaraya（1964）与 Watson（1964）分别提出改进：根据输入位置对输出 $y_i$ 加权：

$$f(x) = \sum_{i=1}^n \frac{K(x - x_i)}{\sum_{j=1}^n K(x - x_j)} y_i,$$

其中 $K$ 是核函数（kernel）。该估计器即Nadaraya-Watson 核回归。这里不展开核函数的理论细节，而是从注意力机制框架（即上图 QKV 模型）出发，将其重写为更通用的注意力汇聚形式：

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

其中 $x$ 是查询，$(x_i, y_i)$ 构成键值对，$\alpha(x, x_i)$ 是分配给对应值 $y_i$ 的注意力权重。与平均汇聚对比可见：注意力汇聚本质上是值 $y_i$ 的加权平均。且对任意查询，所有键值对上的注意力权重构成一个合法的概率分布——非负且总和为 1（这正是 softmax 归一化的体现）。

高斯核推导：注意力权重如何自然涌现

考虑高斯核$K(u) = \frac{1}{\sqrt{2\pi}} \exp(-u^2/2)$，将其代入上式，常数因子 $\frac{1}{\sqrt{2\pi}}$ 在分子分母中约去，得到：

$$\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}$$

注意最后一步：exp(·)/Σexp(·)正是softmax 的精确数学定义。这意味着：越接近查询 $x$ 的键 $x_i$，其对应值 $y_i$ 被分配的注意力权重越大，即"获得更多注意力"。Nadaraya-Watson 核回归没有可训练参数，属于非参数模型，因此上式也是非参数注意力汇聚的典型实例。

向量化实现与注意力热图

借助广播与矩阵乘法，可一次性对所有测试查询完成计算。核心技巧是把每个测试输入复制n_train份，构造成(n_test, n_train)的矩阵：

#@tab pytorch # X_repeat 形状：(n_test, n_train)，每一行包含相同的测试输入（即相同的查询） X_repeat = d2l.reshape(x_test.repeat_interleave(n_train), (-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 = d2l.matmul(attention_weights, y_train) plot_kernel_reg(y_hat)

d2l.matmul即torch.matmul（见 d2l/torch.py 的别名定义），此处完成(n_test, n_train)与(n_train,)的矩阵-向量乘。预测曲线相比平均汇聚明显更平滑、更贴近真实函数。

随后用热图观察注意力权重分布：测试输入是查询（纵轴），训练输入是键（横轴），二者均有序排列：

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

show_heatmaps在 d2l/torch.py 中实现：对输入的 4 维张量（批次×行×列×矩阵）逐个子图调用imshow，默认配色cmap='Reds'，颜色越深代表权重越大。观察可知：查询-键对越接近，对应位置的注意力权重越高，热图中形成一条沿主对角线的高亮带——这正是"按相似度分配注意力"的可视化证据。该函数在 MXNet、TensorFlow、PaddlePaddle 后端中均有同名实现（见 d2l/mxnet.py、d2l/tensorflow.py、d2l/paddle.py），四框架 API 完全对齐。

参数化注意力汇聚：给距离乘上一个可学习参数

非参数 Nadaraya-Watson 核回归具备一致性（consistency）优势：数据足够多时收敛到最优解。但既然注意力权重 $\alpha$ 本质上是一个可微函数，就自然可以给它配备可学习参数。本节沿用 Gaussian 核思路，在查询-键距离上乘一个可学习标量 $w$：

$$\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}$$

批量矩阵乘法：注意力计算的高效引擎

参数化后需要在训练中反复批量计算注意力。深度学习框架提供的批量矩阵乘法（batch matrix multiplication，PyTorch 中为torch.bmm）一次处理 $n$ 对矩阵：给定形状 $(n, a, b)$ 与 $(n, b, c)$ 的两个张量，输出形状为 $(n, a, c)$，对应位置矩阵逐对相乘：

#@tab pytorch X = d2l.ones((2, 1, 4)) Y = d2l.ones((2, 4, 6)) torch.bmm(X, Y).shape # 输出 (2, 1, 6)

在注意力场景中，可用它高效计算小批量内值的加权平均——把权重向量变形为(2, 1, 10)、值向量变形为(2, 10, 1)，bmm一次输出(2, 1, 1)：

#@tab pytorch weights = d2l.ones((2, 10)) * 0.1 values = d2l.reshape(d2l.arange(20.0), (2, 10)) torch.bmm(weights.unsqueeze(1), values.unsqueeze(-1))

这正是注意力汇聚内部"权重 × 值求和"的标准实现路径，后续章节中 Transformer 的缩放点积注意力也沿用同一计算模式。

定义模型：NWKernelRegression

基于上式，用批量矩阵乘法实现参数化 Nadaraya-Watson 核回归模型。PyTorch 版本中w用nn.Parameter声明并自动纳入梯度追踪：

#@tab pytorch 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 = d2l.reshape( queries.repeat_interleave(keys.shape[1]), (-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)

前向过程分三步：① 将每个查询复制成与键数量一致的行，构成(查询数, 键数)的距离矩阵；② 用可学习参数w缩放距离后过 softmax，得到注意力权重并缓存在self.attention_weights（供后验可视化）；③ 用bmm对注意力权重与值做加权求和。MXNet 版本用nn.Block+self.params.get('w', shape=(1,))定义参数，逻辑完全一致（见 nadaraya-waston.md 中#@tab mxnet代码块）。

训练：留一法构造键值对 + SGD

参数化模型需要训练，关键在于如何构造训练数据。与"每个样本用自己预测自己"不同，这里采用留一法：任一训练样本的查询，都要与除自身以外的全部训练样本的键值对交互。实现上用单位矩阵eye的掩码剔除对角线：

#@tab pytorch # 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 = d2l.reshape(X_tile[(1 - d2l.eye(n_train)).type(torch.bool)], (n_train, -1)) # values 形状：(n_train, n_train - 1) values = d2l.reshape(Y_tile[(1 - d2l.eye(n_train)).type(torch.bool)], (n_train, -1))

随后用平方损失 + 随机梯度下降（学习率 0.5，5 轮）训练：

#@tab pytorch 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()))

Animator类在 d2l/torch.py 中实现，负责把逐轮的 loss 增量式绘制成折线动画（调用d2l.plt.subplots建图、add方法追加数据点并刷新显示）。原文档英文版在此处保留了 MXNet 的 L2 损失，而 PyTorch 用 MSE 损失——二者相差系数 2，故部分版本代码中对 loss 额外除以 2（详见英文原版 nadaraya-waston_origin.md 的#@tab pytorch块注释），读者对比多框架实现时需注意这一细节。

预测与注意力热图：参数化让加权区域更尖锐

训练完成后，用全部测试查询进行预测：将每个训练样本复制n_test份作为统一的键与值：

#@tab pytorch # keys 形状：(n_test, n_train)，每一列包含相同的训练输入（即相同的键） keys = x_train.repeat((n_test, 1)) # values 形状：(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)

预测曲线在试图拟合带噪训练数据时，明显比非参数版本的曲线"更不平滑"——这是参数化注意力汇聚的直接体现。再可视化net.attention_weights：

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

与非参数版本对比，参数化模型在注意力权重较大的区域变得更为尖锐：因为可学习的 $w$ 在训练中被放大，等价于把高斯核的带宽压缩，使权重更快地从高值跌落，注意力分布更集中于少数近邻键。

小结

Nadaraya-Watson 核回归是"机器学习 + 注意力机制"的经典范例：1964 年提出，却完美映射现代注意力框架。
其注意力汇聚本质是训练输出的加权平均；从注意力视角看，分配给每个值的注意力权重，是"以该值对应的键和查询为输入"的函数。
注意力汇聚分为非参数型（如高斯核版本，具备一致性、无需训练）与参数型（如本文的 $w$ 缩放版本，可端到端学习）。
批量矩阵乘法是批量注意力计算的关键引擎，本文bmm(权重, 值)的用法在后续章节（加性注意力、缩放点积注意力、Transformer）中会被反复复用。

练习与进阶方向

增加训练样本量：把n_train从 50 增大（如 100、500），非参数 Nadaraya-Watson 核回归是否会学得更好？结合"一致性"性质思考原因。
参数 $w$ 的语义：训练结束后打印net.w的值。为什么放大的 $w$ 会让注意力热图的加权区域更尖锐？提示：观察 $w$ 对exp(-((x - x_i)w)²/2)带宽的影响。
为核回归加超参数：如何在不引入神经网络的前提下，为非参数 Nadaraya-Watson 核回归添加超参数（如核带宽 $\sigma$ 或混合多个核）以提升预测？可以尝试构造 $\frac{1}{\sigma\sqrt{2\pi}}\exp(-u^2/2\sigma^2)$ 形式并手工调参。
设计新的参数化注意力汇聚：仿照本文为核回归设计另一类可学习注意力（例如把距离改为 $|x - x_i|$ 的指数衰减、或引入核函数选择矩阵），训练并可视化其注意力权重，与非参数/本文模型对比。

延伸阅读指引：本文是注意力机制章节的第二篇，建议按序阅读 attention-cues.md（QKV 框架与注意力提示）、attention-scoring-functions.md（加性注意力与缩放点积注意力）以及 self-attention-and-positional-encoding.md（自注意力）。工具库源码可查阅 d2l/torch.py 及其余后端实现；环境安装与运行方式见 chapter_installation/index.md。