下面给出缩放点积自注意力Scaled Dot‑Product Self‑Attention极简 PyTorch 实现全部代码可直接运行注释详细不含多头先实现基础版之后附带多头注意力。import torch import torch.nn as nn import torch.nn.functional as F class ScaledDotProductSelfAttention(nn.Module): def __init__(self, embed_dim): super().__init__() self.d_k embed_dim # Q,K向量维度 # 可训练权重矩阵 Wq, Wk, Wv self.W_q nn.Linear(embed_dim, embed_dim) self.W_k nn.Linear(embed_dim, embed_dim) self.W_v nn.Linear(embed_dim, embed_dim) def forward(self, x, maskNone): :param x: input, shape [batch_size, seq_len, embed_dim] :param mask: 掩码矩阵 [batch_size, seq_len, seq_len] GPT解码器用上三角maskBERT默认None :return out: [batch_size, seq_len, embed_dim] attn_weight: 注意力权重 [batch, seq_len, seq_len] batch, seq_len, _ x.shape # Step1生成 Q, K, V Q self.W_q(x) # (B, L, D) K self.W_k(x) V self.W_v(x) # Step2: Q K^T / sqrt(d_k) attn_score torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtypetorch.float32)) # 如果传入mask将mask为0的位置设置为‑infsoftmax之后权重趋近0 if mask is not None: attn_score attn_score.masked_fill(mask 0, -1e9) # Step3 softmax得到注意力权重 attn_weight F.softmax(attn_score, dim-1) # Step4 加权value得到输出 out torch.matmul(attn_weight, V) return out, attn_weight if __name__ __main__: # 超参设置 batch_size 2 seq_len 4 # 句子4个token embed_dim 64 # 词嵌入维度 # 随机模拟输入向量等价4个单词的embedding x torch.randn(batch_size, seq_len, embed_dim) # 实例化自注意力模块 self_attn ScaledDotProductSelfAttention(embed_dim) output, attn_weights self_attn(x) print(输出shape, output.shape) # torch.Size([2,4,64]) print(注意力权重shape, attn_weights.shape) # torch.Size([2,4,4])2. 拓展多头注意力 Multi‑Head‑AttentionPyTorch 实现class MultiHeadSelfAttention(nn.Module): def __init__(self, embed_dim, num_heads): super().__init__() assert embed_dim % num_heads 0, 嵌入维度必须可以被头数整除 self.d_model embed_dim self.num_heads num_heads self.d_k embed_dim // num_heads self.Wq nn.Linear(embed_dim, embed_dim) self.Wk nn.Linear(embed_dim, embed_dim) self.Wv nn.Linear(embed_dim, embed_dim) self.W_o nn.Linear(embed_dim, embed_dim) def forward(self, x, maskNone): B, L, _ x.shape Q self.Wq(x) K self.Wk(x) V self.Wv(x) # 将维度拆分多头: (B, heads, seq_len, d_k) Q Q.view(B, L, self.num_heads, self.d_k).permute(0,2,1,3) K K.view(B, L, self.num_heads, self.d_k).permute(0,2,1,3) V V.view(B, L, self.num_heads, self.d_k).permute(0,2,1,3) # 缩放点积注意力 attn_score torch.matmul(Q, K.transpose(-2,-1)) / torch.sqrt(torch.tensor(self.d_k, dtypetorch.float32)) if mask is not None: attn_score attn_score.masked_fill(mask 0, -1e9) attn_weight F.softmax(attn_score, dim-1) out torch.matmul(attn_weight, V) # 拼接多头结果 out out.permute(0,2,1,3).contiguous().view(B, L, self.d_model) out self.W_o(out) return out, attn_weight if __name__ __main__: batch 2 seq_len 4 embed_dim 64 heads 8 x torch.randn(batch, seq_len, embed_dim) mh_attn MultiHeadSelfAttention(embed_dim, heads) out, attn mh_attn(x) print(多头注意力输出shape, out.shape)3. 关键点总结Self‑Attention 条件Q、K、V 全部来自输入xCross‑Attention 则 Q 来自 decoder、K‑V 来自 encoderO(n2)复杂度序列长度变长比如几千 token计算量急剧变大该代码没有加入位置编码如果你要做 GPT/BERT 完整版本需要额外加入Positional Encoding。如果你需要我可以给你写带位置编码 残差 LayerNorm完整 Transformer Encoder 单层代码。Q 拿着问题挨个匹配 K匹配分数决定取用 V 的多少通俗拆解Q‑K‑V 逻辑结合数学过程讲明白1. 角色类比Q(Query)提问者每一个 token 带着自己的特征向量提出问题我应该重点关注序列里哪些单词K(Key)候选人的自我介绍序列中每个 token 拿出自身特征用来接受匹配V(Value)候选人真正携带的内容信息。K 只用来打分最终提取内容取自 V。举句子例子He ate appletokenate生成自己的 Q 向量剩下三个 tokenHe、ate、apple各自拿出 KQ 和每个 K 做内积计算相似度Q (ate) 和 K (apple) 内积数值很高Q (ate) 和 K (He) 内积偏低 得到原始分数[scoreHe​, scoreate​, scoreapple​]。2. 缩放 Softmax把分数变成权重除以dk​​缩放防止数值爆炸softmax 归一化所有权重取值 0‑1总和为 1weightssoftmax(dk​​QK⊤​) 这时 w1​(He)≈0.1w2​(ate)≈0.2w3​(apple)≈0.7。 含义ate这个单词 70% 的信息来自 apple10% 来自 He。3. 加权 V 获取输出outputate​w1​⋅VHe​w2​⋅Vate​w3​⋅Vapple​权重越高对应 V 的信息放进来越多权重很低的 V 几乎被忽略。一句话概括你的原话Q 拿着自身特征挨个比对全部 K算出相关性得分softmax 将得分转为权重用权重分配 V 里信息占比这就是自注意力的核心。4. 区分 Self‑Attention 和 Cross‑AttentionSelf‑Attention自注意力Q,K,V全部来自同一组输入句子。句子内部词语互相匹配。Cross‑Attention交叉注意力翻译场景Q 来自 Decoder目标语言K、V 来自 Encoder源语言。例如英译中中文 token 的 Q 去匹配英文单词的 K 和 V。5. 代码层面精简对应关系结合刚才代码attn_score Q K.transpose(-2,-1) / sqrt(d_k) # Q挨个匹配K attn_weight F.softmax(attn_score, dim-1) # 算出权重 out attn_weight V # 根据权重取用V6. 误区纠正Q、K、V 不是原始词向量是输入向量分别乘以可训练矩阵WQ​,WK​,WV​得到模型训练期间自动学习什么样的 Q/K/V 更容易捕捉语义关联内积为什么可以衡量相似度向量夹角越小内积越大特征方向越接近代表语义关联性越强。