Transformer中的自注意力机制是其核心创新，它完全摒弃了RNN的时序递归与CNN的局部感受野，转而通过**并行化、全局感知、可微分的矩阵运算**建模任意两个token之间的语义依赖关系。其实现过程严格遵循“查询-键-值”（Q/K/V）三元范式，并在数学上可精确拆解为**四步标准矩阵流水线**：投影 → 相似度计算 → 加权归一化 → 聚合输出。以下从原理、公式、代码到工程细节进行系统性展开。 --- ### 一、自注意力的完整矩阵运算流程（四步精解） 设输入序列长度为 $ n $，嵌入维度为 $ d_{\text{model}} $，输入张量为 $ X \in \mathbb{R}^{n \times d_{\text{model}}} $（已叠加词嵌入与位置编码）。自注意力执行如下步骤： | 步骤 | 数学表达 | 物理含义 | 关键特性 | |------|----------|-----------|------------| | **1. 线性投影** | $ Q = XW^Q,\ K = XW^K,\ V = XW^V $<br>其中 $ W^Q, W^K \in \mathbb{R}^{d_{\text{model}} \times d_k},\ W^V \in \mathbb{R}^{d_{\text{model}} \times d_v} $ | 将同一输入映射为三组语义角色向量：<br>• $ Q $：当前token“想问什么”（查询）<br>• $ K $：其他token“能回答什么”（键）<br>• $ V $：其他token“实际携带的信息”（值） | 投影矩阵 $ W^Q/W^K/W^V $ 通过反向传播联合学习，使语义相近token的 $ Q $ 与 $ K $ 向量夹角更小，点积更大 [ref_2]；训练完成后即固化用于推理 [ref_4]。 | | **2. 相似度计算（logits）** | $ S = QK^\top \in \mathbb{R}^{n \times n} $<br>$ S_{ij} = \langle q_i, k_j \rangle $ | 计算所有token对之间的“匹配强度”，构成 $ n \times n $ 的相似度矩阵。$ S_{ij} $ 表示第 $ i $ 个token对第 $ j $ 个token的关注程度 | 若无缩放，当 $ d_k $ 较大时，$ \langle q_i, k_j \rangle $ 方差趋近 $ d_k $，导致Softmax饱和、梯度消失；故必须引入缩放因子 $ \frac{1}{\sqrt{d_k}} $ [ref_1]。 | | **3. 掩码 + Softmax归一化** | $ A = \mathrm{softmax}\left( \frac{QK^\top}{\sqrt{d_k}} + M \right) $<br>其中 $ M_{ij} = \begin{cases} 0 & \text{合法位置} \\ -\infty & \text{需屏蔽位置} \end{cases} $ | 对每行 $ i $ 进行Softmax，将相似度转化为概率分布 $ A_{i:} $，确保 $ \sum_j A_{ij} = 1 $；掩码 $ M $ 实现两种关键控制：<br>• **Padding Mask**：屏蔽填充符（如 `[PAD]`）<br>• **Causal Mask**：解码器中强制上三角为 $ -\infty $，防止信息泄露 [ref_5] | 工程实现中，$ -\infty $ 常用极小负数（如 `-1e9`）替代，避免数值不稳定 [ref_1]。 | | **4. 加权聚合输出** | $ O = AV \in \mathbb{R}^{n \times d_v} $<br>$ o_i = \sum_{j=1}^n A_{ij} v_j $ | 按注意力权重 $ A_{ij} $ 对所有 $ v_j $ 线性加权求和，生成第 $ i $ 个token的新表示 $ o_i $，融合了全局上下文信息 | 输出维度 $ d_v $ 可独立于 $ d_k $，为多头注意力提供灵活性 [ref_1]。 | > ✅ **关键结论**：整个过程本质是**三次矩阵乘法 + 一次Softmax**，全部可GPU并行，无任何循环或条件分支，彻底规避RNN的串行瓶颈 [ref_1]。 --- ### 二、PyTorch代码实现（含逐行溯源注释） ```python import torch import torch.nn as nn import torch.nn.functional as F class ScaledDotProductAttention(nn.Module): def __init__(self, d_k: int): super().__init__() self.d_k = d_k def forward(self, Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, mask: torch.Tensor = None): """ Q: (batch, n, d_k) # ref_6 中明确指出Q/K/V形状需一致以支持广播 K: (batch, n, d_k) V: (batch, n, d_v) mask: (batch, 1, n) or (batch, n, n), 值为0/1或0/-inf """ # Step 1: Compute logits: Q @ K^T -> (batch, n, n) # ref_1 强调该步复杂度主导整体O(n^2 d_k)，是长序列瓶颈根源 scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32)) # Step 2: Apply mask (e.g., causal mask for decoder) # ref_1 和 ref_5 均指出mask需在softmax前加至scores，实现padding/因果约束 if mask is not None: scores = scores.masked_fill(mask == 0, -1e9) # 使用-1e9模拟-inf # Step 3: Softmax over last dim -> attention weights # ref_2 解释Softmax使权重和为1，形成概率分布，实现可微分注意力 attn_weights = F.softmax(scores, dim=-1) # (batch, n, n) # Step 4: Weighted sum of V # ref_4 明确o_i = Σ A_ij * v_j，此即上下文感知的token重表征 output = torch.matmul(attn_weights, V) # (batch, n, d_v) return output, attn_weights # Example usage with dummy data batch, n, d_model, d_k, d_v = 2, 5, 8, 4, 6 X = torch.randn(batch, n, d_model) W_Q = nn.Linear(d_model, d_k, bias=False) W_K = nn.Linear(d_model, d_k, bias=False) W_V = nn.Linear(d_model, d_v, bias=False) Q, K, V = W_Q(X), W_K(X), W_V(X) # ref_6 强调投影是独立线性层 attn = ScaledDotProductAttention(d_k) output, weights = attn(Q, K, V) print(f"Output shape: {output.shape}") # torch.Size([2, 5, 6]) ``` --- ### 三、为什么必须是“点积”？——语义相似性的几何本质 自注意力并非凭空设计，其点积操作直指**向量空间中的余弦相似度**： $$ \langle q_i, k_j \rangle = \|q_i\| \|k_j\| \cos\theta_{ij} $$ - 在训练过程中，模型通过梯度下降不断调整 $ W^Q $ 和 $ W^K $，使得**语义相关token的 $ q_i $ 与 $ k_j $ 夹角 $ \theta_{ij} $ 趋近0**，从而点积值显著增大； - 反之，无关token夹角接近90°，点积趋近0； - 因此，$ A_{ij} $ 实质是“第 $ i $ 个词在多大程度上认为第 $ j $ 个词对其当前语义有贡献”，这是一种**完全数据驱动、无需人工规则的软对齐机制** [ref_2]。 例如在句子 `"The cat sat on the mat"` 中： - 当处理 `"sat"` 时，其 $ Q $ 向量会与 `"cat"`（主语）、`"mat"`（宾语）的 $ K $ 向量产生高点积，从而在 $ O $ 中强化这两个实体的语义； - 而 `"the"`（冠词）因语义泛化，点积普遍偏低，被自动抑制。 --- ### 四、自注意力 vs RNN/CNN：性能对比表 | 维度 | 自注意力（Transformer） | RNN（LSTM/GRU） | CNN（扩张卷积） | |------|--------------------------|------------------|-------------------| | **最大路径长度** | $ O(1) $：任意两token间仅1层即可交互 [ref_1] | $ O(n) $：需 $ n $ 步传递信息 | $ O(\log_k n) $：依赖扩张率 $ k $ [ref_1] | | **最小串行操作数** | $ O(1) $：全矩阵并行 [ref_1] | $ O(n) $：严格时间步依赖 | $ O(1) $：卷积核内并行 | | **长程依赖建模能力** | ⭐⭐⭐⭐⭐：全局可见，但受 $ n^2 $ 显存限制 | ⚠️：梯度消失/爆炸严重 | ⚠️：需深层堆叠，感受野增长缓慢 | | **可并行性** | ⭐⭐⭐⭐⭐：无隐状态依赖，完美适配GPU | ❌：无法跨时间步并行 | ⭐⭐⭐⭐：卷积核内并行，但层间仍串行 | | **典型应用场景** | 机器翻译、文本摘要、大语言模型（GPT/BERT） [ref_3] | 语音识别、小规模序列标注 | 图像分类、短文本分类 | > 💡 正因上述优势，Transformer成为现代大模型（LLM）的**唯一基础架构**；而其二次复杂度缺陷，则催生了FlashAttention、RingAttention、StreamingLLM等工业级优化方案 [ref_1]。 --- ### 五、常见误区澄清（基于参考资料交叉验证） | 误区 | 正解 | 依据 | |------|------|------| | “Q/K/V是三个独立网络” | ❌ 实际是**同一输入经三组不同线性权重投影**得到，参数量占比约 $ 4d_{\text{model}}^2 $，是FFN的次要参数来源 [ref_1] | [ref_1] 第二节“复杂度估算” | | “Softmax是必须的非线性” | ❌ Softmax在此仅作**归一化**，不引入非线性表达力；真正的非线性来自FFN层的ReLU [ref_1] | [ref_1] “前馈网络”小节 | | “位置编码必须用正弦” | ❌ 正弦PE是原始论文选择，但**可学习位置嵌入（Learned PE）在实践中更常用**（如BERT），且性能相当甚至更优 [ref_1] | [ref_1] “位置建模”首段 | | “多头注意力提升表达能力” | ✅ 但非因‘多个头’本身，而是因**各头学习不同子空间的注意力模式**（如语法头、指代头、关系头），最终拼接增强鲁棒性 [ref_3] | [ref_3] “多头注意力”分析 | 综上，自注意力机制绝非黑箱魔法，而是一套**严格可推导、可实现、可解释的线性代数系统**。掌握其矩阵本质，是理解BERT、GPT、Claude等一切现代大模型的基石。