从源码级拆解Transformer：自注意力、多头与位置编码的PyTorch实现-平芜编程栈

Transformer 从 2017 年那篇《Attention Is All You Need》开始，就成了深度学习领域绕不开的底座。但说实话，我见过太多人学 Transformer 的路径是：看几篇公众号图解，背下“自注意力”“多头”“位置编码”这几个词，然后就去调库跑模型了。真到要改结构、要排查训练不收敛、要跟人解释为什么你的模型比别人的慢一倍时，脑子里那点图解决不够用了。这篇东西就是写给这类人的——不满足于“知道”，想真正“看懂”源码级实现的人。我会从张量的实际形状出发，把自注意力、多头、前馈、残差、层归一化、位置编码这些模块逐个拆开，配上可运行的 PyTorch 代码，讲清楚每一步为什么这么设计、不这么写会出什么问题。适合已经会 PyTorch 基础、跑过一两个 NLP 任务、但想彻底吃透 Transformer 内部机制的读者。读完你应该能自己从零手写一个可训练的 Transformer，并且知道每一行代码在干什么。

1. 为什么光看图解永远学不会 Transformer

1.1 图解给你的是一张“地图”，源码给你的才是“地形”

市面上大多数 Transformer 图解，画的是这样的流程：输入 → 位置编码 → 编码器堆叠 → 解码器堆叠 → 输出。箭头清晰，模块分明，看着很舒服。但问题在于，图解把张量形状的变化、维度的对齐、mask 的广播方式这些真正决定你能不能跑通的东西全藏起来了。你照着图去写代码，写到torch.matmul(Q, K.transpose(-2, -1))这一步，突然发现 Q 和 K 的维度对不上，或者 softmax 之后数值全变成 0，这时候图帮不了你。

我自己的经验是，学 Transformer 最有效的方式是：先看一遍图建立整体印象，然后立刻去读一份干净的源码实现，逐行对照张量形状。图告诉你“有自注意力这回事”，源码告诉你“Q 是 [batch, heads, seq_len, d_k]，K 转置后是 [batch, heads, d_k, seq_len]，乘完是 [batch, heads, seq_len, seq_len]”。后者才是你能拿来干活的知识。

1.2 一个真实的踩坑：为什么我的 attention 输出全是常数

早期我自己实现的时候，遇到过 attention 输出对所有位置几乎一样的情况。排查了半天，发现是缩放因子写错了——我把1/sqrt(d_k)写成了sqrt(d_k)。这个错误在图解里根本看不出来，因为图解通常只写一个scale的框，不写具体数值。但源码里d_k是多少、开根号后是多少、乘在 softmax 之前还是之后，每一个细节都会影响数值稳定性。

具体来说，当d_k = 64时，sqrt(d_k) = 8。如果 Q 和 K 的点积结果方差大约为d_k（因为每个维度是独立同分布、均值 0 方差 1 的话，点积的方差就是 d_k），那么点积值可能落在[-30, 30]这个范围。softmax 对这么大的输入会变得极其尖锐，几乎变成 one-hot，梯度接近 0。除以 8 之后，范围缩到[-4, 4]，softmax 才有合理的梯度。这就是为什么缩放必须做，而且必须除以sqrt(d_k)而不是别的。

提示：如果你在实现时发现 attention 权重图看起来像一条竖线（某个位置权重接近 1，其余接近 0），先检查缩放因子，再检查 mask 是不是写反了。

1.3 源码阅读的正确姿势：从形状出发，而不是从公式出发

公式Attention(Q,K,V) = softmax(QK^T / sqrt(d_k))V很简洁，但它省略了 batch 维、head 维、mask 的加法和 dropout 的位置。我建议读源码时养成一个习惯：每读一行，就在纸上写出当前张量的形状。比如：

输入x:[batch_size, seq_len, d_model]
经过 Q 线性层后Q:[batch_size, seq_len, d_model]
拆成多头后Q:[batch_size, num_heads, seq_len, d_k]，其中d_k = d_model / num_heads
Q @ K^T:[batch_size, num_heads, seq_len, seq_len]
加 mask、softmax、dropout 后形状不变
再乘 V:[batch_size, num_heads, seq_len, d_k]
合并多头:[batch_size, seq_len, d_model]

把这条形状链背下来，比背公式有用十倍。因为你在调试时，报错信息里全是形状，你脑子里有这条链，一眼就能定位是哪一步错了。

2. 自注意力模块的源码级拆解

2.1 Q、K、V 到底是怎么来的：三个线性层的真实作用

很多人以为 Q、K、V 是输入 x 的某种“变换”，但说不清为什么要三个不同的线性层。从源码角度看，最直接的解释是：如果 Q、K、V 都用同一个 x，那么 attention 就退化成了一种基于自身相似度的加权平均，表达能力受限。用三个独立的线性层，相当于让模型学习三种不同的“视角”去投影输入：Q 负责“我在找什么”，K 负责“我有什么可以被找到”，V 负责“我实际提供什么内容”。

在 PyTorch 里，最常见的写法是用一个大的线性层一次性算出 Q、K、V，然后 split：

class SelfAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() self.d_model = d_model self.num_heads = num_heads self.d_k = d_model // num_heads assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除" self.qkv_proj = nn.Linear(d_model, 3 * d_model) self.out_proj = nn.Linear(d_model, d_model) def forward(self, x, mask=None): batch_size, seq_len, _ = x.shape qkv = self.qkv_proj(x) # [B, L, 3*d_model] qkv = qkv.reshape(batch_size, seq_len, 3, self.num_heads, self.d_k) qkv = qkv.permute(2, 0, 3, 1, 4) # [3, B, H, L, d_k] q, k, v = qkv[0], qkv[1], qkv[2] # 后续计算...

这里有个细节值得说：为什么用reshape再permute，而不是直接切分？因为nn.Linear输出的最后一维是3*d_model，它把 Q、K、V 拼在一起了。reshape成[B, L, 3, H, d_k]之后，第 2 维就是 Q/K/V 的索引，permute把它挪到最前面，方便后面用qkv[0]这种写法。这个顺序不能乱，否则 Q、K、V 会错位。

2.2 缩放点积：为什么是除以 sqrt(d_k) 而不是 d_k

前面提过缩放因子的问题，这里从数值角度再展开一下。假设 Q 和 K 的每个分量都是独立同分布，均值 0、方差 1。那么点积q · k = sum(q_i * k_i)的方差是d_k，标准差是sqrt(d_k)。如果不缩放，点积值的标准差随d_k增大而增大，softmax 的输入范围就越来越大，输出越来越尖锐，梯度越来越小。

除以sqrt(d_k)之后，点积的方差变成 1，标准差变成 1，softmax 的输入范围稳定在合理区间。这就是为什么是sqrt(d_k)而不是d_k：因为方差是d_k，要把它拉回 1，需要除以标准差sqrt(d_k)。

代码上就是一行：

scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)

注意transpose(-2, -1)而不是transpose(1, 2)，因为 q 的形状是[B, H, L, d_k]，最后两维才是L和d_k，转置最后两维得到[B, H, d_k, L]，才能和 q 做矩阵乘法。

2.3 mask 的两种类型：padding mask 和 causal mask 的广播机制

mask 是 Transformer 实现里最容易出错的地方之一。常见的有两种：

padding mask：用来忽略 padding 位置，形状通常是[B, L]，1 表示有效，0 表示 padding。
causal mask（也叫 look-ahead mask）：用来防止解码器看到未来位置，形状是[L, L]的下三角矩阵。

在源码里，这两种 mask 经常需要组合使用。组合方式通常是先扩展维度再相加或相乘。比如：

# padding_mask: [B, L] -> [B, 1, 1, L] # causal_mask: [L, L] -> [1, 1, L, L] # 组合后: [B, 1, L, L]，广播到 [B, H, L, L] combined_mask = padding_mask[:, None, None, :] & causal_mask[None, None, :, :]

然后在 softmax 之前，把 mask 为 0 的位置填成负无穷：

scores = scores.masked_fill(combined_mask == 0, float('-inf')) attn = torch.softmax(scores, dim=-1)

这里有个坑：masked_fill填的是-inf，不是 0。如果填 0，softmax 之后这些位置会得到非零权重，模型就会“看到”不该看的位置。填-inf之后，softmax 输出严格为 0。

注意：如果你用混合精度训练（fp16），-inf可能会导致 NaN。这时候可以用一个很大的负数代替，比如-1e9。但要注意，如果所有位置都被 mask 掉，softmax 会输出 NaN，所以至少要保证每行有一个有效位置。

2.4 多头注意力的 reshape 与 transpose：维度顺序决定成败

多头注意力的核心思想是：把d_model维拆成num_heads个d_k维，每个头独立做 attention，最后拼回来。但“拆”和“拼”的顺序非常讲究。

假设d_model = 512，num_heads = 8，d_k = 64。输入x形状[B, L, 512]。经过 Q 线性层后还是[B, L, 512]。现在要拆成 8 个头：

q = q.reshape(batch_size, seq_len, num_heads, d_k) # [B, L, 8, 64] q = q.permute(0, 2, 1, 3) # [B, 8, L, 64]

注意reshape之后是[B, L, H, d_k]，permute把 H 挪到 L 前面。为什么不能直接reshape(B, H, L, d_k)？因为reshape是按内存顺序重新解释数据，[B, L, 512]的内存布局是L维连续，直接 reshape 成[B, H, L, d_k]会把原本属于不同位置的维度混在一起。必须先 reshape 成[B, L, H, d_k]，再 permute，才能保证每个头拿到的是正确的子空间。

合并的时候反过来：

context = context.permute(0, 2, 1, 3) # [B, L, H, d_k] context = context.reshape(batch_size, seq_len, d_model) # [B, L, 512]

这个顺序错了，模型可能也能训，但学到的表示是乱的，效果会明显变差。我见过有人直接reshape(B, L, d_model)而不 permute，结果训练 loss 降得很慢，排查很久才发现是这里的问题。

3. 前馈网络、残差与层归一化的工程细节

3.1 前馈网络为什么先升维再降维

Transformer 里的前馈网络（FFN）通常是两层线性层加一个激活函数：

class FeedForward(nn.Module): def __init__(self, d_model, d_ff, dropout=0.1): super().__init__() self.linear1 = nn.Linear(d_model, d_ff) self.linear2 = nn.Linear(d_ff, d_model) self.dropout = nn.Dropout(dropout) self.activation = nn.GELU() def forward(self, x): return self.linear2(self.dropout(self.activation(self.linear1(x))))

d_ff通常是d_model的 4 倍，比如d_model=512时d_ff=2048。为什么要先升维再降维？直观理解是：升维提供了一个更高维的空间，让模型能在这个空间里做非线性变换，捕捉更复杂的特征交互；降维再把它压回原维度，方便和残差连接相加。如果直接在同维度做非线性，表达能力会受限。

另一个细节是激活函数的选择。原始论文用的是 ReLU，但后来很多实现改用 GELU。GELU 在 0 附近更平滑，梯度性质更好，实测在 Transformer 里通常比 ReLU 略好。不过这个差异不是决定性的，ReLU 也能用。

3.2 残差连接的位置：Pre-LN 和 Post-LN 的区别

原始论文用的是 Post-LN：LayerNorm(x + Sublayer(x))。但后来大家发现 Post-LN 在深层 Transformer 里训练不稳定，需要 warmup 很久。于是 Pre-LN 变得流行：x + Sublayer(LayerNorm(x))。

两者的区别在源码里就是一行顺序：

# Post-LN x = LayerNorm(x + Sublayer(x)) # Pre-LN x = x + Sublayer(LayerNorm(x))

Pre-LN 的好处是梯度可以直接通过残差路径回传，不需要经过 LayerNorm，训练更稳定，对 warmup 的依赖更小。现在大多数开源实现（包括很多大模型）都用 Pre-LN。如果你自己从零写，建议直接用 Pre-LN，省去调 warmup 的麻烦。

提示：Pre-LN 的最后一层通常还需要一个额外的 LayerNorm，因为残差路径上的值没有被归一化过。这个细节容易被忽略，但不加的话输出分布可能偏大。

3.3 LayerNorm 和 BatchNorm 在序列任务上的取舍

LayerNorm 是对每个样本的最后一维做归一化，BatchNorm 是对整个 batch 的同一维度做归一化。在序列任务里，序列长度可变，BatchNorm 的统计量会随 batch 内序列长度变化而波动，不稳定。LayerNorm 对每个样本独立归一化，不受 batch 组成影响，更适合序列任务。

源码上，LayerNorm 就是：

class LayerNorm(nn.Module): def __init__(self, d_model, eps=1e-6): super().__init__() self.gamma = nn.Parameter(torch.ones(d_model)) self.beta = nn.Parameter(torch.zeros(d_model)) self.eps = eps def forward(self, x): mean = x.mean(-1, keepdim=True) std = x.std(-1, keepdim=True, unbiased=False) return self.gamma * (x - mean) / (std + self.eps) + self.beta

注意unbiased=False，因为这里算的是整个维度的统计量，不是采样估计，用有偏方差更合适。另外eps放在分母里而不是外面，是为了防止std为 0 时除零。

4. 位置编码：从正弦到可学习，再到旋转

4.1 正弦位置编码的公式推导与实现

Transformer 本身没有循环结构，也没有卷积结构，如果不加位置编码，它对输入顺序是完全不敏感的——把序列打乱，输出不变。所以必须显式注入位置信息。原始论文用的是正弦位置编码：

class PositionalEncoding(nn.Module): def __init__(self, d_model, max_len=5000): super().__init__() pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) self.register_buffer('pe', pe) def forward(self, x): return x + self.pe[:x.size(1)]

这里div_term是10000^(-2i/d_model)，用exp和log实现是为了数值稳定。偶数维用 sin，奇数维用 cos。这样设计的好处是：对于任意固定偏移k，PE(pos+k)可以表示为PE(pos)的线性变换，模型容易学到相对位置关系。

4.2 可学习位置编码的适用场景与局限

可学习位置编码就是用一个nn.Embedding(max_len, d_model)，把位置索引映射成向量，和词嵌入相加。BERT 用的就是这种。它的好处是实现简单，能自适应数据；局限是最大长度固定，超过max_len就没法处理，而且不能外推到训练时没见过的长度。

如果你的任务序列长度固定或变化不大，可学习位置编码够用。但如果需要处理变长序列，或者希望模型能外推到更长序列，正弦编码或旋转位置编码（RoPE）更合适。

4.3 旋转位置编码（RoPE）的核心思想

RoPE 是近几年比较流行的位置编码方式，核心思想是把位置信息编码成旋转矩阵，作用在 Q 和 K 上。具体来说，对于位置m的向量q，把它按二维一组拆分，每组的两个分量看成复平面上的一个点，乘以e^{i m theta}，相当于旋转了一个角度。这样两个位置的内积只依赖于它们的相对距离，天然具备相对位置编码的性质。

RoPE 的实现稍微复杂一些，但好处是不增加额外参数，而且外推性比正弦编码好。很多大模型都用它。如果你要实现，关键是把d_k维分成d_k/2组，每组用不同的旋转频率。

5. 编码器与解码器的堆叠逻辑

5.1 编码器层：自注意力加前馈的串联

一个编码器层就是：自注意力 → 残差 + LayerNorm → 前馈 → 残差 + LayerNorm。用 Pre-LN 写出来：

class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout=0.1): super().__init__() self.self_attn = SelfAttention(d_model, num_heads) self.ffn = FeedForward(d_model, d_ff, dropout) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) def forward(self, x, mask=None): x = x + self.dropout1(self.self_attn(self.norm1(x), mask)) x = x + self.dropout2(self.ffn(self.norm2(x))) return x

注意 dropout 的位置：在子层输出之后、残差相加之前。这个顺序不能反，否则 dropout 会作用在残差和上，影响梯度回传。

5.2 解码器层：masked 自注意力与交叉注意力的配合

解码器层比编码器多了一个交叉注意力子层，用来让解码器关注编码器的输出。结构是：masked 自注意力 → 残差 + LayerNorm → 交叉注意力 → 残差 + LayerNorm → 前馈 → 残差 + LayerNorm。

class DecoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout=0.1): super().__init__() self.self_attn = SelfAttention(d_model, num_heads) self.cross_attn = CrossAttention(d_model, num_heads) self.ffn = FeedForward(d_model, d_ff, dropout) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.norm3 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.dropout3 = nn.Dropout(dropout) def forward(self, x, enc_out, self_mask=None, cross_mask=None): x = x + self.dropout1(self.self_attn(self.norm1(x), self_mask)) x = x + self.dropout2(self.cross_attn(self.norm2(x), enc_out, cross_mask)) x = x + self.dropout3(self.ffn(self.norm3(x))) return x

交叉注意力的 Q 来自解码器，K 和 V 来自编码器输出。mask 通常用编码器的 padding mask，形状需要广播到[B, H, L_dec, L_enc]。

5.3 堆叠层数、头数与维度的经验取值

原始论文的 base 模型是 6 层编码器、6 层解码器、8 个头、d_model=512、d_ff=2048。大模型通常更深、更宽。经验上：

d_model和num_heads要匹配，d_k = d_model / num_heads通常取 64 或 128。
层数增加能提升表达能力，但训练难度也增加，需要配合 Pre-LN 和合适的初始化。
d_ff一般是d_model的 4 倍，但有些实现用 2 倍或 8 倍，取决于任务。

如果你资源有限，优先增加d_model而不是层数，因为宽模型的并行效率更高。

6. 从零跑通一个 Transformer 的实操路径

6.1 数据准备与词表构建的最小闭环

先准备一个简单的序列到序列任务，比如把数字序列反转。数据量不用大，几百条就够验证代码正确性。词表用最简单的字符级或数字级映射：

vocab = {str(i): i for i in range(10)} vocab['<pad>'] = 10 vocab['<sos>'] = 11 vocab['<eos>'] = 12

然后写一个 Dataset，把输入序列和输出序列都转成索引，padding 到固定长度。注意输出序列要右移一位作为解码器输入，这是 teacher forcing 的标准做法。

6.2 训练循环中的关键细节：warmup、label smoothing、梯度裁剪

训练 Transformer 有几个细节直接影响收敛：

warmup：学习率从 0 线性增加到峰值，再按1/sqrt(step)衰减。Post-LN 对 warmup 很敏感，Pre-LN 可以不用，但用了通常更稳。
label smoothing：把目标分布从 one-hot 改成(1-eps)的 one-hot 加eps/V的均匀分布，eps通常取 0.1。能防止模型过度自信，提升泛化。
梯度裁剪：把梯度范数限制在 1.0 或 5.0，防止梯度爆炸。

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

这三个细节在论文里可能只是一句话，但在实操里是能不能训起来的关键。

6.3 验证模型是否学对的几个信号

跑通之后，怎么知道模型真的学对了？看几个信号：

loss 是否稳定下降，而不是震荡或卡住。
训练集上的准确率是否接近 100%（对于简单任务）。
生成的序列是否符合预期，比如反转任务里输出是不是输入的反转。
attention 权重图是否有合理的模式，比如解码器在生成某个位置时是否关注到编码器对应位置。

如果 loss 不降，先检查 mask 和缩放因子；如果 loss 降但生成结果乱，检查位置编码和 teacher forcing 的移位是否正确。

7. 几个容易翻车的实现细节

7.1 dropout 加错位置导致训练和推理不一致

dropout 必须加在子层输出之后、残差相加之前，而且推理时要model.eval()关掉。如果加在残差之后，训练时残差路径被随机置零，梯度回传受影响；如果推理时忘了 eval，输出会随机波动，看起来像模型没训好。

7.2 权重初始化：为什么不能全用默认初始化

PyTorch 的nn.Linear默认用 Kaiming 均匀初始化，对 Transformer 来说偏大。原始论文用的是 Xavier 均匀初始化，并且对残差路径上的层做了缩放。如果全用默认初始化，深层模型可能一开始就梯度爆炸或消失。建议手动初始化：

def init_weights(m): if isinstance(m, nn.Linear): nn.init.xavier_uniform_(m.weight) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.LayerNorm): nn.init.ones_(m.weight) nn.init.zeros_(m.bias)

7.3 学习率调度：不是所有任务都需要 warmup

warmup 对 Post-LN 是必需的，对 Pre-LN 是可选的。如果你的任务小、数据少，warmup 可能反而拖慢收敛。可以先不加 warmup 跑一下，如果 loss 前期震荡再考虑加。另外，峰值学习率不要设太大，Transformer 通常用 1e-4 到 5e-4 这个量级，比 CNN 小。

7.4 序列长度与显存的关系：attention 的 O(L^2) 复杂度

自注意力的计算复杂度是O(L^2 * d_model)，显存占用也是O(L^2 * num_heads)。序列长度翻倍，显存翻四倍。如果显存不够，优先减小 batch size 或序列长度，而不是减模型维度。也可以用梯度累积来模拟大 batch。

8. 从理解到改造：下一步可以做什么

把上面这些跑通之后，你就有了一个可以随意改造的 Transformer。接下来可以尝试：

把自注意力换成稀疏注意力或线性注意力，看能不能在长序列上省显存。
把正弦位置编码换成 RoPE，对比外推能力。
把 Pre-LN 换成 Post-LN，观察训练稳定性的差异。
加一个 adapter 层，做参数高效微调。

我自己在实际操作中的体会是，Transformer 的很多设计选择不是“最优解”，而是“在当时的计算条件下够用的解”。理解了这一点，你就不会把它当成黑盒，而是当成一个可以按需调整的工具。源码读多了会发现，真正决定模型能不能跑起来的，往往不是那些高大上的概念，而是 mask 有没有写对、缩放因子有没有漏、dropout 位置有没有放错这些细节。把这些细节吃透，比多背几个公式有用得多。

从源码级拆解Transformer：自注意力、多头与位置编码的PyTorch实现