PyTorch 2.0 实现 Nadaraya-Watson 核回归构建可复用的注意力层模块1. 理解注意力机制的本质在深度学习领域注意力机制的核心思想是模拟人类的选择性关注能力——从大量信息中筛选出关键部分进行处理。Nadaraya-Watson核回归作为最早的注意力形式之一完美诠释了这一理念通过计算查询与键的相似度来决定值的权重分配。让我们用一个直观的例子来说明假设你正在嘈杂的咖啡馆里专注于与朋友的对话。你的大脑会自动降低背景音乐、其他顾客谈话声的权重提高朋友语音信号的权重这种动态调整就是注意力机制的本质体现在数学表达上Nadaraya-Watson核回归可以表示为$$ f(x) \sum_{i1}^n \frac{K(x, x_i)}{\sum_{j1}^n K(x, x_j)} y_i $$其中$K(x, x_i)$是核函数衡量查询$x$与键$x_i$的相似度。当使用高斯核时这个公式就转化为我们熟悉的softmax加权形式。2. 数据准备与基础实现2.1 生成非线性数据集我们首先生成一个具有明显非线性特征的数据集加入适量噪声以模拟真实场景import torch import matplotlib.pyplot as plt n_train 50 # 训练样本数 x_train, _ torch.sort(torch.rand(n_train) * 5) # 排序后的训练输入 def true_function(x): return 2 * torch.sin(x) x**0.8 y_train true_function(x_train) torch.normal(0.0, 0.5, (n_train,)) # 带噪声的训练输出 x_test torch.arange(0, 5, 0.1) # 测试输入 y_true true_function(x_test) # 真实函数值 # 可视化 plt.figure(figsize(10, 6)) plt.scatter(x_train, y_train, labelTraining data, alpha0.7) plt.plot(x_test, y_true, r-, labelTrue function) plt.legend() plt.xlabel(x) plt.ylabel(y) plt.title(Generated Dataset) plt.show()2.2 非参数化注意力实现我们先实现一个非参数化的版本直接使用高斯核计算注意力权重def nonparametric_nw_kernel_regression(x_train, y_train, x_test, bandwidth1.0): 非参数化Nadaraya-Watson核回归 参数: x_train: 训练输入 (n_train,) y_train: 训练输出 (n_train,) x_test: 测试输入 (n_test,) bandwidth: 高斯核带宽 返回: y_pred: 预测输出 (n_test,) attention_weights: 注意力权重 (n_test, n_train) # 扩展维度以便广播计算 x_test_exp x_test.unsqueeze(1) # (n_test, 1) x_train_exp x_train.unsqueeze(0) # (1, n_train) # 计算高斯核注意力权重 distances -(x_test_exp - x_train_exp)**2 / (2 * bandwidth**2) attention_weights torch.softmax(distances, dim1) # (n_test, n_train) # 加权平均得到预测 y_pred torch.matmul(attention_weights, y_train) # (n_test,) return y_pred, attention_weights # 测试非参数化实现 y_pred, attn_weights nonparametric_nw_kernel_regression(x_train, y_train, x_test) # 可视化结果 plt.figure(figsize(10, 6)) plt.scatter(x_train, y_train, labelTraining data, alpha0.5) plt.plot(x_test, y_true, r-, labelTrue function) plt.plot(x_test, y_pred.detach(), g--, labelNW Prediction) plt.legend() plt.title(Nonparametric Nadaraya-Watson Kernel Regression) plt.show()3. 构建可训练的注意力层3.1 参数化注意力层实现现在我们将核回归封装为PyTorch模块引入可学习的带宽参数import torch.nn as nn class NadarayaWatsonAttention(nn.Module): def __init__(self, init_bandwidth1.0, learnableTrue): super().__init__() if learnable: # 使用log带宽确保正值 self.log_bandwidth nn.Parameter(torch.log(torch.tensor(init_bandwidth))) else: self.register_buffer(log_bandwidth, torch.log(torch.tensor(init_bandwidth))) self.learnable learnable property def bandwidth(self): return torch.exp(self.log_bandwidth) def forward(self, queries, keys, values): 参数: queries: (n_queries, query_dim) keys: (n_keys, key_dim) values: (n_keys, value_dim) 返回: output: (n_queries, value_dim) attention_weights: (n_queries, n_keys) # 计算查询-键相似度 distances torch.cdist(queries.unsqueeze(0), keys.unsqueeze(0)).squeeze(0) # (n_queries, n_keys) attention_scores -0.5 * (distances / self.bandwidth)**2 # 计算注意力权重 attention_weights torch.softmax(attention_scores, dim1) # (n_queries, n_keys) # 加权求和 output torch.matmul(attention_weights, values) # (n_queries, value_dim) return output, attention_weights3.2 训练注意力层我们构建完整的训练流程优化带宽参数# 准备训练数据 X_tile x_train.repeat((n_train, 1)) # (n_train, n_train) Y_tile y_train.repeat((n_train, 1)) # (n_train, n_train) # 创建掩码排除对角线元素避免自注意力 mask ~torch.eye(n_train, dtypetorch.bool) keys X_tile[mask].reshape(n_train, -1) # (n_train, n_train-1) values Y_tile[mask].reshape(n_train, -1) # (n_train, n_train-1) # 初始化模型和优化器 model NadarayaWatsonAttention(init_bandwidth0.5) criterion nn.MSELoss() optimizer torch.optim.SGD(model.parameters(), lr0.5) # 训练循环 n_epochs 100 losses [] for epoch in range(n_epochs): optimizer.zero_grad() # 前向传播 preds, _ model(x_train.unsqueeze(1), keys.unsqueeze(2), values.unsqueeze(2)) loss criterion(preds.squeeze(), y_train) # 反向传播 loss.backward() optimizer.step() losses.append(loss.item()) if (epoch1) % 10 0: print(fEpoch {epoch1}, Loss: {loss.item():.4f}, Bandwidth: {model.bandwidth.item():.4f}) # 绘制训练损失 plt.plot(losses) plt.xlabel(Epoch) plt.ylabel(MSE Loss) plt.title(Training Loss Curve) plt.show()4. 与多头注意力的对比分析4.1 实现PyTorch内置多头注意力为了进行公平比较我们使用PyTorch内置的多头注意力模块class MultiHeadAttentionWrapper(nn.Module): def __init__(self, embed_dim64, num_heads4): super().__init__() self.attention nn.MultiheadAttention(embed_dim, num_heads) self.query_proj nn.Linear(1, embed_dim) self.key_proj nn.Linear(1, embed_dim) self.value_proj nn.Linear(1, embed_dim) self.out_proj nn.Linear(embed_dim, 1) def forward(self, x, keys, values): # 投影到嵌入空间 q self.query_proj(x.unsqueeze(-1)).transpose(0, 1) # (seq_len, batch, embed_dim) k self.key_proj(keys.unsqueeze(-1)).transpose(0, 1) v self.value_proj(values.unsqueeze(-1)).transpose(0, 1) # 多头注意力 attn_output, _ self.attention(q, k, v) # 投影回原始空间 output self.out_proj(attn_output.transpose(0, 1)).squeeze(-1) return output # 训练多头注意力模型 mha_model MultiHeadAttentionWrapper() mha_optimizer torch.optim.Adam(mha_model.parameters(), lr0.001) mha_losses [] for epoch in range(n_epochs): mha_optimizer.zero_grad() preds mha_model(x_train, keys, values) loss criterion(preds, y_train) loss.backward() mha_optimizer.step() mha_losses.append(loss.item()) # 比较两种注意力机制 plt.figure(figsize(10, 5)) plt.plot(losses, labelNW Attention) plt.plot(mha_losses, labelMulti-Head Attention) plt.xlabel(Epoch) plt.ylabel(MSE Loss) plt.title(Comparison of Attention Mechanisms) plt.legend() plt.show()4.2 性能对比与适用场景通过实验对比我们可以总结两种注意力机制的差异特性Nadaraya-Watson 注意力多头注意力参数量极少(仅带宽参数)较多(投影矩阵注意力权重)计算复杂度O(n²)O(n²)可解释性高中等适合任务小规模回归、解释性要求高大规模序列建模训练数据需求少多特征提取能力有限强大在实际项目中NW注意力更适合需要模型解释性的场景小规模数据集简单的特征交互建模而多头注意力则更适合大规模预训练任务复杂的关系建模需要捕获长距离依赖的场景5. 高级应用与扩展5.1 整合到神经网络架构我们可以将NW注意力层作为更大网络的一部分class NWAttentionNetwork(nn.Module): def __init__(self, input_dim1, hidden_dim32, output_dim1): super().__init__() self.feature_extractor nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Linear(hidden_dim, hidden_dim) ) self.attention NadarayaWatsonAttention() self.predictor nn.Linear(hidden_dim, output_dim) def forward(self, x, memory_keys, memory_values): # 提取特征 features self.feature_extractor(x.unsqueeze(-1)).squeeze(-1) mem_features self.feature_extractor(memory_keys.unsqueeze(-1)).squeeze(-1) # 注意力聚合 context, _ self.attention(features, mem_features, memory_values) # 预测 output self.predictor(context) return output5.2 可视化注意力模式理解模型如何分配注意力至关重要def plot_attention(attention_weights, x_train, x_test): plt.figure(figsize(10, 8)) plt.imshow(attention_weights.detach().numpy(), cmapReds, aspectauto, extent[x_train.min(), x_train.max(), x_test.max(), x_test.min()]) plt.colorbar(labelAttention weight) plt.scatter(x_train, [x_test.max()]*len(x_train), cblue, alpha0.3, labelTraining points) plt.xlabel(Training inputs (keys)) plt.ylabel(Test inputs (queries)) plt.title(Attention Weights Visualization) plt.legend() plt.show() # 使用训练好的模型计算测试集的注意力 with torch.no_grad(): _, test_attn model(x_test.unsqueeze(1), x_train.unsqueeze(1), y_train.unsqueeze(1)) plot_attention(test_attn, x_train, x_test)这种可视化能清晰展示模型如何根据输入相似度分配注意力权重对于调试和解释模型行为非常有价值。