当前位置：首页 > news >正文

Self Attention

news 2026/3/27 0:59:25

一、什么是自注意力机制？

二、自注意力的计算步骤（以单个头为例）

步骤 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（键）：K=XW

Value（值）：V=XW

💡 类比搜索引擎：

Query：你在搜什么？

Key：网页的关键词

Value：网页的实际内容

步骤 2：计算注意力分数（Attention Scores）

计算每个 Query 与所有 Key 的点积，衡量“匹配度”：

为了防止点积过大导致 Softmax 梯度饱和，进行缩放（Scaled）：

其中，是Key向量的维度

步骤 3：应用 Softmax 得到注意力权重

输出是一个概率分布，表示每个位置对其他位置的关注程度。

步骤 4：加权求和 Value

用注意力权重对 Value 向量加权求和，得到最终输出：

Output = Weights.V

这个输出向量融合了整个序列的信息，并根据语义重要性进行了加权。

代码实现（PyTorch）

import torch import math def self_attention(Q, K, V, mask=None): # 获取 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？

指标	RNN	self-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.SDPA（Scaled Dot Product Attention）

Pytorch 内置的高效实现，自动选择最优后端

from torch.nn.functional import scaled_dot_product_attention output = scaled_dot_product_attention(Q, K, V, attn_mask=mask)

2.FlashAttention-2

更快（减少显存读写）
支持 fp16/bf16
需要安装：pip install flash-attn --no-build-isolation

启用方式：

model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-3.1-8B", attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16, device_map="auto" )

⚠️ 注意：不支持 padding token，建议在训练时使用 packed dataset。

六、自注意力的应用场景

场景	说明
自然语言处理	BERT、GPT、LLaMA 等大模型核心
图像处理	Vision Transformer（ViT）中将图像分块视为序列
语音识别	处理音频帧序列
推荐系统	用户行为序列建模

七、总结

特性	说明
核心思想	动态关注序列中重要位置
三大向量	Q（查什么）、K（有什么）、V（给什么）
缩放点积	防止梯度饱和
并行优势	比 RNN 更快、更适合 GPU
长度依赖	能直接建模任意距离的依赖关系
现代优化	SDPA、FlashAttention-2 提升推理效率