目录一、什么是自注意力机制二、自注意力的计算步骤以单个头为例步骤 1生成 Q、K、V 向量步骤 2计算注意力分数Attention Scores步骤 3应用 Softmax 得到注意力权重步骤 4加权求和 Value代码实现PyTorch三、为什么自注意力优于 RNN四、多头注意力Multi-Head Attention五、现代优化FlashAttention 与 SDPA六、自注意力的应用场景七、总结一、什么是自注意力机制自注意力机制Self-Attention 是 Transformer 模型的核心组件它允许模型在处理序列数据如文本、语音时动态地关注序列中不同位置之间的相关性。核心思想“当前词应该关注序列中的哪些词”比如在句子 “The cat sat on the mat” 中“mat” 和 “cat” 语义相关自注意力会让 “mat” 更加关注 “cat”。二、自注意力的计算步骤以单个头为例假设输入是一个序列的向量表示,其中 n 是序列长度d 是向量维度。步骤 1生成 Q、K、V 向量通过三个可学习的权重矩阵将输入 X 映射为Query查询Key键KXWValue值VXW 类比搜索引擎Query你在搜什么Key网页的关键词Value网页的实际内容步骤 2计算注意力分数Attention Scores计算每个 Query 与所有 Key 的点积衡量“匹配度”为了防止点积过大导致 Softmax 梯度饱和进行缩放Scaled其中是Key向量的维度步骤 3应用 Softmax 得到注意力权重输出是一个概率分布表示每个位置对其他位置的关注程度。步骤 4加权求和 Value用注意力权重对 Value 向量加权求和得到最终输出Output Weights.V这个输出向量融合了整个序列的信息并根据语义重要性进行了加权。代码实现PyTorchimport torch import math def self_attention(Q, K, V, maskNone): # 获取 Q 矩阵最后一个维度的大小即 d_k键向量的维度 d_k Q.size(-1) # 计算注意力分数Q K^T / sqrt(d_k) # 目的是防止点积过大导致 softmax 梯度消失 scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k) # 如果有掩码如 padding mask 或 causal mask将无效位置设为极小值-inf # softmax 后这些位置的权重会趋近于 0 if mask is not None: scores scores.masked_fill(mask 0, -1e9) # 对注意力分数进行 softmax得到注意力权重权重和为1 weights torch.softmax(scores, dim-1) # 加权求和用注意力权重乘以值向量 V得到输出 return torch.matmul(weights, V)三、为什么自注意力优于 RNN指标RNNself-attention并行性串行处理t1→ t2 → t3)所有位置同时计算训练更快长度依赖梯度消失/爆炸难以捕捉远距离消息直接建模任意两个位置的关系计算复杂度O(n) per step$$O(n^2)$$但 GPU 并行弥补灵活性固定顺序更新动态调整关注点 举例在句子 “I arrived at the bank after crossing the river” 中RNN 需要一步步传递信息到 “bank”可能丢失上下文Self-Attention 可以直接让 “bank” 关注 “river”立刻判断是“河岸”而非“银行”。四、多头注意力Multi-Head Attention为了捕捉不同子空间的语义信息Transformer 使用 多头注意力Multi-Head Attention将Q、K、V分成多个“头”如8个每个头独立计算自注意力最后将所有头的输出拼接并线性变换其中✅ 好处不同头可以关注不同类型的依赖如语法、语义、指代等五、现代优化FlashAttention 与 SDPA随着模型变大标准自注意力计算成本高。Hugging Face 和 PyTorch 提供了优化方案1.SDPAScaled Dot Product AttentionPytorch 内置的高效实现自动选择最优后端from torch.nn.functional import scaled_dot_product_attention output scaled_dot_product_attention(Q, K, V, attn_maskmask)2.FlashAttention-2更快减少显存读写支持 fp16/bf16需要安装pip install flash-attn --no-build-isolation启用方式model AutoModelForCausalLM.from_pretrained( meta-llama/Llama-3.1-8B, attn_implementationflash_attention_2, torch_dtypetorch.bfloat16, device_mapauto )⚠️ 注意不支持 padding token建议在训练时使用 packed dataset。六、自注意力的应用场景场景说明自然语言处理BERT、GPT、LLaMA 等大模型核心图像处理Vision TransformerViT中将图像分块视为序列语音识别处理音频帧序列推荐系统用户行为序列建模七、总结特性说明核心思想动态关注序列中重要位置三大向量Q查什么、K有什么、V给什么缩放点积防止梯度饱和并行优势比 RNN 更快、更适合 GPU长度依赖能直接建模任意距离的依赖关系现代优化SDPA、FlashAttention-2 提升推理效率 一句话总结自注意力机制让模型像“阅读时划重点”一样动态决定哪些词更重要从而更高效地理解和生成语言。