注意力机制的核心是动态加权求和,三要素为Query(查询)、Key(键)、Value(值):Query与Key计算相似度得分数,softmax归一化为权重,再加权求和Value得到输出;自注意力通过全连接匹配突破距离限制,多头机制并行捕获多维特征;PyTorch手写实现含线性投影、缩放点积、softmax及加权求和;可视化注意力权重热力图可分析模型关注模式。
注意力机制不是让模型“看哪里”,而是让模型在处理信息时,动态分配不同权重——哪些词、哪些位置、哪些特征更关键,就给它更高权重。核心在于“加权求和”+“可学习的权重计算”。
这是理解所有注意力变体的起点:
简单说:Query 和 Key 算相似度 → 得到注意力分数(attention scores)→ 归一化成权重(softmax)→ 加权求和 Value → 输出上下文向量。这个过程可微、可训练。
传统RNN/LSTM靠逐步传递隐藏状态,距离越远信号越弱;CNN靠局部感受野,需多层堆叠才能覆盖全局。而自注意力一步到位:
立即学习“Python免费学习笔记(深入)”;
这也是 Transformer 能替代 RNN 成为大模型基座的关键原因。
不依赖 nn.MultiheadAttention,从零实现有助于理解每一步的张量形状和计算逻辑:
import torch import torch.nn as nn import torch.nn.functional as Fclass SimpleAttention(nn.Module): def init(self, dim): super().init() self.W_q = nn.Linear(dim, dim) self.W_k = nn.Linear(dim, dim) self.W_v = nn.Linear(dim, dim)
def forward(self, x): # x: [batch, seq_len, dim] Q = self.W_q(x) # [b, s, d] K = self.W_k(x) # [b, s, d] V = self.W_v(x) # [b, s, d] # 计算相似度(缩放点积) scores = torch.bmm(Q, K.transpose(1, 2)) / (K.size(-1) ** 0.5) # [b, s, s] attn_weights = F.softmax(scores, dim=-1) # 每行和为1 output = torch.bmm(attn_weights, V) # [b, s, d] return output, attn_weights
注意:实际训练中通常会加入 dropout、LayerNorm 和残差连接,构成一个完整的 Attention Block。
训练完模型后,提取某次前向传播中的 attention_weights(形状为 [batch, heads, seq_len, seq_len]),可做如下分析:
指向其先行词;matplotlib.pyplot.imshow 或 seaborn.heatmap 快速出图,配合 xticks/yticks 标注词元。这不仅是调试工具,更是验证模型是否学到合理语言结构的直接证据。