很多初学 Transformer 的人第一眼看到那幅经典的架构图时通常会有两种感受要么觉得它过于复杂注意力机制、多头、位置编码、残差连接、层归一化一大堆概念堆在一起不知道从哪入手要么觉得它不过如此核心不就是算一个加权平均吗有什么值得大书特书的这两种感受都有道理但也都错过了最关键的东西。Transformer 真正革命性的地方不是那个公式有多复杂而是它彻底改变了深度学习中处理序列数据的方式。在它之前RNN、LSTM 是处理序列事实上的标准选择它们像人类阅读一样逐字处理输入顺序执行、难以并行长距离依赖问题虽然在 LSTM 中有所缓解但没有被真正解决。Transformer 则完全抛弃了循环结构用一种叫自注意力Self-Attention的机制让序列中的任意两个位置直接建立联系。这一改变带来的连锁反应是巨大的并行计算成为可能训练效率大幅提升长距离依赖不再是瓶颈BERT、GPT 这类超大规模预训练模型才有了存在的根基。可以说自注意力机制是理解整个 Transformer 大厦的第一块地基也是从传统深度学习走向现代大模型必须跨越的一道门槛。这篇文章会从最朴素的问题出发用场景化的方式讲清楚自注意力机制到底解决了什么问题、它的计算过程每一步在做什么、为什么需要缩放点积、多头注意力又为什么要存在并配合完整的 PyTorch 代码实现帮你把理论和代码对上号。读完这篇文章你再看 Transformer 的任何变体、再看大模型的原始论文都会感觉顺很多。1. 这篇文章真正要解决的问题先做一个判断自注意力机制是 Transformer 的核心创新也是理解全部现代大模型的最低门槛。很多人学 Transformer 时遇到的困境不是不想学而是被资料淹没。打开搜索引擎教程铺天盖地但大体上有两类一类是纯概念科普把注意力机制讲成了一个神奇的黑盒放几张热力图告诉你模型在关注哪些词看起来很直观但你合上页面依然不知道 QQuery、KKey、VValue到底怎么来的更不知道代码从哪一行开始写。另一类是纯数学推导把注意力公式用大篇幅展开算给你看内积、归一化、加权求和逻辑没问题但你不知道这些矩阵运算放在神经网络里到底承担了什么角色为什么要除一个根号 dk这个操作不合理会怎样。更常见的一个痛点是很多人习惯了 CNN 和 RNN 的思维会把自注意力理解为某种“更复杂的卷积”或“更聪明的 LSTM”这种先入为主的类比框架反而成了理解 Transformer 的最大障碍。这篇文章想解决的问题就是把心法和招式结合起来用场景讲直觉为什么需要自注意力它比 RNN 强在哪用公式讲原理Q、K、V 的本质是什么缩放点积为什么要有缩放用代码讲实现从零实现一个完整的自注意力模块让你真的跑起来。用对比讲延展多头注意力、因果自注意力是怎么从基础版本演化出来的如果你正在学深度学习或者刚接触大模型相关的技术甚至是你已经在用 BERT、GPT 的 API 但想深入了解它们内部机制这篇文章都适合你。它不要求你有深厚的数学背景只要会基本的矩阵乘法和 python 编程就能跟上节奏。2. 从 Seq2Seq 到 Self-Attention注意力机制到底在解决什么问题要理解自注意力必须先从它的前身——经典的注意力机制Attention讲起。2.1 没有注意力机制的时代信息瓶颈问题在注意力机制出现之前序列到序列Seq2Seq任务的标准架构是这样的一个编码器EncoderRNN/LSTM 把输入序列逐字读进去最后汇总成一个固定长度的向量这个向量被称作“上下文向量”。一个解码器DecoderRNN/LSTM 从这个上下文向量出发一步步生成输出序列。这个架构存在一个致命问题信息瓶颈。无论输入序列多长编码器都必须把所有信息压进最后一个隐藏状态里。如果句子很长开头的信息在传递过程中会被不断稀释最终到达解码器时可能已经损失殆尽。这就是长距离依赖问题的本质。2.2 注意力机制的诞生按需获取信息注意力机制的突破性想法是解码器在生成每一步的时候不再只依赖一个固定的上下文向量而是可以用一个“查询”去编码器的所有位置上检索信息然后根据相关性决定要看哪里、看多少。这里借用一下信息检索的类比非常直观Query查询你可以理解为你大脑中正在思考的问题。比如你要翻译一个词当前的生成状态就是你的 Query相当于你在问“现在该看源句子的哪个词”Key键源序列中每个词的一个索引标签。你会拿 Query 和每个 Key 做匹配判断“这个词跟我当前的问题有关系吗”Value值源序列中每个词的实际内容。一旦根据匹配程度确定了该关注哪些词就去取这些词的内容来用。整个过程就是三步把 Query 和每个 Key 做相似度计算通过 Softmax 转成权重再用权重对 Value 加权求和。注意力得分 Softmax(Query 与每个 Key 的相似度) × Value2.3 什么是自注意力关注的不是“对方”而是“自己”经典注意力发生在编码器和解码器之间建立在两个不同序列上。而自注意力Self-Attention则是一个序列内部自己和自己计算注意力所以它还有另一个名字内部注意力Intra-Attention。自注意力的核心思想是对于一个句子中的每个词通过计算它和其他所有词的相关性来更新它自己的语义表示。举个例子处理这句话The animal didnt 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、VQ 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 本身是词向量它是静态的而通过不同的线性变换模型可以在不同的语义空间中衡量“提问”和“回答”的关系。没有这个变换所有词都在同一个空间里直接做内积表达能力会受到很大限制。从工程上讲如果不经过变换直接让 QX、KX那么计算出的注意力矩阵就是固定的。比如某个位置和其他位置的相似度完全由词向量的静态相似度决定相当于模型没有学习余地。引入三个可学习的权重矩阵后这个相似度计算方式本身可以通过反向传播被优化模型的能力上限就完全不同了。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 把分数转换成概率分布保证每一行的权重和为 1A 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_k2数值是编造的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.414S_scaled: [[0.523, 0.325, 0.226], [0.573, 0.438, 0.410], [0.403, 0.481, 0.523]]第三步对每一行做 Softmax。以第一行为例计算 expexp(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 中通常设置 h8 个头每个头的维度 d_k d_v d_model / h。比如 d_model512 时每个头的维度是 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.8PyTorch 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, maskNone): 缩放点积注意力。 参数: 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, dtypetorch.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_kNone, d_vNone, dropout0.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, biasFalse) self.W_K nn.Linear(d_model, self.d_k, biasFalse) self.W_V nn.Linear(d_model, self.d_v, biasFalse) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): 参数: 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。biasFalse 是遵循原始 Transformer 论文的设置实际中也可以加上 bias效果差别不明显。7.4 实现多头自注意力模块多头注意力是在单头基础上组合出来的class MultiHeadSelfAttention(nn.Module): def __init__(self, d_model, num_heads, dropout0.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, biasFalse) self.W_K nn.Linear(d_model, d_model, biasFalse) self.W_V nn.Linear(d_model, d_model, biasFalse) self.W_O nn.Linear(d_model, d_model, biasFalse) 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, maskNone): 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_modeld_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_modeld_model, num_headsnum_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_fnSelectBackward0) 多头输出形状: 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_gradFalse 的设置或者前向传播中没有可学习参数参与计算。这里再补充一个测试因果掩码的完整示例torch.manual_seed(7) # 构造一个小批次 x torch.randn(1, 5, 16) mask create_causal_mask(5) print(因果掩码:) print(mask) # 使用掩码运行多头注意力 output, attn_weights multi_attn(x, maskmask) 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[i1:], torch.zeros_like(row[i1:])), 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 NetworkFFN通常是先升维再降维中间使用 ReLU 或 GELU 激活函数。在标准 Transformer 中FFN 的中间层维度通常是 d_model 的 4 倍比如 d_model512 时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。ViTVision Transformer把图像切成 patch当作序列输入 Transformer用自注意力捕捉图像中不同区域的全局关系。Swin Transformer引入窗口注意力来降低视觉任务的注意力计算量。这些都是同一个思想在不同领域的延伸。12. 总结与后续学习路径这篇文章从“为什么 Transformer 会出现”出发讲清楚了自注意力机制的核心原理、数学公式、代码实现和工程注意事项。现在你应该能够回答这几个问题了自注意力机制解决了什么根本问题——让序列中的任意两个位置直接交互克服了 RNN 顺序处理和长距离衰减的缺陷。Q、K、V 的本质是什么——通过线性变换从输入中提取三种不同的特征视角动态计算相关性并加权聚合信息。为什么要缩放点积——保持 Softmax 的梯度稳定避免 d_k 增大后分数过大导致的梯度消失。多头注意力和单头有什么区别——多头让模型在多个语义子空间中并行捕捉不同类型的关注关系。因果自注意力是干什么的——保证生成模型只能看到过去的信息是实现自回归生成的基础。如果你要续深入学习 Transformer建议按下面的路径走完整手写一遍 Transformer 层包含残差连接、层归一化、前馈网络。用一个小型数据集如机器翻译或文本分类训练一个真实的 Transformer 模型。读懂 BERT 的代码看它是如何在 Transformer 编码器上做预训练的。读懂 GPT 的代码重点理解因果自注意力和自回归训练方式。如果有条件跑一下大模型的微调和推理流程理解 KV Cache、增量推理等工程优化。自注意力机制是整个深度学习迈向大模型时代的一把钥匙。这一步真正理解了后面的路会顺畅很多。建议把这篇文章收藏起来写代码碰到问题的时候回来对着公式和实现逐行核对会比反复查资料更有效。 自注意力机制深度解析：从原理到PyTorch实现
很多初学 Transformer 的人第一眼看到那幅经典的架构图时通常会有两种感受要么觉得它过于复杂注意力机制、多头、位置编码、残差连接、层归一化一大堆概念堆在一起不知道从哪入手要么觉得它不过如此核心不就是算一个加权平均吗有什么值得大书特书的这两种感受都有道理但也都错过了最关键的东西。Transformer 真正革命性的地方不是那个公式有多复杂而是它彻底改变了深度学习中处理序列数据的方式。在它之前RNN、LSTM 是处理序列事实上的标准选择它们像人类阅读一样逐字处理输入顺序执行、难以并行长距离依赖问题虽然在 LSTM 中有所缓解但没有被真正解决。Transformer 则完全抛弃了循环结构用一种叫自注意力Self-Attention的机制让序列中的任意两个位置直接建立联系。这一改变带来的连锁反应是巨大的并行计算成为可能训练效率大幅提升长距离依赖不再是瓶颈BERT、GPT 这类超大规模预训练模型才有了存在的根基。可以说自注意力机制是理解整个 Transformer 大厦的第一块地基也是从传统深度学习走向现代大模型必须跨越的一道门槛。这篇文章会从最朴素的问题出发用场景化的方式讲清楚自注意力机制到底解决了什么问题、它的计算过程每一步在做什么、为什么需要缩放点积、多头注意力又为什么要存在并配合完整的 PyTorch 代码实现帮你把理论和代码对上号。读完这篇文章你再看 Transformer 的任何变体、再看大模型的原始论文都会感觉顺很多。1. 这篇文章真正要解决的问题先做一个判断自注意力机制是 Transformer 的核心创新也是理解全部现代大模型的最低门槛。很多人学 Transformer 时遇到的困境不是不想学而是被资料淹没。打开搜索引擎教程铺天盖地但大体上有两类一类是纯概念科普把注意力机制讲成了一个神奇的黑盒放几张热力图告诉你模型在关注哪些词看起来很直观但你合上页面依然不知道 QQuery、KKey、VValue到底怎么来的更不知道代码从哪一行开始写。另一类是纯数学推导把注意力公式用大篇幅展开算给你看内积、归一化、加权求和逻辑没问题但你不知道这些矩阵运算放在神经网络里到底承担了什么角色为什么要除一个根号 dk这个操作不合理会怎样。更常见的一个痛点是很多人习惯了 CNN 和 RNN 的思维会把自注意力理解为某种“更复杂的卷积”或“更聪明的 LSTM”这种先入为主的类比框架反而成了理解 Transformer 的最大障碍。这篇文章想解决的问题就是把心法和招式结合起来用场景讲直觉为什么需要自注意力它比 RNN 强在哪用公式讲原理Q、K、V 的本质是什么缩放点积为什么要有缩放用代码讲实现从零实现一个完整的自注意力模块让你真的跑起来。用对比讲延展多头注意力、因果自注意力是怎么从基础版本演化出来的如果你正在学深度学习或者刚接触大模型相关的技术甚至是你已经在用 BERT、GPT 的 API 但想深入了解它们内部机制这篇文章都适合你。它不要求你有深厚的数学背景只要会基本的矩阵乘法和 python 编程就能跟上节奏。2. 从 Seq2Seq 到 Self-Attention注意力机制到底在解决什么问题要理解自注意力必须先从它的前身——经典的注意力机制Attention讲起。2.1 没有注意力机制的时代信息瓶颈问题在注意力机制出现之前序列到序列Seq2Seq任务的标准架构是这样的一个编码器EncoderRNN/LSTM 把输入序列逐字读进去最后汇总成一个固定长度的向量这个向量被称作“上下文向量”。一个解码器DecoderRNN/LSTM 从这个上下文向量出发一步步生成输出序列。这个架构存在一个致命问题信息瓶颈。无论输入序列多长编码器都必须把所有信息压进最后一个隐藏状态里。如果句子很长开头的信息在传递过程中会被不断稀释最终到达解码器时可能已经损失殆尽。这就是长距离依赖问题的本质。2.2 注意力机制的诞生按需获取信息注意力机制的突破性想法是解码器在生成每一步的时候不再只依赖一个固定的上下文向量而是可以用一个“查询”去编码器的所有位置上检索信息然后根据相关性决定要看哪里、看多少。这里借用一下信息检索的类比非常直观Query查询你可以理解为你大脑中正在思考的问题。比如你要翻译一个词当前的生成状态就是你的 Query相当于你在问“现在该看源句子的哪个词”Key键源序列中每个词的一个索引标签。你会拿 Query 和每个 Key 做匹配判断“这个词跟我当前的问题有关系吗”Value值源序列中每个词的实际内容。一旦根据匹配程度确定了该关注哪些词就去取这些词的内容来用。整个过程就是三步把 Query 和每个 Key 做相似度计算通过 Softmax 转成权重再用权重对 Value 加权求和。注意力得分 Softmax(Query 与每个 Key 的相似度) × Value2.3 什么是自注意力关注的不是“对方”而是“自己”经典注意力发生在编码器和解码器之间建立在两个不同序列上。而自注意力Self-Attention则是一个序列内部自己和自己计算注意力所以它还有另一个名字内部注意力Intra-Attention。自注意力的核心思想是对于一个句子中的每个词通过计算它和其他所有词的相关性来更新它自己的语义表示。举个例子处理这句话The animal didnt 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、VQ 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 本身是词向量它是静态的而通过不同的线性变换模型可以在不同的语义空间中衡量“提问”和“回答”的关系。没有这个变换所有词都在同一个空间里直接做内积表达能力会受到很大限制。从工程上讲如果不经过变换直接让 QX、KX那么计算出的注意力矩阵就是固定的。比如某个位置和其他位置的相似度完全由词向量的静态相似度决定相当于模型没有学习余地。引入三个可学习的权重矩阵后这个相似度计算方式本身可以通过反向传播被优化模型的能力上限就完全不同了。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 把分数转换成概率分布保证每一行的权重和为 1A 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_k2数值是编造的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.414S_scaled: [[0.523, 0.325, 0.226], [0.573, 0.438, 0.410], [0.403, 0.481, 0.523]]第三步对每一行做 Softmax。以第一行为例计算 expexp(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 中通常设置 h8 个头每个头的维度 d_k d_v d_model / h。比如 d_model512 时每个头的维度是 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.8PyTorch 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, maskNone): 缩放点积注意力。 参数: 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, dtypetorch.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_kNone, d_vNone, dropout0.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, biasFalse) self.W_K nn.Linear(d_model, self.d_k, biasFalse) self.W_V nn.Linear(d_model, self.d_v, biasFalse) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): 参数: 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。biasFalse 是遵循原始 Transformer 论文的设置实际中也可以加上 bias效果差别不明显。7.4 实现多头自注意力模块多头注意力是在单头基础上组合出来的class MultiHeadSelfAttention(nn.Module): def __init__(self, d_model, num_heads, dropout0.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, biasFalse) self.W_K nn.Linear(d_model, d_model, biasFalse) self.W_V nn.Linear(d_model, d_model, biasFalse) self.W_O nn.Linear(d_model, d_model, biasFalse) 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, maskNone): 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_modeld_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_modeld_model, num_headsnum_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_fnSelectBackward0) 多头输出形状: 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_gradFalse 的设置或者前向传播中没有可学习参数参与计算。这里再补充一个测试因果掩码的完整示例torch.manual_seed(7) # 构造一个小批次 x torch.randn(1, 5, 16) mask create_causal_mask(5) print(因果掩码:) print(mask) # 使用掩码运行多头注意力 output, attn_weights multi_attn(x, maskmask) 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[i1:], torch.zeros_like(row[i1:])), 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 NetworkFFN通常是先升维再降维中间使用 ReLU 或 GELU 激活函数。在标准 Transformer 中FFN 的中间层维度通常是 d_model 的 4 倍比如 d_model512 时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。ViTVision Transformer把图像切成 patch当作序列输入 Transformer用自注意力捕捉图像中不同区域的全局关系。Swin Transformer引入窗口注意力来降低视觉任务的注意力计算量。这些都是同一个思想在不同领域的延伸。12. 总结与后续学习路径这篇文章从“为什么 Transformer 会出现”出发讲清楚了自注意力机制的核心原理、数学公式、代码实现和工程注意事项。现在你应该能够回答这几个问题了自注意力机制解决了什么根本问题——让序列中的任意两个位置直接交互克服了 RNN 顺序处理和长距离衰减的缺陷。Q、K、V 的本质是什么——通过线性变换从输入中提取三种不同的特征视角动态计算相关性并加权聚合信息。为什么要缩放点积——保持 Softmax 的梯度稳定避免 d_k 增大后分数过大导致的梯度消失。多头注意力和单头有什么区别——多头让模型在多个语义子空间中并行捕捉不同类型的关注关系。因果自注意力是干什么的——保证生成模型只能看到过去的信息是实现自回归生成的基础。如果你要续深入学习 Transformer建议按下面的路径走完整手写一遍 Transformer 层包含残差连接、层归一化、前馈网络。用一个小型数据集如机器翻译或文本分类训练一个真实的 Transformer 模型。读懂 BERT 的代码看它是如何在 Transformer 编码器上做预训练的。读懂 GPT 的代码重点理解因果自注意力和自回归训练方式。如果有条件跑一下大模型的微调和推理流程理解 KV Cache、增量推理等工程优化。自注意力机制是整个深度学习迈向大模型时代的一把钥匙。这一步真正理解了后面的路会顺畅很多。建议把这篇文章收藏起来写代码碰到问题的时候回来对着公式和实现逐行核对会比反复查资料更有效。