自注意力机制详解从零开始拆解 Transformer 的核心原理这次我们直接切入一个绕不开的话题自注意力Self-Attention。如果你看过 Transformer 相关的论文、博客或者源码一定见过那张经典的 Attention 计算图Q、K、V 三个矩阵经过 MatMul、Scale、Mask、Softmax、MatMul 最后得到输出。问题在于图能看懂但真正动手实现时维度怎么对齐、Mask 怎么加、多头怎么拆很多细节一上手就乱。这篇文章不绕弯子直接拆解自注意力机制的基础原理和工程实现。你会看到自注意力到底在算什么、为什么需要缩放、Padding Mask 和因果 Mask 有什么区别、多头注意力Multi-Head Attention又是怎么组织的。最后我会用 PyTorch 从零手写一个完整的自注意力模块并给出可运行的示例代码。适合正在学习 Transformer、准备面试、或者想自己实现一遍注意力机制的读者。1. 核心概念速览概念说明核心作用建模序列中任意两个位置之间的依赖关系不受距离限制输入形式一个序列的 Embedding 向量形状为[batch_size, seq_len, d_model]核心参数d_model特征维度、num_heads注意力头数、d_k每个头的维度计算流程生成 Q/K/V - 缩放点积 - Softmax - 加权求和关键机制缩放Scale、掩码Mask、多头Multi-Head与 RNN 区别并行计算所有位置不需要按时间步逐个处理时间复杂度O(n²·d)n 为序列长度d 为特征维度适用场景文本、语音、图像序列等任意序列数据的特征提取常见实现PyTorchnn.MultiheadAttention、Hugging Face Transformer 实现学习难度中等数学门槛不高理解矩阵运算即可实现2. 自注意力要解决什么问题在自注意力出现之前处理序列数据的主流方案是 RNN 和 CNN。RNN 的问题在于串行计算。以 LSTM 为例处理第 t 个词的时候必须依赖第 t-1 个词的隐状态输出。这就导致两个问题第一训练速度慢长序列没办法并行第二长距离依赖难建模虽然 LSTM 引入了门控机制但信息在传递过程中仍然会衰减相隔很远的两个词之间的依赖关系很难被捕捉。CNN 的问题在于感受野受限。虽然可以通过堆叠多层来扩大感受野但要建模长距离依赖需要堆很多层计算量不小而且卷积核的局部性限制了它对全局信息的直接感知。自注意力机制一次性解决了这两个问题。给定一个序列自注意力直接计算序列中任意两个位置之间的关联权重也就是说第 1 个词和第 100 个词之间的依赖关系可以直接建立不需要通过中间步骤传递。同时所有位置的注意力权重可以并行计算训练效率大大提升。一句话总结自注意力让每个位置都能直接看到序列中的所有位置并且这个计算过程可以被高效并行化。这就是为什么 Transformer 能取代 RNN 成为大模型的主流架构。3. 自注意力机制的数学原理自注意力的输入是一个序列的向量表示。假设输入的序列长度为 n每个 token 的向量维度是 d_model那么输入矩阵 X 的形状是[seq_len, d_model]为简化说明先忽略 batch 维度。3.1 生成 Q、K、V首先对输入 X 做三次线性变换得到三个矩阵QQuery查询向量表示“我想找什么”KKey键向量表示“我是什么”VValue值向量表示“我能提供什么内容”计算公式为Q X W_Q K X W_K V X W_V其中 W_Q、W_K、W_V 都是可学习的权重矩阵形状为[d_model, d_k]。这里的 d_k 是每个注意力头中 Q、K、V 的维度。3.2 计算注意力分数注意力分数的计算方式是 Q 和 K 的点积scores Q K^T得到的结果形状是[seq_len, seq_len]表示序列中每个位置对每个位置的匹配程度。第 i 行第 j 列的值越大表示第 i 个 token 越关注第 j 个 token。3.3 缩放直接使用点积结果会有一个问题当 d_k 比较大时点积结果的方差会变大导致 softmax 之后梯度极小训练不稳定。所以需要除以一个缩放因子scores scores / sqrt(d_k)这样做的原因是当两个 d_k 维向量各分量都是独立随机变量均值 0、方差 1时它们的点积的方差约为 d_k。除以 sqrt(d_k) 后方差被拉回 1 左右softmax 的梯度更稳定。3.4 Softmax 归一化对每一行做 softmax让同一位置对所有位置的注意力权重之和为 1weights softmax(scores, dim-1)3.5 加权求和用注意力权重对 V 进行加权求和得到输出output weights V输出形状为[seq_len, d_k]。这一步可以理解为根据“每个位置应该关注谁”把 V 中的信息按权重融合起来。3.6 整体计算图把上面的步骤连起来Attention(Q, K, V) softmax(Q K^T / sqrt(d_k)) V这就是缩放点积注意力Scaled Dot-Product Attention的完整公式。4. 为什么需要缩放和掩码这两个细节是面试高频考点也是实现时最容易出错的地方。4.1 缩放的意义前面已经提到缩放是为了让点积结果的方差稳定避免 softmax 进入饱和区。具体来说如果不缩放当 d_k 很大时点积结果的数值会很大导致 softmax 的输出非常接近 one-hot 分布最大值接近 1其他接近 0。这样梯度过小训练几乎无法进行。除以 sqrt(d_k) 后数值范围被拉回到合理区间。这是 Google 在《Attention Is All You Need》论文中给出的解释实际工程中也验证了缩放的必要性。4.2 Mask 的作用Mask 有两种常见类型Padding Mask 和 Causal Mask。Padding Mask用于处理变长序列。一个 batch 中不同样本的序列长度不同短的序列尾部需要 padding 成同一个长度。padding 的位置没有实际含义在计算注意力权重时应该被屏蔽。实现方式是在 softmax 之前把 padding 位置的分数设为一个非常大的负数如 -1e9这样 softmax 后对应权重趋近于 0。Causal Mask用于自回归生成任务如 GPT 系列。每个位置只能看到当前及之前的位置不能看到未来的 token。这通过在分数矩阵上叠加一个上三角矩阵的负无穷来实现强制未来位置的分数为 -infsoftmax 后权重为 0。两种 mask 可以同时使用。例如在训练 decoder 时既要屏蔽 padding 位置也要屏蔽未来位置。5. 多头注意力机制多头注意力不是做一次注意力而是把 Q、K、V 分成 h 个头每个头独立计算注意力最后把结果拼起来再做一次线性投影。5.1 为什么需要多头单头注意力只能建模一种“关注模式”。但句子中可能同时存在多种依赖关系语法关系、语义关系、指代关系。多头注意力让模型能够同时关注不同维度的信息每个头学到不同的表示子空间。5.2 实现方式假设 d_model 512num_heads 8那么每个头的维度 d_k 512 / 8 64。实现上有两种方式第一种把 W_Q、W_K、W_V 拆成 8 份每份形状为[512, 64]分别计算 8 次注意力最后拼接。第二种也是常用做法一次性生成整个 Q、K、V形状都是[batch_size, seq_len, 512]然后 reshape 成[batch_size, seq_len, num_heads, d_k]再交换维度变成[batch_size, num_heads, seq_len, d_k]直接批量计算注意力。第二种方式在 PyTorch 中更容易实现后面会在代码中演示。5.3 输出投影多头注意力的输出拼接后形状为[batch_size, seq_len, d_model]再经过一个线性层 W_O 进行输出投影output concat(head_1, ..., head_h) W_O6. PyTorch 从零实现自注意力现在我们用 PyTorch 实现一个完整的自注意力模块。不依赖nn.MultiheadAttention从零手写方便看清每一步的维度变化。6.1 单头缩放点积注意力import torch import torch.nn as nn import torch.nn.functional as F def scaled_dot_product_attention(query, key, value, maskNone): 缩放点积注意力。 参数: query: [batch_size, seq_len_q, d_k] key: [batch_size, seq_len_k, d_k] value: [batch_size, seq_len_k, d_v] mask: [batch_size, seq_len_q, seq_len_k] 或可广播版本 返回: output: [batch_size, seq_len_q, d_v] weights: [batch_size, seq_len_q, seq_len_k] d_k query.size(-1) # 1. 计算 Q 和 K 的点积 scores torch.matmul(query, key.transpose(-2, -1)) # [B, 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, -1e9) # 4. Softmax weights F.softmax(scores, dim-1) # 5. 加权求和 output torch.matmul(weights, value) # [B, seq_len_q, d_v] return output, weights6.2 多头自注意力完整实现class MultiHeadSelfAttention(nn.Module): def __init__(self, d_model, num_heads): 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 # 生成 Q、K、V 和输出投影的权重 self.w_q nn.Linear(d_model, d_model) self.w_k nn.Linear(d_model, d_model) self.w_v nn.Linear(d_model, d_model) self.w_o nn.Linear(d_model, d_model) def split_heads(self, x): 输入: [batch_size, seq_len, d_model] 输出: [batch_size, num_heads, seq_len, d_k] batch_size, seq_len, _ x.size() x x.view(batch_size, seq_len, self.num_heads, self.d_k) return x.transpose(1, 2) def combine_heads(self, x): 输入: [batch_size, num_heads, seq_len, d_k] 输出: [batch_size, seq_len, d_model] batch_size, _, seq_len, _ x.size() x x.transpose(1, 2).contiguous() return x.view(batch_size, seq_len, self.d_model) def forward(self, x, maskNone): # x: [batch_size, seq_len, d_model] # 1. 线性变换生成 Q、K、V q self.w_q(x) k self.w_k(x) v self.w_v(x) # 2. 拆分成多头 q self.split_heads(q) # [B, num_heads, seq_len, d_k] k self.split_heads(k) v self.split_heads(v) # 3. 对每个头独立计算注意力 # 注意 mask 需要能广播到 [B, num_heads, seq_len, seq_len] attn_output, weights scaled_dot_product_attention(q, k, v, mask) # 4. 合并多头输出投影 output self.combine_heads(attn_output) output self.w_o(output) return output, weights6.3 使用示例# 测试输入 batch_size 2 seq_len 10 d_model 512 num_heads 8 x torch.randn(batch_size, seq_len, d_model) model MultiHeadSelfAttention(d_model, num_heads) output, weights model(x) print(f输入形状: {x.shape}) print(f输出形状: {output.shape}) print(f注意力权重形状: {weights.shape})输出结果输入形状: torch.Size([2, 10, 512]) 输出形状: torch.Size([2, 10, 512]) 注意力权重形状: torch.Size([2, 8, 10, 10])7. 自注意力与 RNN、CNN 的对比特性RNN / LSTMCNN自注意力并行性差串行计算好好长距离依赖弱信息衰减弱需堆叠层数强任意位置直接关联计算复杂度每步 O(d²)总 O(n·d²)O(k·n·d²)O(n²·d)可解释性一般不直接观察特征图可视化注意力权重可直接可视化对序列长度扩展性较好较好差n² 复杂度从表格可以看到自注意力的最大瓶颈是O(n²) 的序列长度复杂度。当序列长度 n 很大时比如处理长文档或高分辨率图像注意力矩阵的计算和存储开销会迅速膨胀。这也是为什么后续出现了各种优化变体比如稀疏注意力Sparse Attention、局部注意力Local Attention、线性注意力Linear Attention等目的都是降低 n² 复杂度。8. 资源占用与性能观察先看一个常见问题为什么显存经常爆自注意力的中间张量weights的形状是[batch_size, num_heads, seq_len, seq_len]。假设 batch_size4num_heads8seq_len2048这个注意力矩阵就有4 × 8 × 2048 × 2048 1.34 亿个 float32 元素占 512 MB 显存。如果 seq_len 翻倍到 4096显存占用直接变成原来的 4 倍达到 2 GB 以上。观察资源占用可以从这几个维度入手显存占用使用 Nvidia 的nvidia-smi命令实时查看。重点关注训练或推理过程中的显存峰值而不是启动时的静态占用。计算耗时分别用不同 seq_len 跑几次前向观察耗时增长曲线。你会发现 seq_len 从 512 到 1024耗时大约增加 4 倍这和 O(n²) 的复杂度是吻合的。批量大小的影响batch_size 增大会让显存占用接近线性增长但计算单元利用率也更高。如果显存不够优先减小 batch_size而不是减小 seq_len。一个实用的降低显存手段是使用梯度检查点Gradient Checkpointing在反向传播时重新计算前向结果而不是全部保存。以时间换显存适合长序列训练。9. 常见问题与排查方法问题现象可能原因排查方式解决方案训练 loss 不收敛或震荡缩放因子未实现softmax 进入饱和区检查是否除以 sqrt(d_k)加上缩放确认数值范围注意力权重全是同一个分布初始化不当或训练不充分打印 attention weights 观察调整初始化方式增加训练步数序列长度不一致导致报错缺少 Padding Mask查看 batch 中 tensor 形状使用pad_sequence并对 padding 位置做 Maskdecoder 训练时未来信息泄漏缺少 Causal Mask检查 loss 是否异常低叠加上三角 Mask屏蔽未来位置显存溢出OOM序列长度过长或 batch_size 过大查看nvidia-smi显存占用减小 batch_size、梯度检查点、混合精度训练多头 reshape 报错d_model 无法被 num_heads 整除检查d_model % num_heads 0调整 d_model 或 num_heads权重形状不对维度转置或 view 顺序理解错误打印每一步张量形状按照split_heads/combine_heads的流程逐步检查10. 最佳实践与使用建议10.1 从标准实现开始先不用急着上 Flash Attention 等优化版。用 PyTorch 从零写一遍标准自注意力确认每一步维度都正确后再考虑优化。这样对后续读源码、调试模型帮助最大。10.2 注意数值稳定性软注意力的实现里masked_fill使用的负值要够大但不能大到造成数值溢出。常见的做法是使用-1e9或直接用 PyTorch 的float(-inf)配合masked_fill操作。10.3 序列长度和显存要提前预估在开始训练前用一个小规模的输入先跑一次前向和反向记录显存占用。再按公式估算更大序列长度时的显存需求避免训练中途 OOM。10.4 合理选择掩码编码器Encoder只用 Padding Mask。解码器DecoderPadding Mask Causal Mask 同时使用。10.5 可视化注意力权重训练完成后把注意力权重提取出来做可视化可以帮助定位模型学到了哪些依赖关系。如果权重分布接近均匀分布说明模型没有学到有效的信息如果某个位置出现明显的“关注集中”说明该位置在语义上很重要。11. 总结与下一步自注意力机制是 Transformer 的基石理解它的关键点可以归纳为四条第一它通过 Q、K、V 三组向量完成“查询-匹配-提取”的过程让序列中的任意两个位置可以直接建立关联。第二缩放因子 sqrt(d_k) 不是公式修饰而是稳定梯度、保证训练收敛的必要设计。第三Mask 机制解决了两类实际问题变长序列的 padding 处理和自回归模型的未来信息泄漏。第四多头注意力让模型在多个表示子空间中并行建模不同的依赖关系这是 Transformer 表达能力的重要来源。最容易踩的坑有两个一是忽略缩放导致训练不收敛二是多头 reshape 时维度没理清导致报错。建议你按照本文的代码把每一步的中间张量形状打印出来逐一对照很快就能建立手感。下一步你可以去读 Transformer 论文的原文重点看 3.2 节的注意力公式和 3.2.2 节的多头注意力解释然后再读 PyTorch 官方实现的nn.MultiheadAttention源码看官方是怎么做维度变换和 mask 处理的。如果想深入做性能优化可以继续学习 Flash Attention 和 KV Cache 的原理这两块在推理加速中非常关键。 自注意力机制详解：从零拆解Transformer核心原理
自注意力机制详解从零开始拆解 Transformer 的核心原理这次我们直接切入一个绕不开的话题自注意力Self-Attention。如果你看过 Transformer 相关的论文、博客或者源码一定见过那张经典的 Attention 计算图Q、K、V 三个矩阵经过 MatMul、Scale、Mask、Softmax、MatMul 最后得到输出。问题在于图能看懂但真正动手实现时维度怎么对齐、Mask 怎么加、多头怎么拆很多细节一上手就乱。这篇文章不绕弯子直接拆解自注意力机制的基础原理和工程实现。你会看到自注意力到底在算什么、为什么需要缩放、Padding Mask 和因果 Mask 有什么区别、多头注意力Multi-Head Attention又是怎么组织的。最后我会用 PyTorch 从零手写一个完整的自注意力模块并给出可运行的示例代码。适合正在学习 Transformer、准备面试、或者想自己实现一遍注意力机制的读者。1. 核心概念速览概念说明核心作用建模序列中任意两个位置之间的依赖关系不受距离限制输入形式一个序列的 Embedding 向量形状为[batch_size, seq_len, d_model]核心参数d_model特征维度、num_heads注意力头数、d_k每个头的维度计算流程生成 Q/K/V - 缩放点积 - Softmax - 加权求和关键机制缩放Scale、掩码Mask、多头Multi-Head与 RNN 区别并行计算所有位置不需要按时间步逐个处理时间复杂度O(n²·d)n 为序列长度d 为特征维度适用场景文本、语音、图像序列等任意序列数据的特征提取常见实现PyTorchnn.MultiheadAttention、Hugging Face Transformer 实现学习难度中等数学门槛不高理解矩阵运算即可实现2. 自注意力要解决什么问题在自注意力出现之前处理序列数据的主流方案是 RNN 和 CNN。RNN 的问题在于串行计算。以 LSTM 为例处理第 t 个词的时候必须依赖第 t-1 个词的隐状态输出。这就导致两个问题第一训练速度慢长序列没办法并行第二长距离依赖难建模虽然 LSTM 引入了门控机制但信息在传递过程中仍然会衰减相隔很远的两个词之间的依赖关系很难被捕捉。CNN 的问题在于感受野受限。虽然可以通过堆叠多层来扩大感受野但要建模长距离依赖需要堆很多层计算量不小而且卷积核的局部性限制了它对全局信息的直接感知。自注意力机制一次性解决了这两个问题。给定一个序列自注意力直接计算序列中任意两个位置之间的关联权重也就是说第 1 个词和第 100 个词之间的依赖关系可以直接建立不需要通过中间步骤传递。同时所有位置的注意力权重可以并行计算训练效率大大提升。一句话总结自注意力让每个位置都能直接看到序列中的所有位置并且这个计算过程可以被高效并行化。这就是为什么 Transformer 能取代 RNN 成为大模型的主流架构。3. 自注意力机制的数学原理自注意力的输入是一个序列的向量表示。假设输入的序列长度为 n每个 token 的向量维度是 d_model那么输入矩阵 X 的形状是[seq_len, d_model]为简化说明先忽略 batch 维度。3.1 生成 Q、K、V首先对输入 X 做三次线性变换得到三个矩阵QQuery查询向量表示“我想找什么”KKey键向量表示“我是什么”VValue值向量表示“我能提供什么内容”计算公式为Q X W_Q K X W_K V X W_V其中 W_Q、W_K、W_V 都是可学习的权重矩阵形状为[d_model, d_k]。这里的 d_k 是每个注意力头中 Q、K、V 的维度。3.2 计算注意力分数注意力分数的计算方式是 Q 和 K 的点积scores Q K^T得到的结果形状是[seq_len, seq_len]表示序列中每个位置对每个位置的匹配程度。第 i 行第 j 列的值越大表示第 i 个 token 越关注第 j 个 token。3.3 缩放直接使用点积结果会有一个问题当 d_k 比较大时点积结果的方差会变大导致 softmax 之后梯度极小训练不稳定。所以需要除以一个缩放因子scores scores / sqrt(d_k)这样做的原因是当两个 d_k 维向量各分量都是独立随机变量均值 0、方差 1时它们的点积的方差约为 d_k。除以 sqrt(d_k) 后方差被拉回 1 左右softmax 的梯度更稳定。3.4 Softmax 归一化对每一行做 softmax让同一位置对所有位置的注意力权重之和为 1weights softmax(scores, dim-1)3.5 加权求和用注意力权重对 V 进行加权求和得到输出output weights V输出形状为[seq_len, d_k]。这一步可以理解为根据“每个位置应该关注谁”把 V 中的信息按权重融合起来。3.6 整体计算图把上面的步骤连起来Attention(Q, K, V) softmax(Q K^T / sqrt(d_k)) V这就是缩放点积注意力Scaled Dot-Product Attention的完整公式。4. 为什么需要缩放和掩码这两个细节是面试高频考点也是实现时最容易出错的地方。4.1 缩放的意义前面已经提到缩放是为了让点积结果的方差稳定避免 softmax 进入饱和区。具体来说如果不缩放当 d_k 很大时点积结果的数值会很大导致 softmax 的输出非常接近 one-hot 分布最大值接近 1其他接近 0。这样梯度过小训练几乎无法进行。除以 sqrt(d_k) 后数值范围被拉回到合理区间。这是 Google 在《Attention Is All You Need》论文中给出的解释实际工程中也验证了缩放的必要性。4.2 Mask 的作用Mask 有两种常见类型Padding Mask 和 Causal Mask。Padding Mask用于处理变长序列。一个 batch 中不同样本的序列长度不同短的序列尾部需要 padding 成同一个长度。padding 的位置没有实际含义在计算注意力权重时应该被屏蔽。实现方式是在 softmax 之前把 padding 位置的分数设为一个非常大的负数如 -1e9这样 softmax 后对应权重趋近于 0。Causal Mask用于自回归生成任务如 GPT 系列。每个位置只能看到当前及之前的位置不能看到未来的 token。这通过在分数矩阵上叠加一个上三角矩阵的负无穷来实现强制未来位置的分数为 -infsoftmax 后权重为 0。两种 mask 可以同时使用。例如在训练 decoder 时既要屏蔽 padding 位置也要屏蔽未来位置。5. 多头注意力机制多头注意力不是做一次注意力而是把 Q、K、V 分成 h 个头每个头独立计算注意力最后把结果拼起来再做一次线性投影。5.1 为什么需要多头单头注意力只能建模一种“关注模式”。但句子中可能同时存在多种依赖关系语法关系、语义关系、指代关系。多头注意力让模型能够同时关注不同维度的信息每个头学到不同的表示子空间。5.2 实现方式假设 d_model 512num_heads 8那么每个头的维度 d_k 512 / 8 64。实现上有两种方式第一种把 W_Q、W_K、W_V 拆成 8 份每份形状为[512, 64]分别计算 8 次注意力最后拼接。第二种也是常用做法一次性生成整个 Q、K、V形状都是[batch_size, seq_len, 512]然后 reshape 成[batch_size, seq_len, num_heads, d_k]再交换维度变成[batch_size, num_heads, seq_len, d_k]直接批量计算注意力。第二种方式在 PyTorch 中更容易实现后面会在代码中演示。5.3 输出投影多头注意力的输出拼接后形状为[batch_size, seq_len, d_model]再经过一个线性层 W_O 进行输出投影output concat(head_1, ..., head_h) W_O6. PyTorch 从零实现自注意力现在我们用 PyTorch 实现一个完整的自注意力模块。不依赖nn.MultiheadAttention从零手写方便看清每一步的维度变化。6.1 单头缩放点积注意力import torch import torch.nn as nn import torch.nn.functional as F def scaled_dot_product_attention(query, key, value, maskNone): 缩放点积注意力。 参数: query: [batch_size, seq_len_q, d_k] key: [batch_size, seq_len_k, d_k] value: [batch_size, seq_len_k, d_v] mask: [batch_size, seq_len_q, seq_len_k] 或可广播版本 返回: output: [batch_size, seq_len_q, d_v] weights: [batch_size, seq_len_q, seq_len_k] d_k query.size(-1) # 1. 计算 Q 和 K 的点积 scores torch.matmul(query, key.transpose(-2, -1)) # [B, 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, -1e9) # 4. Softmax weights F.softmax(scores, dim-1) # 5. 加权求和 output torch.matmul(weights, value) # [B, seq_len_q, d_v] return output, weights6.2 多头自注意力完整实现class MultiHeadSelfAttention(nn.Module): def __init__(self, d_model, num_heads): 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 # 生成 Q、K、V 和输出投影的权重 self.w_q nn.Linear(d_model, d_model) self.w_k nn.Linear(d_model, d_model) self.w_v nn.Linear(d_model, d_model) self.w_o nn.Linear(d_model, d_model) def split_heads(self, x): 输入: [batch_size, seq_len, d_model] 输出: [batch_size, num_heads, seq_len, d_k] batch_size, seq_len, _ x.size() x x.view(batch_size, seq_len, self.num_heads, self.d_k) return x.transpose(1, 2) def combine_heads(self, x): 输入: [batch_size, num_heads, seq_len, d_k] 输出: [batch_size, seq_len, d_model] batch_size, _, seq_len, _ x.size() x x.transpose(1, 2).contiguous() return x.view(batch_size, seq_len, self.d_model) def forward(self, x, maskNone): # x: [batch_size, seq_len, d_model] # 1. 线性变换生成 Q、K、V q self.w_q(x) k self.w_k(x) v self.w_v(x) # 2. 拆分成多头 q self.split_heads(q) # [B, num_heads, seq_len, d_k] k self.split_heads(k) v self.split_heads(v) # 3. 对每个头独立计算注意力 # 注意 mask 需要能广播到 [B, num_heads, seq_len, seq_len] attn_output, weights scaled_dot_product_attention(q, k, v, mask) # 4. 合并多头输出投影 output self.combine_heads(attn_output) output self.w_o(output) return output, weights6.3 使用示例# 测试输入 batch_size 2 seq_len 10 d_model 512 num_heads 8 x torch.randn(batch_size, seq_len, d_model) model MultiHeadSelfAttention(d_model, num_heads) output, weights model(x) print(f输入形状: {x.shape}) print(f输出形状: {output.shape}) print(f注意力权重形状: {weights.shape})输出结果输入形状: torch.Size([2, 10, 512]) 输出形状: torch.Size([2, 10, 512]) 注意力权重形状: torch.Size([2, 8, 10, 10])7. 自注意力与 RNN、CNN 的对比特性RNN / LSTMCNN自注意力并行性差串行计算好好长距离依赖弱信息衰减弱需堆叠层数强任意位置直接关联计算复杂度每步 O(d²)总 O(n·d²)O(k·n·d²)O(n²·d)可解释性一般不直接观察特征图可视化注意力权重可直接可视化对序列长度扩展性较好较好差n² 复杂度从表格可以看到自注意力的最大瓶颈是O(n²) 的序列长度复杂度。当序列长度 n 很大时比如处理长文档或高分辨率图像注意力矩阵的计算和存储开销会迅速膨胀。这也是为什么后续出现了各种优化变体比如稀疏注意力Sparse Attention、局部注意力Local Attention、线性注意力Linear Attention等目的都是降低 n² 复杂度。8. 资源占用与性能观察先看一个常见问题为什么显存经常爆自注意力的中间张量weights的形状是[batch_size, num_heads, seq_len, seq_len]。假设 batch_size4num_heads8seq_len2048这个注意力矩阵就有4 × 8 × 2048 × 2048 1.34 亿个 float32 元素占 512 MB 显存。如果 seq_len 翻倍到 4096显存占用直接变成原来的 4 倍达到 2 GB 以上。观察资源占用可以从这几个维度入手显存占用使用 Nvidia 的nvidia-smi命令实时查看。重点关注训练或推理过程中的显存峰值而不是启动时的静态占用。计算耗时分别用不同 seq_len 跑几次前向观察耗时增长曲线。你会发现 seq_len 从 512 到 1024耗时大约增加 4 倍这和 O(n²) 的复杂度是吻合的。批量大小的影响batch_size 增大会让显存占用接近线性增长但计算单元利用率也更高。如果显存不够优先减小 batch_size而不是减小 seq_len。一个实用的降低显存手段是使用梯度检查点Gradient Checkpointing在反向传播时重新计算前向结果而不是全部保存。以时间换显存适合长序列训练。9. 常见问题与排查方法问题现象可能原因排查方式解决方案训练 loss 不收敛或震荡缩放因子未实现softmax 进入饱和区检查是否除以 sqrt(d_k)加上缩放确认数值范围注意力权重全是同一个分布初始化不当或训练不充分打印 attention weights 观察调整初始化方式增加训练步数序列长度不一致导致报错缺少 Padding Mask查看 batch 中 tensor 形状使用pad_sequence并对 padding 位置做 Maskdecoder 训练时未来信息泄漏缺少 Causal Mask检查 loss 是否异常低叠加上三角 Mask屏蔽未来位置显存溢出OOM序列长度过长或 batch_size 过大查看nvidia-smi显存占用减小 batch_size、梯度检查点、混合精度训练多头 reshape 报错d_model 无法被 num_heads 整除检查d_model % num_heads 0调整 d_model 或 num_heads权重形状不对维度转置或 view 顺序理解错误打印每一步张量形状按照split_heads/combine_heads的流程逐步检查10. 最佳实践与使用建议10.1 从标准实现开始先不用急着上 Flash Attention 等优化版。用 PyTorch 从零写一遍标准自注意力确认每一步维度都正确后再考虑优化。这样对后续读源码、调试模型帮助最大。10.2 注意数值稳定性软注意力的实现里masked_fill使用的负值要够大但不能大到造成数值溢出。常见的做法是使用-1e9或直接用 PyTorch 的float(-inf)配合masked_fill操作。10.3 序列长度和显存要提前预估在开始训练前用一个小规模的输入先跑一次前向和反向记录显存占用。再按公式估算更大序列长度时的显存需求避免训练中途 OOM。10.4 合理选择掩码编码器Encoder只用 Padding Mask。解码器DecoderPadding Mask Causal Mask 同时使用。10.5 可视化注意力权重训练完成后把注意力权重提取出来做可视化可以帮助定位模型学到了哪些依赖关系。如果权重分布接近均匀分布说明模型没有学到有效的信息如果某个位置出现明显的“关注集中”说明该位置在语义上很重要。11. 总结与下一步自注意力机制是 Transformer 的基石理解它的关键点可以归纳为四条第一它通过 Q、K、V 三组向量完成“查询-匹配-提取”的过程让序列中的任意两个位置可以直接建立关联。第二缩放因子 sqrt(d_k) 不是公式修饰而是稳定梯度、保证训练收敛的必要设计。第三Mask 机制解决了两类实际问题变长序列的 padding 处理和自回归模型的未来信息泄漏。第四多头注意力让模型在多个表示子空间中并行建模不同的依赖关系这是 Transformer 表达能力的重要来源。最容易踩的坑有两个一是忽略缩放导致训练不收敛二是多头 reshape 时维度没理清导致报错。建议你按照本文的代码把每一步的中间张量形状打印出来逐一对照很快就能建立手感。下一步你可以去读 Transformer 论文的原文重点看 3.2 节的注意力公式和 3.2.2 节的多头注意力解释然后再读 PyTorch 官方实现的nn.MultiheadAttention源码看官方是怎么做维度变换和 mask 处理的。如果想深入做性能优化可以继续学习 Flash Attention 和 KV Cache 的原理这两块在推理加速中非常关键。