从零手写Transformer：解开多头注意力与位置编码的代码实现

之前有段时间，我觉得自己把Transformer相关的东西学得差不多了：QKV、自注意力、位置编码，一套一套都能讲。结果某天想读一份开源实现，打开文件没看几行就被各种 reshape、transpose 和 mask 绕晕了。后来我换了个思路——不去硬啃别人的完整源码，而是自己动手写一个最简版本。这篇文章就是我用深度学习框架手写一个“能跑版” Transformer 的完整记录，代码量不大，但足够把核心模块串起来。如果你和我一样，属于“论文看懂不少、代码没写过几行”的状态，建议跟着把代码敲一遍，比看十篇解析都管用。

1. 先从整体结构下手：拿到一份Transformer代码先看什么

1.1 先明确这个Transformer是哪种Transformer

Transformer 在原始论文里是 Encoder-Decoder 结构，但后来大家在代码里见到了更多变体：Encoder-only（比如做分类）、Decoder-only（比如做文本生成）。我这次实现的是一个 Decoder-only 的简化版本，因为它包含了 Transformer 最核心的几个组件：多头自注意力、位置编码、前馈网络、LayerNorm、残差连接。你把这套代码看懂以后，想改成 Encoder-only 其实非常简单——把因果 mask 拿掉，让每个位置看到全序列就行；想加 Cross-Attention，也只需要在 TransformerBlock 里再插一个注意力模块。

我对“解析代码”这件事的标准是：必须让我能在笔记本上把模型跑起来，而不是只贴一个永远没办法运行的半成品。所以下面所有代码都是完整可复现的，你把它按顺序拼到一个文件里就能训练。模型很小，但“麻雀虽小五脏俱全”，核心链路一个都不少。

1.2 整条数据流和shape变化

先别急着看代码，闭眼想一遍数据是怎么流的。输入是一个整数序列，假设形状是 (B, T)：

1. 通过 token embedding 查表，变成 (B, T, C)；
2. 加上位置编码，还是 (B, T, C)；
3. 经过 N 个 Transformer Block，每个 Block 内部都是“注意力 + 残差 + LayerNorm + 前馈 + 残差”，输出形状不变；
4. 最后过一个线性层，变成 (B, T, vocab_size)，得到每个位置预测下一个 token 的 logits；
5. 算交叉熵损失。

所有 shape 变化我建议用一张表记下来，后面调试代码时对照着看：

代码里我习惯用 B、T、C 三个字母表示 batch size、序列长度、模型维度。C 就是 d_model。很多开源实现也这么命名，你抓住这套维度关系，后面任何 Transformer 代码都能拆开。

2. 把多头注意力拆成代码：别再背诵QKV公式

2.1 为什么需要多头而不是一个头

在动手写代码之前，先建立一点直觉。单头注意力就是一组 Q、K、V，它只能学一种“谁和谁有关”的模式。多头相当于把模型分成多个子空间，让不同的头分别关注不同类型的关系：有的头关注相邻词，有的头关注跨长距离的指代，有的头关注某种句法结构。最后再把所有头的输出拼起来，让模型拥有更丰富的表达能力。

从公式上看，多头注意力就是缩放点积注意力：

Attention(Q, K, V) = softmax(QK^T / sqrt(d_head)) V

差别在于，Q、K、V 先被拆成 n_heads 份，每份维度是 d_head = d_model / n_heads，各部分独立计算注意力，最后把结果 concat 回去。这也是为什么代码里会有大量 view 和 transpose。

2.2 MultiHeadAttention类的实现与逐行解释

直接看代码。这个类大概是我们整个 Transformer 里最核心的部分，我建议你亲手敲一遍：

import torch import torch.nn as nn import torch.nn.functional as F class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout=0.1): super().__init__() assert d_model % n_heads == 0 self.d_model = d_model self.n_heads = n_heads self.head_dim = d_model // n_heads self.wq = nn.Linear(d_model, d_model, bias=False) self.wk = nn.Linear(d_model, d_model, bias=False) self.wv = nn.Linear(d_model, d_model, bias=False) self.out_proj = nn.Linear(d_model, d_model, bias=False) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): B, T, C = x.shape # 线性投影并拆成多头 q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) k = self.wk(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) v = self.wv(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) # 缩放点积注意力 attn_scores = q @ k.transpose(-2, -1) / (self.head_dim ** 0.5) if mask is not None: attn_scores = attn_scores.masked_fill(mask == 0, float('-inf')) attn_weights = F.softmax(attn_scores, dim=-1) attn_weights = self.dropout(attn_weights) # 加权聚合，然后把多头合并 out = attn_weights @ v out = out.transpose(1, 2).contiguous().view(B, T, C) # 输出投影 return self.out_proj(out)

有几个点必须拆开讲。

第一，为什么线性投影不加 bias？有些实现会加，但后面有 LayerNorm 和残差连接，bias 在多数情况下不是必须的。这里我习惯关闭，减少一点参数，也不会影响收敛。

第二，view 和 transpose 的顺序。x 从 (B, T, C) 变成 (B, T, n_heads, head_dim)，这一步只是拆成 4 维，但此时 head 还在第 2 维，我们需要把它换到第 1 维，让每个 head 独立计算，所以 transpose(1, 2)。最终变成 (B, n_heads, T, head_dim)。所有头在 batch 维之后并行计算，这正是多头注意力的精髓。

第三，缩放因子是 head_dim ** 0.5，不是 d_model ** 0.5。这一点容易被忽略。为什么要缩放？因为 QK^T 的数值会随着维度和长度变大而变大，不缩放的话 softmax 会进入饱和区，梯度变得很小。d_head 决定了每个头上点积的尺度，所以用 head_dim 来缩放。

2.3 因果mask：让Decoder承担自回归任务的关键

Decoder 在生成每个 token 时，绝对不能看到未来的 token，否则就“作弊”了。训练的时候虽然我们可以一次性把整条序列喂进去，但通过一个上三角全为 0 的 mask，让注意力分数在未来的位置变成负无穷，softmax 之后这些位置的概率就变成 0。

因果 mask 的构造非常简单：

def create_causal_mask(T): mask = torch.tril(torch.ones(T, T)).bool() return mask

mask 的形状是 (T, T)，其中 mask[i][j] = True 表示位置 i 可以 attend 到位置 j，False 表示不能。这个 mask 在传入前向函数后会自动广播到 (B, n_heads, T, T) 上，不需要额外扩展。

一个常见的疑问是：为什么训练时不能像推理那样一个个 token 循环生成？因为训练需要高效并行。Transformer 本身是定长序列到定长序列的计算，mask 把它伪装成“因果”的，但内部所有位置的注意力计算还是同时进行的。这也是自回归语言模型训练效率高的原因。

3. 位置编码：代码里最容易忽略却最讲究的一层

3.1 经典sin-cos位置编码的代码落地

Transformer 没有循环结构，也没有卷积，如果直接把 token embedding 输入注意力，模型完全分不清“我喜欢你”和“你喜欢我”在位置上的差别。位置编码就是为了给每个位置一个独特的“指纹”。

经典正弦位置编码的公式是：

PE(pos, 2i) = sin(pos / 10000^(2i/d_model)) PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

代码实现如下：

import math def sinusoidal_position_encoding(max_len, d_model): assert d_model % 2 == 0 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, dtype=torch.float) * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) return pe

代码中有个细节：计算10000^(2i/d_model)时，直接调pow容易出现数值问题，尤其是 i 很大的时候。先取对数再取指数，也就是exp(-log(10000)*2i/d_model)，在数值上更稳定。div_term的 shape 是 d_model/2，分别作用在偶数维和奇数维上。

实际加到输入上时，只需要截取当前序列长度：

x = x + self.position_encoding[:T]

在注册为 buffer 的情况下，position_encoding会跟着模型一起切换设备，不需要手工.to(device)。

3.2 为什么不能直接用数字1、2、3

最容易想到的位置编码是给每个位置加一个标量：位置 0 加 0，位置 1 加 1，位置 2 加 2。这个方案的主要问题有两个。

第一，无界。位置编码的值会随序列长度无限增长，可能掩盖 token 本身的语义信息。一个数值很大的位置向量会“淹没” embedding。

第二，模型很难从无界整数中泛化到更长的序列。sin-cos 编码把位置映射到一个有界区间内，并且相邻位置之间保持平滑的相似性，模型更容易学习相对位置的模式。事实上，sin-cos 编码还具备一些漂亮的线性性质，通过变换可以表达相对位置关系，这是直接加序号做不到的。

3.3 可学习位置编码和正弦位置编码怎么选

在实际代码里，你还会看到另一种做法：把位置编码当作一个nn.Embedding，随机初始化，然后跟着模型训练。

self.position_embedding = nn.Embedding(max_len, d_model)

两种方案各有取舍。正弦编码完全不需要额外参数，有更好的外推能力；可学习编码在训练数据足够多时表现也不错，且实现更简单。近年很多大模型直接用可学习位置编码，因为它们的训练语料极其充足，外推需求不是首要考虑。但如果你是解析经典 Transformer 代码，肯定会遇到 sin-cos 版本，所以这一节务必要吃透。

4. 残差、LayerNorm和前馈网络：Transformer块的骨架

4.1 Pre-LN还是Post-LN：一个顺序决定训练难度

原始论文里的顺序是先做注意力，再加残差，最后 LayerNorm，也就是：

• Post-LN：x = LayerNorm(x + Attention(x))

后来大量实践发现，Post-LN 在深层模型上训练不稳定，往往需要学习率预热（warmup）来稳定训练。现代实现更常见的是 Pre-LN：

• Pre-LN：x = x + Attention(LayerNorm(x))

Pre-LN 把归一化放在前面，残差路径始终保持“干净”的恒等映射，梯度可以直接从最后一层回流到第一层，训练稳定很多，也不需要复杂的预热策略。代码上只是几行的顺序差异，但影响非常大。我们这里的示例使用 Pre-LN，这也是我现在写模块的默认选择。

4.2 FFN的d_ff为什么要膨胀

多头注意力本质上是一个加权求和，权重来自输入之间的相似度，这个过程是线性的。如果堆叠很多层注意力，模型整体表达能力还是受限。前馈网络的作用就是给每个 token 做一次非线性特征变换，通常的做法是先把 d_model 映射到一个更大的中间维度 d_ff，激活后再映射回来。

为什么中间维度要变大？可以把它理解为一个“展开再压缩”的过程。展开阶段把每个位置的局部特征扩展成更高维的表达，激活函数引入非线性，压缩阶段再把高维特征提炼成下一层需要的语义信息。经典配置是 d_ff = 4 * d_model，这个比例在绝大多数代码里都成立。

实现也很简洁：

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

原始论文用的是 ReLU，现在很多实现喜欢 GELU。如果看到老代码里的 ReLU，也不要惊讶，功能上都是非线性激活，只是曲线平滑度不同。

4.3 把模块拼成TransformerBlock

有了注意力、LayerNorm、前馈网络，一个 Transformer Block 就是把他们按 Pre-LN 的顺序粘起来：

class TransformerBlock(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout=0.1): super().__init__() self.attn = MultiHeadAttention(d_model, n_heads, dropout) self.ffn = FeedForward(d_model, d_ff, dropout) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) def forward(self, x, mask=None): x = x + self.attn(self.norm1(x), mask=mask) x = x + self.ffn(self.norm2(x)) return x

这里的 LayerNorm 只对最后一维 C 做归一化，不会跨 batch 或序列长度。这个细节很重要，因为很多刚接触的人会把 LayerNorm 和 BatchNorm 搞混。LayerNorm 是对每个 token 的向量内部做归一化，和序列中其他 token 无关。

残差连接的价值也值得强调。没有残差连接，几十层堆叠之后梯度很容易消失；有了残差，每一层都保证输入信息能直接向后传递。你可以把残差理解成一个“高速公路”，模型可以选择只在这个高速公路上加一点微小的变换，也可以利用它把信息传到很深的层。

5. 从零拼出一个能跑的Transformer：训练小例子

5.1 设计一个能验证模型是否真在学习的任务

理论学习到最后，最重要的就是模型能跑起来。为了快速验证代码的正确性，我找了一个简单到不能再简单的任务：让模型学习循环序列。具体来说，我们用数字 0 到 4 不断重复：

0, 1, 2, 3, 4, 0, 1, 2, 3, 4, ...

然后给模型一段长度为 T 的窗口，让它预测下一个 token。因为规律非常明显，只要模型真的学会了位置编码和注意力，loss 应该快速下降。这个任务不需要大模型，也不需要大规模训练数据，非常适合验证 Transformer 代码链路。

数据生成代码不需要写 DataLoader 类，几行函数就够了：

def random_batch(seq, batch_size, T): starts = torch.randint(0, len(seq) - T - 1, (batch_size,)) xs, ys = [], [] for s in starts: xs.append(seq[s:s + T - 1]) ys.append(seq[s + 1:s + T]) return torch.stack(xs), torch.stack(ys)

5.2 模型结构参数怎么定

因为任务很简单，我用非常小的模型。选择参数要遵守一个原则：尽量小到 CPU 能跑，但大到足够体现多头和深度的作用。下面这组参数是我试过比较舒服的配置：

• vocab_size = 5（只有 0 到 4 五个数字）
• d_model = 32
• n_heads = 4
• head_dim = 8
• d_ff = 64
• num_layers = 2
• block_size = 64（最大序列长度）

注意 d_model 必须能被 n_heads 整除，否则注意力模块的 assert 会直接报错。d_model 选 32 而 head_dim 是 8，这个维度对 CPU 上的小任务很友好，同时还能展示多头并行。

完整模型类：

class Transformer(nn.Module): def __init__(self, vocab_size, d_model, n_heads, d_ff, num_layers, max_len=64, dropout=0.1): super().__init__() self.token_embedding = nn.Embedding(vocab_size, d_model) self.position_encoding = sinusoidal_position_encoding(max_len, d_model) self.register_buffer('pe', self.position_encoding) self.blocks = nn.ModuleList([ TransformerBlock(d_model, n_heads, d_ff, dropout) for _ in range(num_layers) ]) self.norm_final = nn.LayerNorm(d_model) self.lm_head = nn.Linear(d_model, vocab_size) def forward(self, tokens): B, T = tokens.shape x = self.token_embedding(tokens) x = x + self.pe[:T] mask = torch.tril(torch.ones(T, T, dtype=torch.bool, device=tokens.device)) for block in self.blocks: x = block(x, mask) x = self.norm_final(x) logits = self.lm_head(x) return logits

位置编码在__init__里通过register_buffer注册，这样模型调用.to(device)或.cuda()时，位置编码也会自动跟着移动，省去手动转移设备的麻烦。

5.3 训练循环和loss观察

训练循环非常常规。优化器用 Adam，学习率设 1e-3，损失函数是交叉熵：

model = Transformer(vocab_size=5, d_model=32, n_heads=4, d_ff=64, num_layers=2, max_len=64) opt = torch.optim.Adam(model.parameters(), lr=1e-3) crit = nn.CrossEntropyLoss() pattern = torch.arange(5).repeat(100) for step in range(500): x, y = random_batch(pattern, batch_size=32, T=8) opt.zero_grad() logits = model(x) loss = crit(logits.transpose(1, 2), y) loss.backward() opt.step() if step % 100 == 0: print(f"step {step}: loss {loss.item():.4f}")

运行之后你会发现，loss 一开始大约在 1.6 附近，因为 5 个类别的均匀分布熵就是 log5≈1.609。这意味着模型初期完全没有学习。几百步内 loss 会明显下降，最后降到很低，说明模型已经掌握了循环规律。

如果 loss 一直不下降，别怀疑任务，先回头查上一节的那些 shape 和 mask 细节。

5.4 生成：让模型自己预测下一个token

训练完之后，可以写一个最简单的贪心生成函数，看看模型能不能自己续写：

def generate(model, start_tokens, max_new=20, block_size=64): model.eval() tokens = list(start_tokens) with torch.no_grad(): for _ in range(max_new): inp = torch.tensor([tokens[-block_size:]], dtype=torch.long) logits = model(inp) next_id = logits[0, -1].argmax().item() tokens.append(next_id) return tokens

这段代码每次把当前已有的 token 序列重新送进模型，只取最后一个位置的预测结果，然后拼回去。这个“重新前向计算整个序列”的做法，也是很多生成模型推理速度慢的原因，因为每生成一个 token 都要全量计算一遍。gpt 类模型生成新词时能缓存历史的 K/V 从而加速，但这是优化层面的问题，不影响理解。

6. 调试Transformer代码时踩过的坑：shape、mask、NaN与顺序

6.1 shape不匹配：先画表再写代码

我写这份代码时犯过最多的错误就是 shape 问题。最常见的是view和transpose的顺序搞反。很多新手会把x.view(B, n_heads, T, head_dim)直接写出来，但这是错的，因为原来的最后一个维度是 C，要先把 C 拆成 (n_heads, head_dim)，然后再把 n_heads 换到第 1 维。

我的调试习惯是：在关键位置打印 shape。比如注意力模块里，打印q.shape、k.shape、attn_scores.shape和out.shape，一行行对表。写代码不是一次就能跑通的，打印调试是最快的路径。

6.2 mask的广播一点要小心

mask 的形状是 (T, T)，但注意力分数是 (B, n_heads, T, T)。很多新手会觉得必须手动把 mask 扩展成 (B, n_heads, T, T) 才能用，实际上masked_fill会自动广播，不需要额外处理。真正容易翻车的是 mask 的逻辑方向。

我习惯用 bool 类型的 mask，True 表示可见、False 表示不可见。这样在代码里写masked_fill(mask == 0, float('-inf'))非常直观。如果你用的是 float mask，0 和 1 的含义就要格外小心，很容易弄反。还有一个隐蔽问题：如果某个位置的所有注意力分数都被 mask 成-inf，softmax 之后可能出现 NaN。正常情况下因果 mask 不会出现整行都被 mask 的情况，但如果你在建 mask 时方向写反，就很有可能。

6.3 loss不降和NaN的排查顺序

遇到 loss 不降，我一般按这个顺序排查：

1. 检查 label 是否右移正确。自回归模型里，输入是前 T-1 个 token，目标必须是后 T-1 个 token。
2. 检查 mask 是否让每个位置都能至少看到一个 token。
3. 检查多头合并时有没有用 contiguous 再 view。把(B, n_heads, T, head_dim)转回(B, T, C)时，先transpose再view会报错，必须先.contiguous()。
4. 检查激活函数和初始化。小模型一般不需要特殊初始化，但如果 loss 变成 NaN，最常见原因是学习率太大。加一行torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)能避免大部分梯度爆炸。

6.4 结构性建议：先跑最小case再放大

最后分享一个我自己一直在用的调试技巧：先跑 batch_size=1、序列长度=1，确认 forward 能过；再跑 batch_size=2、序列长度=3，确认多头和 mask 没有逻辑错误；最后才用完整配置训练。这样做可以把问题从“模型结构错误”和“训练参数问题”中快速分开。

还有一个小经验：写 Transformer 模块时，不要把 LayerNorm 放在forward函数里临时创建，而要在__init__里注册为子模块。否则模型的参数无法被优化器管理，训练时会出现“模型 loss 下降但参数不更新”的诡异现象。听起来很蠢，但我真的见过有人把 norm 写成每次 forward 都新实例化，然后怎么调都不收敛。代码这个东西，跑通了并不难，难的是把每个“为什么”都想明白。

如果你按这篇文章把代码敲下来，并且真的让 loss 下降了，我相信你再去读那些大而全的开源实现，至少不会一看到 reshape 就发怵。写 Transformer 代码没有捷径，能做的就是把每一步 shape 和每一块 mask 都亲手走一遍。