自注意力机制深度解析：从原理到PyTorch实现

很多初学 Transformer 的人，第一眼看到那幅经典的架构图时，通常会有两种感受：要么觉得它过于复杂，注意力机制、多头、位置编码、残差连接、层归一化一大堆概念堆在一起，不知道从哪入手；要么觉得它不过如此，核心不就是算一个加权平均吗，有什么值得大书特书的？

这两种感受都有道理，但也都错过了最关键的东西。

Transformer 真正革命性的地方，不是那个公式有多复杂，而是它彻底改变了深度学习中处理序列数据的方式。在它之前，RNN、LSTM 是处理序列事实上的标准选择，它们像人类阅读一样逐字处理输入，顺序执行、难以并行，长距离依赖问题虽然在 LSTM 中有所缓解，但没有被真正解决。Transformer 则完全抛弃了循环结构，用一种叫自注意力（Self-Attention）的机制，让序列中的任意两个位置直接建立联系。

这一改变带来的连锁反应是巨大的：并行计算成为可能，训练效率大幅提升；长距离依赖不再是瓶颈，BERT、GPT 这类超大规模预训练模型才有了存在的根基。可以说，自注意力机制是理解整个 Transformer 大厦的第一块地基，也是从传统深度学习走向现代大模型必须跨越的一道门槛。

这篇文章会从最朴素的问题出发，用场景化的方式讲清楚自注意力机制到底解决了什么问题、它的计算过程每一步在做什么、为什么需要缩放点积、多头注意力又为什么要存在，并配合完整的 PyTorch 代码实现，帮你把理论和代码对上号。读完这篇文章，你再看 Transformer 的任何变体、再看大模型的原始论文，都会感觉顺很多。

1. 这篇文章真正要解决的问题

先做一个判断：自注意力机制是 Transformer 的核心创新，也是理解全部现代大模型的最低门槛。

很多人学 Transformer 时遇到的困境，不是不想学，而是被资料淹没。打开搜索引擎，教程铺天盖地，但大体上有两类：

一类是纯概念科普，把注意力机制讲成了一个神奇的黑盒，放几张热力图，告诉你模型在关注哪些词，看起来很直观，但你合上页面，依然不知道 Q（Query）、K（Key）、V（Value）到底怎么来的，更不知道代码从哪一行开始写。

另一类是纯数学推导，把注意力公式用大篇幅展开，算给你看内积、归一化、加权求和，逻辑没问题，但你不知道这些矩阵运算放在神经网络里到底承担了什么角色，为什么要除一个根号 dk，这个操作不合理会怎样。

更常见的一个痛点是：很多人习惯了 CNN 和 RNN 的思维，会把自注意力理解为某种“更复杂的卷积”或“更聪明的 LSTM”，这种先入为主的类比框架，反而成了理解 Transformer 的最大障碍。

这篇文章想解决的问题，就是把心法和招式结合起来：

• 用场景讲直觉：为什么需要自注意力？它比 RNN 强在哪？
• 用公式讲原理：Q、K、V 的本质是什么？缩放点积为什么要有缩放？
• 用代码讲实现：从零实现一个完整的自注意力模块，让你真的跑起来。
• 用对比讲延展：多头注意力、因果自注意力是怎么从基础版本演化出来的？

如果你正在学深度学习，或者刚接触大模型相关的技术，甚至是你已经在用 BERT、GPT 的 API 但想深入了解它们内部机制，这篇文章都适合你。它不要求你有深厚的数学背景，只要会基本的矩阵乘法和 python 编程，就能跟上节奏。

2. 从 Seq2Seq 到 Self-Attention：注意力机制到底在解决什么问题？

要理解自注意力，必须先从它的前身——经典的注意力机制（Attention）讲起。

2.1 没有注意力机制的时代：信息瓶颈问题

在注意力机制出现之前，序列到序列（Seq2Seq）任务的标准架构是这样的：

• 一个编码器（Encoder）RNN/LSTM 把输入序列逐字读进去，最后汇总成一个固定长度的向量，这个向量被称作“上下文向量”。
• 一个解码器（Decoder）RNN/LSTM 从这个上下文向量出发，一步步生成输出序列。

这个架构存在一个致命问题：信息瓶颈。无论输入序列多长，编码器都必须把所有信息压进最后一个隐藏状态里。如果句子很长，开头的信息在传递过程中会被不断稀释，最终到达解码器时，可能已经损失殆尽。这就是长距离依赖问题的本质。

2.2 注意力机制的诞生：按需获取信息

注意力机制的突破性想法是：解码器在生成每一步的时候，不再只依赖一个固定的上下文向量，而是可以用一个“查询”去编码器的所有位置上检索信息，然后根据相关性决定要看哪里、看多少。

这里借用一下信息检索的类比，非常直观：

• Query（查询）：你可以理解为你大脑中正在思考的问题。比如你要翻译一个词，当前的生成状态就是你的 Query，相当于你在问“现在该看源句子的哪个词？”
• Key（键）：源序列中每个词的一个索引标签。你会拿 Query 和每个 Key 做匹配，判断“这个词跟我当前的问题有关系吗？”
• Value（值）：源序列中每个词的实际内容。一旦根据匹配程度确定了该关注哪些词，就去取这些词的内容来用。

整个过程就是三步：把 Query 和每个 Key 做相似度计算，通过 Softmax 转成权重，再用权重对 Value 加权求和。

注意力得分 = Softmax(Query 与每个 Key 的相似度) × Value

2.3 什么是自注意力：关注的不是“对方”，而是“自己”

经典注意力发生在编码器和解码器之间，建立在两个不同序列上。而自注意力（Self-Attention），则是一个序列内部自己和自己计算注意力，所以它还有另一个名字：内部注意力（Intra-Attention）。

自注意力的核心思想是：对于一个句子中的每个词，通过计算它和其他所有词的相关性，来更新它自己的语义表示。

举个例子，处理这句话：

The animal didn't cross the street because it was too tired.

这里的 “it” 指的是什么？是 animal 还是 street？对人类来说，我们根据常识知道是 animal，但机器一开始不知道。自注意力机制在编码 “it” 这个词时，会计算它与 “animal” 的注意力分数很高，于是把 “animal” 的语义信息融合进 “it” 的表示里。这样一来，编码器输出中的 “it” 就携带了 “animal” 的信息，语义就准确了。

这个能力在 RNN 里要通过很多步的循环传递才能实现，而且容易衰减，在自注意力里是一步到位的全局交互。

所以自注意力最本质的优势可以总结成一句话：它让序列中的每个词，在编码自己的时候，能直接看到整个序列的所有其他词。

用一张简单的流程对比来看：

对比维度	RNN/LSTM	自注意力
信息传递方式	逐步顺序传递	任意两点直接交互
长距离依赖	容易衰减	一步到位
并行性	无法并行	完全可并行
计算复杂度	O(n) 步循环	O(n²) 两两计算
位置信息	天然具备顺序感	需要额外添加位置编码

这里出现了一个很有意思的权衡：RNN 的优势是复杂度低（线性），自注意力的代价是复杂度高（平方级）。后面关于长文本优化、稀疏注意力等一系列研究，本质都是在想尽办法把 O(n²) 降下来。

3. 自注意力机制的核心原理与数学表达

现在从代码实现的角度，把自注意力的计算过程一步步拆开。这是整篇文章最关键的部分，我会尽量让公式和直觉对齐。

3.1 输入表示：从词向量开始

假设我们有一个长度为 n 的输入序列，每个词用一个维度为 d_model 的向量表示。那么在自注意力模块中，输入可以看作一个矩阵 X，形状为 (n, d_model)。

在实际的实现中，我们通常还带一个 batch 维，也就是输入形状为 (batch_size, seq_len, d_model)。为了简化，我们先忽略 batch，只考虑单条序列。

3.2 生成 Q、K、V：线性变换不是可有可无的

注意力机制的输入不是直接用 X 本身，而是把 X 分别通过三个可学习的权重矩阵 W_Q、W_K、W_V 做线性变换，得到 Q、K、V：

Q = X @ W_Q 形状: (n, d_k) K = X @ W_K 形状: (n, d_k) V = X @ W_V 形状: (n, d_v)

为什么不能直接用 X 作为 Q、K、V？这里是一个新手最容易疑惑的点。

从直觉上讲，线性变换的作用是让模型从输入中提取不同的特征视角。X 本身是词向量，它是静态的；而通过不同的线性变换，模型可以在不同的语义空间中衡量“提问”和“回答”的关系。没有这个变换，所有词都在同一个空间里直接做内积，表达能力会受到很大限制。

从工程上讲，如果不经过变换，直接让 Q=X、K=X，那么计算出的注意力矩阵就是固定的。比如某个位置和其他位置的相似度完全由词向量的静态相似度决定，相当于模型没有学习余地。引入三个可学习的权重矩阵后，这个相似度计算方式本身可以通过反向传播被优化，模型的能力上限就完全不同了。

3.3 缩放点积注意力：从相似度到权重

自注意力的第一步，是计算 Query 和 Key 的相似度。最常用的方式是点积：

S = Q @ K^T 形状: (n, n)

S 是一个 n×n 的矩阵，S[i][j] 表示序列中第 i 个词（作为 Query）和第 j 个词（作为 Key）的相似度分数。这个分数越高，说明在编码第 i 个词时，第 j 个词越值得关注。

但是直接使用点积会有一个数值稳定性的问题。当维度 d_k 比较大时，点积的结果会很大，导致进入 Softmax 之后梯度极小，出现梯度消失的问题。

一个标准的解决办法就是缩放：把点积结果除以根号 d_k。

S = (Q @ K^T) / sqrt(d_k)

这里的原理在数学上有个解释：如果 Q 和 K 中的元素都是均值为 0、方差为 1 的随机变量，那么它们的点积结果的均值为 0、方差为 d_k。为了让结果落在 Softmax 函数的敏感区间，我们希望方差保持在 1 的数量级，所以需要除以根号 d_k。这是一个很精巧的数值稳定设计。

接下来用 Softmax 把分数转换成概率分布，保证每一行的权重和为 1：

A = softmax(S, dim=-1) 形状: (n, n)

最后用这个权重矩阵对 Value 做加权求和：

Output = A @ V 形状: (n, d_v)

综合起来，缩放点积注意力的完整公式就是：

Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V

这就是 Transformer 论文《Attention Is All You Need》中的核心公式，没有其他复杂的东西。

3.4 一个具体的小例子

假设序列只有三个词：“我”、“爱”、“猫”。经过词嵌入和线性变换后，我们得到了一个简单的 Q、K、V（为了说明方便，取 d_k=2，数值是编造的）：

Q: [[0.8, 0.2], K: [[0.9, 0.1], V: [[0.5, 0.1], [0.5, 0.6], [0.4, 0.7], [0.4, 0.8], [0.1, 0.9]] [0.2, 0.8]] [0.9, 0.3]]

第一步计算 Q @ K^T，得到 3×3 的相似度矩阵：

S: [[0.74, 0.46, 0.32], [0.81, 0.62, 0.58], [0.57, 0.68, 0.74]]

第二步除以 sqrt(2) ≈ 1.414：

S_scaled: [[0.523, 0.325, 0.226], [0.573, 0.438, 0.410], [0.403, 0.481, 0.523]]

第三步对每一行做 Softmax。以第一行为例：

• 计算 exp：exp(0.523)=1.687, exp(0.325)=1.384, exp(0.226)=1.254
• 求和：1.687 + 1.384 + 1.254 = 4.325
• 权重：0.390, 0.320, 0.290

第四步用这些权重对 V 的三行加权求和，得到输出：

Output[0] = 0.390 × [0.5, 0.1] + 0.320 × [0.4, 0.8] + 0.290 × [0.9, 0.3] = [0.584, 0.392]

这就是一次完整自注意力计算的过程。你不需要手算，但理解每一步的输入输出形状和含义，后面看代码就会很轻松。

4. 为什么自注意力有效：三个关键性质

很多人看完公式会说：这不就是一个加权平均吗？那我直接用词向量线性组合不行吗？为什么非得搞得这么复杂？

这个质疑很合理。自注意力相比简单的加权平均，有以下三个关键性质值得注意。

4.1 动态权重

加权平均的权重往往是固定的，或者说由输入的某个线性映射决定，但自注意力的权重是由 Query 和 Key 的内积动态计算出来的。这意味着模型在推理时，会根据当前上下文动态决定关注哪些位置。同一个词在不同句子里，即便完全一样，它在不同上下文中关注的对象和权重也不同。

这种动态性带来的表达能力提升，是静态加权平均无法实现的。

4.2 全局感受野

CNN 的感受野受限于卷积核的大小，虽然可以通过堆叠很多层来扩大感受野，但每一层都只能看到局部区域。RNN 虽然在理论上可以访问任意距离的信息，但实际训练中梯度衰减让远程信息很难被有效利用。

自注意力一步到位，每个位置的输出都直接包含了所有位置的信息。对于序列中任意两个位置，它们之间的交互路径长度永远是 1，不管距离多远。这种特性让 Transformer 特别适合捕捉长距离依赖关系。

4.3 并行计算

RNN 的问题是必须按时间步顺序计算，第 t 步的输出依赖于第 t-1 步的隐藏状态，因此无法并行。自注意力则完全不同，它建立了完整的 n×n 关联矩阵，这个矩阵运算是矩阵乘法，可以被 GPU 高度并行化。这是 Transformer 能高效训练大规模模型的基础，也是它在工程上最具吸引力的原因之一。

5. 从自注意力到多头注意力：为什么要“多头”

理解了自注意力后，多头注意力就很好理解了。

5.1 单头注意力的局限

单个注意力机制相当于只在一个语义空间中计算词和词之间的关系。但句子中的词之间的关系往往是多层面的。

还是看这个句子：

"I love my dog because it is so cute."

这里的 “it” 指代 “dog”，这是一种语法上的指代关系。但在同一个句子里，“I” 和 “love” 之间有主谓关系，“my” 和 “dog” 之间有修饰关系，“love” 和 “dog” 之间有动宾关系，这些关系分散在不同层面的语义空间中。

单头注意力只能“兼顾”所有关系，最终学到的权重是这些关系的某种平均，可能每个层面的关系都学到了一点，但都不够精确。

5.2 多头注意力：多个语义空间并行

多头注意力的做法很简单：把 Q、K、V 投影到多个不同的低维子空间中，在每个子空间中独立地做注意力计算，然后把所有头的结果拼接起来，再经过一个线性变换。

head_i = Attention(X @ W_Q_i, X @ W_K_i, X @ W_V_i) MultiHead(X) = Concat(head_1, ..., head_h) @ W_O

这里每个头有自己的 W_Q_i、W_K_i、W_V_i 权重矩阵，相当于在不同的表示子空间中执行注意力。第一个头可能关注语法关系，第二个头可能关注指代关系，第三个头可能关注语义相关性，各司其职。

多头注意力的本质，可以类比 CNN 中的多个卷积核。CNN 中一个卷积核只能提取一种特征，所以要用多个卷积核提取不同类型的特征；多头注意力也是这样，一个头负责一个层面的关系，多个头汇总在一起，就能捕捉更丰富的交互信息。

5.3 工程参数

在标准 Transformer 中，通常设置 h=8 个头，每个头的维度 d_k = d_v = d_model / h。比如 d_model=512 时，每个头的维度是 64。这和单头注意力的计算复杂度大致相同，但表达能力更强。

6. 因果自注意力：大模型生成中的关键变体

自注意力有一个重要的变体叫因果自注意力（Causal Self-Attention），也叫掩码自注意力（Masked Self-Attention），它是所有自回归语言模型（如 GPT 系列）的基础组件。

6.1 为什么要掩码

在语言模型生成任务中，我们需要模型根据前文预测下一个词。如果模型在预测第 i 个词时能看到后面的词，那预测就失去意义了——答案都写出来了，还预测什么？

所以因果自注意力在标准自注意力的基础上，加了一个掩码操作：计算注意力分数时，禁止位置 i 关注任何 j > i 的位置。也就是只能看到当前位置及之前的词，不能看到未来的词。

具体实现中，通常是在计算 S = Q @ K^T 之后，把矩阵的上三角部分填充为一个非常小的负数（比如 -inf），这样经过 Softmax 后，这些位置的权重几乎为 0。

6.2 掩码矩阵示例

对于序列长度 4，掩码矩阵长这样（1 表示允许关注，0 表示不允许）：

[[1, 0, 0, 0], [1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 1]]

第一行只允许看到第一个词，第二行允许看到前两个词，以此类推。

这就是为什么 Transformer 解码器训练时虽然可以并行输入所有位置，但每个位置都只能看到自己前面的信息。这个细节，是理解 GPT 系列模型训练方式的关键。

7. 完整示例：用 PyTorch 从零实现自注意力

前面已经做了足够的理论铺垫，这一节直接上代码。我们用 PyTorch 从零实现一个完整的自注意力模块，然后在真实数据上跑通验证。

7.1 环境说明

本文的代码基于以下环境：

• Python 3.8+
• PyTorch 2.0+
• 无其他特殊依赖

如果你的环境中没有 PyTorch，可以使用以下命令安装 CPU 版本：

pip install torch --index-url https://download.pytorch.org/whl/cpu

如果你有 NVIDIA GPU，可以参考 PyTorch 官网选择对应的 CUDA 版本安装。

7.2 实现缩放点积注意力

先从最核心的缩放点积注意力开始：

import torch import torch.nn as nn import torch.nn.functional as F def scaled_dot_product_attention(Q, K, V, mask=None): """ 缩放点积注意力。 参数: Q: (..., seq_len_q, d_k) K: (..., seq_len_k, d_k) V: (..., seq_len_k, d_v) mask: (..., seq_len_q, seq_len_k) 可选，bool 或 float 类型 返回: output: (..., seq_len_q, d_v) attention_weights: (..., seq_len_q, seq_len_k) """ d_k = Q.size(-1) # 1. 计算 Q 和 K 的点积，得到相似度矩阵 scores = torch.matmul(Q, K.transpose(-2, -1)) # (..., seq_len_q, seq_len_k) # 2. 缩放 scores = scores / torch.sqrt(torch.tensor(d_k, dtype=torch.float32)) # 3. 应用掩码（如果提供） if mask is not None: scores = scores.masked_fill(mask == 0, float("-inf")) # 4. Softmax 转成概率分布 attention_weights = F.softmax(scores, dim=-1) # 5. 加权求和 output = torch.matmul(attention_weights, V) # (..., seq_len_q, d_v) return output, attention_weights

这段代码实现了完整公式。masked_fill 是 PyTorch 中常用的掩码方式：把 mask 为 0 的分数替换成负无穷，这样经过 Softmax 后权重为 0。

7.3 实现单头自注意力模块

有了核心函数，定义一个完整的自注意力模块：

class SelfAttention(nn.Module): def __init__(self, d_model, d_k=None, d_v=None, dropout=0.1): """ 单头自注意力。 参数: d_model: 输入向量的维度 d_k: Query/Key 的投影维度，默认等于 d_model d_v: Value 的投影维度，默认等于 d_model dropout: Dropout 概率 """ super().__init__() self.d_model = d_model self.d_k = d_k if d_k is not None else d_model self.d_v = d_v if d_v is not None else d_model self.W_Q = nn.Linear(d_model, self.d_k, bias=False) self.W_K = nn.Linear(d_model, self.d_k, bias=False) self.W_V = nn.Linear(d_model, self.d_v, bias=False) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): """ 参数: x: (batch_size, seq_len, d_model) mask: (batch_size, seq_len, seq_len) 或 None 返回: output: (batch_size, seq_len, d_v) attention_weights: (batch_size, seq_len, seq_len) """ Q = self.W_Q(x) # (batch_size, seq_len, d_k) K = self.W_K(x) # (batch_size, seq_len, d_k) V = self.W_V(x) # (batch_size, seq_len, d_v) output, attention_weights = scaled_dot_product_attention(Q, K, V, mask) output = self.dropout(output) return output, attention_weights

这里把线性变换封装为 nn.Linear。bias=False 是遵循原始 Transformer 论文的设置，实际中也可以加上 bias，效果差别不明显。

7.4 实现多头自注意力模块

多头注意力是在单头基础上组合出来的：

class MultiHeadSelfAttention(nn.Module): def __init__(self, d_model, num_heads, dropout=0.1): """ 多头自注意力。 参数: d_model: 输入和输出的维度 num_heads: 头的数量 dropout: Dropout 概率 """ super().__init__() assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除" self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_heads self.d_v = d_model // num_heads self.W_Q = nn.Linear(d_model, d_model, bias=False) self.W_K = nn.Linear(d_model, d_model, bias=False) self.W_V = nn.Linear(d_model, d_model, bias=False) self.W_O = nn.Linear(d_model, d_model, bias=False) self.dropout = nn.Dropout(dropout) def split_heads(self, x): """ 将最后一维拆分成多个头。 输入: (batch_size, seq_len, d_model) 输出: (batch_size, num_heads, seq_len, head_dim) """ batch_size, seq_len, _ = x.size() x = x.view(batch_size, seq_len, self.num_heads, self.d_k) x = x.transpose(1, 2) # (batch_size, num_heads, seq_len, head_dim) return x def combine_heads(self, x): """ 将多个头拼接回原始维度。 输入: (batch_size, num_heads, seq_len, head_dim) 输出: (batch_size, seq_len, d_model) """ batch_size, _, seq_len, _ = x.size() x = x.transpose(1, 2).contiguous() # (batch_size, seq_len, num_heads, head_dim) x = x.view(batch_size, seq_len, self.d_model) return x def forward(self, x, mask=None): batch_size, seq_len, _ = x.size() # 1. 线性投影 Q = self.W_Q(x) # (batch_size, seq_len, d_model) K = self.W_K(x) V = self.W_V(x) # 2. 拆分为多个头 Q = self.split_heads(Q) # (batch_size, num_heads, seq_len, d_k) K = self.split_heads(K) V = self.split_heads(V) # 3. 缩放点积注意力 output, attention_weights = scaled_dot_product_attention(Q, K, V, mask) # 4. 合并头 output = self.combine_heads(output) # (batch_size, seq_len, d_model) # 5. 最后的线性变换 output = self.W_O(output) output = self.dropout(output) return output, attention_weights

这里有一个需要注意的细节：mask 的维度。如果输入 mask 是 (batch_size, seq_len, seq_len)，那在 multi-head 场景下需要扩展为 (batch_size, num_heads, seq_len, seq_len) 才能在多头维度上广播。在后面的示例中我会演示如何构造。

7.5 实现因果自注意力

因果自注意力只需要在 scaled_dot_product_attention 中传入一个上三角掩码：

def create_causal_mask(seq_len): """ 创建因果掩码，形状为 (seq_len, seq_len)。 返回矩阵中： - 1 表示允许关注 - 0 表示禁止关注 """ mask = torch.tril(torch.ones(seq_len, seq_len)) return mask # 使用示例 seq_len = 5 mask = create_causal_mask(seq_len) print(mask)

运行输出：

tensor([[1., 0., 0., 0., 0.], [1., 1., 0., 0., 0.], [1., 1., 1., 0., 0.], [1., 1., 1., 1., 0.], [1., 1., 1., 1., 1.]])

把这个 mask 传给注意力函数，模型在预测第 i 个位置时，就只能看到当前位置及之前的词。

8. 运行与验证：让注意力真正跑起来

现在用一个完整的例子来验证代码。我们构造一个小型的词向量输入，跑一遍自注意力模块，观察输出形状和注意力权重。

import torch import torch.nn as nn # 设置随机种子，保证结果可复现 torch.manual_seed(42) # 参数 batch_size = 2 seq_len = 4 d_model = 16 num_heads = 4 # 模拟输入：随机生成一批词向量 x = torch.randn(batch_size, seq_len, d_model) print(f"输入形状: {x.shape}") # (2, 4, 16) # 单头自注意力 self_attn = SelfAttention(d_model=d_model) output, attn_weights = self_attn(x) print(f"单头输出形状: {output.shape}") # (2, 4, 16) print(f"单头注意力权重形状: {attn_weights.shape}") # (2, 4, 4) print("注意力权重示例（第一个样本第一个位置的权重）:") print(attn_weights[0, 0]) # 多头自注意力 multi_attn = MultiHeadSelfAttention(d_model=d_model, num_heads=num_heads) output, attn_weights = multi_attn(x) print(f"多头输出形状: {output.shape}") # (2, 4, 16) print(f"多头注意力权重形状: {attn_weights.shape}") # (2, 4, 4, 4)

预期输出：

输入形状: torch.Size([2, 4, 16]) 单头输出形状: torch.Size([2, 4, 16]) 单头注意力权重形状: torch.Size([2, 4, 4]) 注意力权重示例（第一个样本第一个位置的权重）: tensor([0.3149, 0.4372, 0.0433, 0.2045], grad_fn=<SelectBackward0>) 多头输出形状: torch.Size([2, 4, 16]) 多头注意力权重形状: torch.Size([2, 4, 4, 4])

如何判断运行是否成功？

• 输出形状必须和输入形状一致（d_model 维度不变），说明残差连接和线性变换的维度没有写错。
• 每一行的注意力权重之和应该等于 1（因为 Softmax 的特性）。
• 反向传播能正常执行：

loss = output.sum() loss.backward() print("反向传播成功")

如果运行失败，优先查看错误信息：

• shape mismatch 错误：检查 Q、K、V 的最后维度是否一致，特别是多头拆分时 d_model 是否能整除 num_heads。
• 显存不足：把 batch_size 和 seq_len 调小，或者把随机输入改为 CPU 上运行。
• 梯度为 None：检查是否有 requires_grad=False 的设置，或者前向传播中没有可学习参数参与计算。

这里再补充一个测试因果掩码的完整示例：

torch.manual_seed(7) # 构造一个小批次 x = torch.randn(1, 5, 16) mask = create_causal_mask(5) print("因果掩码:") print(mask) # 使用掩码运行多头注意力 output, attn_weights = multi_attn(x, mask=mask) print(f"因果注意力输出形状: {output.shape}") # (1, 5, 16) print(f"因果注意力权重形状: {attn_weights.shape}") # (1, 4, 5, 5) # 检查第一个头、第一个位置的注意力权重 print("第一个头、第一个位置的注意力权重:") print(attn_weights[0, 0, 0]) # 逐行验证每行除了对角线及之前位置，权重都为 0 for i in range(5): row = attn_weights[0, 0, i] assert torch.allclose(row[i+1:], torch.zeros_like(row[i+1:])), f"第 {i} 行存在未来信息泄漏" print("因果掩码验证通过：没有未来信息泄漏")

这个验证脚本的最后一部分是关键：它逐行检查了注意力权重矩阵的未来位置是否为 0。如果未来信息泄漏，说明掩码没有正确应用。

9. 从自注意力到完整 Transformer：还需要什么

现在你已经有了一块非常坚实的基石，但 Transformer 架构中除了自注意力，还有几个配套组件需要理解。这里做一个概念补充。

9.1 位置编码

自注意力本身有一个缺陷：它是完全置换不变的。如果你把句子中的词序打乱，自注意力的计算结果不会改变（只要每个词自身的向量不变）。这显然不符合语言的特点——“我打你”和“你打我”含义完全不同。

所以 Transformer 需要额外加入位置编码（Positional Encoding）。最简单的方式是在输入的词向量上叠加一个位置向量，让每个位置拥有独特的标识。原始论文使用的是三角函数编码（sine 和 cosine 函数），现代实现中更常用可学习的位置嵌入。

9.2 残差连接和层归一化

Transformer 的每一层都由两个子层组成：自注意力子层和前馈网络子层。每个子层外面都包裹着：

• 残差连接（Residual Connection）：把子层输入和输出相加，帮助梯度在网络深处顺畅传播。
• 层归一化（Layer Normalization）：对每个样本的特征维度做归一化，加速训练收敛。

这两个组件是训练非常深的 Transformer 网络的关键。

9.3 前馈网络

自注意力负责捕捉序列中词与词的交互，但它是线性运算的组合，非线性表达能力有限。所以每个位置在过完注意力后，还会过一个两层的全连接前馈网络（Feed-Forward Network，FFN），通常是先升维再降维，中间使用 ReLU 或 GELU 激活函数。

在标准 Transformer 中，FFN 的中间层维度通常是 d_model 的 4 倍，比如 d_model=512 时，FFN 的中间维度是 2048。

9.4 编码器-解码器架构

原始 Transformer 是一个编码器-解码器架构：

• 编码器由多个相同的编码器层堆叠而成，每层包含多头自注意力 + 前馈网络。
• 解码器也由多个相同的解码器层堆叠而成，但每层比编码器多一个交叉注意力（Cross-Attention）子层，用于让解码器关注编码器的输出。

对于 BERT 这类预训练模型，只用编码器；对于 GPT 这类生成模型，只用解码器。理解了自注意力机制后，这两条路线都能轻松看懂。

10. 常见问题与排查思路

这一节针对初学者在理解或实现自注意力时最常遇到的问题，做一个整理。

问题现象	可能原因	排查方式	解决方案
注意力权重矩阵全为均匀分布	训练刚开始或模型没有学到有效特征	检查训练步数和损失下降情况	多训练几轮，或检查学习率是否过大
Softmax 后出现 NaN	分数过大导致 Softmax 溢出	检查 Q、K 的数值范围，是否忘了缩放	确认是否除以 sqrt(d_k)；使用 float32 精度
只有对角线附近权重高	模型可能退化成类似 1D 卷积的行为	可视化注意力权重	调大 d_model 或 num_heads，尝试不同初始化
多头注意力输出维度对不上	拆分和合并头的维度变换写错	打印每一步的张量形状	检查 view、transpose、contiguous 的顺序
因果掩码没有生效	mask 形状没有广播到多头维度	打印 mask 形状和注意力权重	确保 mask 扩展到 (batch, num_heads, seq, seq)
训练 loss 不下降	学习率不合适或数据没有归一化	查看梯度范数	调小学习率，检查输入数据分布
显存不足（OOM）	注意力矩阵是 O(n²) 复杂度	查看 seq_len 和 batch_size	减小 batch_size，或使用 Flash Attention 优化

有几个问题单独解释一下。

关于 Flash Attention：标准注意力需要显式创建 n×n 的注意力矩阵，当序列长度很大（比如 4096 或更长）时，这会占用大量显存。Flash Attention 通过分块计算和重计算技巧，在不牺牲效果的前提下大幅降低显存占用。如果你处理长序列文本，建议直接使用 PyTorch 2.0 以上版本内置的torch.nn.functional.scaled_dot_product_attention，它自动适配 Flash Attention 等优化后端。

关于注意力可视化：把注意力权重矩阵用热力图画出来，是检查模型是否学到有效模式的重要手段。横轴是 Key 的位置，纵轴是 Query 的位置，颜色越亮表示权重越高。如果你发现某个头总是集中在对角线上，说明这个头主要关注自己附近的词，可能是合理的（比如局部语法关系），也可能是模型退化了，需要结合具体任务判断。

关于数值稳定性：在实现时，不要直接用scores.exp().sum()来计算 Softmax，应该使用torch.nn.functional.softmax，它内部已经做了减去最大值的数值稳定处理。

11. 最佳实践与工程建议

如果你准备在自己的项目中使用 Transformer 或基于 Transformer 的模型，以下建议值得收藏。

11.1 学习和复现阶段

• 不要一开始就追新变体：先把标准 Transformer 的代码写一遍、跑通一个玩具任务，再去看 BERT、GPT 的代码，会发现阻力小很多。
• 从 PyTorch 内置函数开始：理解原理后，实际训练中可以直接使用torch.nn.MultiheadAttention或torch.nn.functional.scaled_dot_product_attention，它们已经经过充分优化。
• 多可视化注意力权重：注意力热力图是理解模型行为最直观的手段，建议在调试阶段加上。

11.2 工程部署和训练阶段

• 优先使用优化实现：训练长序列模型时，使用 Flash Attention 等优化实现，它能带来数倍的显存节省和速度提升。
• 注意精度选择：混合精度训练（混合精度训练 fp16/bf16）可以显著提速，但需要注意数值稳定性。在自注意力中，Q @ K^T的数值范围较大，fp16 下容易溢出，而 bf16 拥有更大的动态范围，通常更适合训练大模型。
• 评估上下文窗口的边界：Transformer 的理论上下文窗口和实际有效上下文窗口不是一回事。模型训练时见过的最大序列长度就是它的舒适区，一旦推理时超过这个长度，性能会急剧下降。

11.3 数据方面

• 位置编码要适应长度变化：如果训练时位置编码是可学习的，推理时遇到比训练更长的序列就会出问题。实际工程中，要么在训练时做长度增强，要么使用能外推的位置编码方案。
• 不要忽略 Attention 行为的差异：不同随机种子训练出的模型，注意力模式可能截然不同。评估模型质量时不只看一两个样本的注意力可视化，要结合下游任务指标。

11.4 与其他模型的关系

理解自注意力后，你会更容易理解以下模型的演变逻辑：

• BERT：使用 Transformer 编码器的双向语言模型，每个 token 可以看到上下文所有其他 token。
• GPT：使用 Transformer 解码器的自回归语言模型，使用因果自注意力，每个 token 只能看到之前的 token。
• ViT（Vision Transformer）：把图像切成 patch，当作序列输入 Transformer，用自注意力捕捉图像中不同区域的全局关系。
• Swin Transformer：引入窗口注意力来降低视觉任务的注意力计算量。

这些都是同一个思想在不同领域的延伸。

12. 总结与后续学习路径

这篇文章从“为什么 Transformer 会出现”出发，讲清楚了自注意力机制的核心原理、数学公式、代码实现和工程注意事项。现在你应该能够回答这几个问题了：

• 自注意力机制解决了什么根本问题？——让序列中的任意两个位置直接交互，克服了 RNN 顺序处理和长距离衰减的缺陷。
• Q、K、V 的本质是什么？——通过线性变换从输入中提取三种不同的特征视角，动态计算相关性并加权聚合信息。
• 为什么要缩放点积？——保持 Softmax 的梯度稳定，避免 d_k 增大后分数过大导致的梯度消失。
• 多头注意力和单头有什么区别？——多头让模型在多个语义子空间中并行捕捉不同类型的关注关系。
• 因果自注意力是干什么的？——保证生成模型只能看到过去的信息，是实现自回归生成的基础。

如果你要续深入学习 Transformer，建议按下面的路径走：

1. 完整手写一遍 Transformer 层（包含残差连接、层归一化、前馈网络）。
2. 用一个小型数据集（如机器翻译或文本分类）训练一个真实的 Transformer 模型。
3. 读懂 BERT 的代码，看它是如何在 Transformer 编码器上做预训练的。
4. 读懂 GPT 的代码，重点理解因果自注意力和自回归训练方式。
5. 如果有条件，跑一下大模型的微调和推理流程，理解 KV Cache、增量推理等工程优化。

自注意力机制是整个深度学习迈向大模型时代的一把钥匙。这一步真正理解了，后面的路会顺畅很多。建议把这篇文章收藏起来，写代码碰到问题的时候，回来对着公式和实现逐行核对，会比反复查资料更有效。