用PyTorch从零手写Transformer：多头注意力、mask与编解码器实现详解-平芜编程栈

Transformer 是很多算法工程师和研究生绕不开的一个模型。看了大量讲解图，收藏了不少经典文章，但真到了自己动手实现时，往往会在三个地方卡住：多头注意力的张量维度怎么组织、mask 矩阵如何广播、编码器和解码器之间到底怎么传递数据。这篇文章就用 PyTorch 从零手写一个完整可运行的 Transformer，包含位置编码、多头注意力、层归一化、前馈网络、编码器和解码器，每个类都带详细注释，最后在一个倒序复制任务上训练验证。如果你对 Transformer 已经有概念，但一直没真正敲过代码，这篇教程就是为你准备的。

1. Transformer 背景与核心概念

1.1 为什么需要 Transformer

在 Transformer 出现之前，序列建模任务主要依赖 RNN、LSTM 和 GRU。这类循环神经网络的核心问题是：当前时刻的输出依赖上一时刻的隐状态，天然只能顺序计算。这带来两个问题：一是长序列中信息衰减严重，很难捕捉距离很远的依赖关系；二是并行性很差，训练效率低。

Transformer 直接放弃了循环结构，提出一种基于注意力机制的网络架构。它允许每个位置的输出同时关注整个输入序列中的任意位置，全局依赖建模能力更强，而且不同位置的计算可以并行完成。这也是后来 BERT、GPT 以及各种多模态大模型都基于 Transformer 演进的原因。

1.2 从人的注意力到自注意力机制

先来一个通俗的理解。你在看这张图时，眼睛会优先关注信息量最大的区域，而不是把整张图从上到下每个像素都均匀扫一遍。注意力机制本质上就是让模型学会“对重要信息赋予更高权重”。

Transformer 使用自注意力机制，核心是三个名词：查询 Query、键 Key、值 Value。

Query 表示“我想找什么信息”。
Key 表示“我有哪些信息可以匹配”。
Value 表示“真正被提取出来的内容”。

计算过程可以分成四步：

输入序列通过三个线性层，分别得到 Q、K、V。
用 Query 和所有 Key 做点积，得到注意力分数，表示当前位置应该关注哪些位置。
除以缩放因子，再经过 softmax 归一化成概率。
用概率对 Value 做加权求和，得到注意力输出。

为什么要除以根号 d_k？因为当维度较大时，点积结果方差会变大，softmax 输出的分布会过于尖锐，梯度容易消失。缩放一下可以让训练更稳定。

1.3 编解码器框架

Transformer 最初是机器翻译模型，整体结构分为编码器和解码器。

编码器负责读取源序列。它内部使用自注意力，也就是每个位置都可以关注源序列中的所有位置。比如处理英文句子 “I love you” 时，“love” 可以通过注意力机制和 “I”“you” 建立联系。

解码器负责生成目标序列。它内部有两层注意力：

掩码自注意力：目标序列当前位置只能看到当前位置以及它之前的位置，不能看到未来信息，否则就是作弊。
交叉注意力：解码器每个位置都会关注编码器输出的整个源序列表示，这样生成目标时才能对齐输入信息。

整体训练过程可以理解为：编码器把源句子编码成一个“语义记忆”，解码器在这个记忆的指导下逐步生成目标句子。

2. 环境准备与版本说明

本文所有代码基于 PyTorch 实现，建议准备一个干净的 Python 环境。版本不同影响不大，主要 API 都是稳定的，下面以常见环境为例。

python -m venv transformer_env source transformer_env/bin/activate pip install torch numpy

创建环境后，可以先验证 PyTorch 是否安装成功：

import torch print(torch.__version__) print(torch.cuda.is_available())

如果输出一长串版本号，并且torch.cuda.is_available()为 True，说明 GPU 可用。没有 GPU 也没关系，本文的演示任务很小，使用 CPU 也可以训练。

建议 IDE 使用 PyCharm 或者 VSCode，新建一个项目文件夹，结构如下：

transformer_from_scratch/ ├── model.py ├── train.py └── README.md

其中model.py放模型代码，train.py放训练代码。后续所有代码都按照这个结构组织。

3. 逐行手写多头注意力机制

3.1 Q、K、V 线性投影

多头注意力是整个 Transformer 最核心的部分。所谓“多头”，就是不要只让模型用一种注意力模式，而是用多组 Q、K、V 并行计算，每组关注不同的子空间信息。

先定义 MultiHeadAttention 类：

import math import torch import torch.nn as nn class MultiHeadAttention(nn.Module): def __init__(self, d_model, n_heads, dropout=0.1): super().__init__() assert d_model % n_heads == 0, "d_model 必须能被 n_heads 整除" self.d_model = d_model self.n_heads = n_heads self.head_dim = d_model // n_heads 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.out_proj = nn.Linear(d_model, d_model) self.dropout = nn.Dropout(dropout) def forward(self, q, k, v, mask=None): batch_size, q_len, _ = q.size() k_len = k.size(1) Q = self.w_q(q).view(batch_size, q_len, self.n_heads, self.head_dim) K = self.w_k(k).view(batch_size, k_len, self.n_heads, self.head_dim) V = self.w_v(v).view(batch_size, k_len, self.n_heads, self.head_dim) Q = Q.transpose(1, 2) K = K.transpose(1, 2) V = V.transpose(1, 2) scores = Q @ K.transpose(-2, -1) / math.sqrt(self.head_dim) if mask is not None: scores = scores.masked_fill(mask == 0, float("-inf")) attn = torch.softmax(scores, dim=-1) attn = self.dropout(attn) out = attn @ V out = out.transpose(1, 2).contiguous() out = out.view(batch_size, q_len, self.d_model) return self.out_proj(out)

这里需要重点理解张量维度的变化。

输入 q 的形状是[batch_size, seq_len, d_model]。经过线性层后，用view把它拆成[batch_size, seq_len, n_heads, head_dim]，然后再用transpose(1, 2)把 n_heads 提到第二个维度，变成[batch_size, n_heads, seq_len, head_dim]。

为什么要 transpose 到第二个维度？因为 PyTorch 的 batch 维度在第一位，记忆和参数都在后面。我们把每个头独立成一组矩阵，方便对每个头分别做点积和 softmax。

注意力分数的计算是Q @ K.transpose(-2, -1)，也就是说[batch, heads, q_len, head_dim]与[batch, heads, head_dim, k_len]做矩阵乘法，结果形状是[batch, heads, q_len, k_len]。这正是每个 query 和每个 key 的相似度矩阵。

mask 的作用后面会详细讲。最后把多头的输出重新拼接回[batch, q_len, d_model]，再经过一个输出线性层。

这里顺带说明：PyTorch 官方也提供了nn.MultiheadAttention，开箱即用，但它的参数batch_first可能让初学者困惑，而且我们手写一遍能更清楚地理解内部流程。等手写版本跑通后，再切换成官方 API 就很简单了。

3.2 位置编码

Transformer 没有循环结构，如果我们直接把 token 向量输入进模型，那么序列顺序信息就完全丢失了。比如“我喜欢猫”和“猫喜欢我”，词一样，但含义不同。必须把位置信息加入模型。

论文中使用的是正余弦位置编码：

偶数维度用 sin 函数。
奇数维度用 cos 函数。

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).unsqueeze(1).float() 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) pe = pe.unsqueeze(0) self.register_buffer("pe", pe) def forward(self, x): return x + self.pe[:, : x.size(1)]

这里的核心思想是：不同位置的向量在不同维度上具有不同的相位差，模型可以通过注意力计算来学习位置之间的关系。

register_buffer的作用是把pe注册为模型的持久缓冲区，它不会参与梯度更新，但会随着模型一起迁移到 GPU 或者保存到 checkpoint 中。

forward中直接做加法，因为词嵌入之后通过位置编码叠加位置信息，这是最常见的使用方式。

3.3 前馈网络与层归一化

完成注意力计算后，每个位置还要经过一个前馈网络。这个网络就是两次全连接变换，中间使用 ReLU 激活函数：

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.relu = nn.ReLU() self.linear2 = nn.Linear(d_ff, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x): return self.linear2(self.dropout(self.relu(self.linear1(x))))

前馈网络的作用是对每个位置独立地进行非线性特征变换。注意力负责在位置之间交换信息，前馈网络负责在每个位置上做更复杂的特征映射。

关于层归一化 LayerNorm，PyTorch 直接提供了nn.LayerNorm(d_model)，它会对最后一个维度做归一化。LayerNorm 和 BatchNorm 的区别在于：BatchNorm 在 batch 维度上做归一化，依赖 batch 内样本统计量；LayerNorm 在特征维度上做归一化，不受 batch 大小影响，更适合 NLP 任务。

4. 编码器与解码器

4.1 编码器层

编码器层由两部分组成：多头自注意力 + 前馈网络。每部分都带有残差连接和层归一化。

class EncoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout=0.1): super().__init__() self.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) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) def forward(self, x, mask=None): attn_out = self.self_attn(x, x, x, mask) x = self.norm1(x + self.dropout1(attn_out)) ffn_out = self.ffn(x) x = self.norm2(x + self.dropout2(ffn_out)) return x

注意自注意力三个输入都是x，也就是 query、key、value 来自同一个序列，所以叫“自”注意力。

残差连接的作用是让梯度可以直接从后面层传到前面层，解决深层网络训练困难的问题。实际操作顺序是：先计算注意力输出，然后和原始输入相加，再送入 LayerNorm。

4.2 解码器层

解码器层稍微复杂一些，包含三部分：

掩码自注意力，防止看到未来信息。
交叉注意力，query 来自解码器，key 和 value 来自编码器输出。
前馈网络。

class DecoderLayer(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout=0.1): super().__init__() self.self_attn = MultiHeadAttention(d_model, n_heads, dropout) self.cross_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) 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, memory, tgt_mask=None, memory_mask=None): attn_out = self.self_attn(x, x, x, tgt_mask) x = self.norm1(x + self.dropout1(attn_out)) cross_out = self.cross_attn(x, memory, memory, memory_mask) x = self.norm2(x + self.dropout2(cross_out)) ffn_out = self.ffn(x) x = self.norm3(x + self.dropout3(ffn_out)) return x

这里memory就是编码器输出的完整表示。交叉注意力的 query 是解码器当前状态，key 和 value 来自编码器输出，这样生成目标 token 时可以关注源序列中相关的部分。

掩码自注意力中的 mask 是关键。训练时我们一次性传入完整的目标序列，不希望模型偷看未来的 token，所以用下三角 mask。

4.3 完整 Transformer 组装

把以上所有模块组装起来：

class Transformer(nn.Module): def __init__( self, src_vocab_size, tgt_vocab_size, d_model=512, n_heads=8, d_ff=2048, n_layers=6, dropout=0.1, max_len=5000, ): super().__init__() self.d_model = d_model self.src_embedding = nn.Embedding(src_vocab_size, d_model) self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model) self.pos_encoding = PositionalEncoding(d_model, max_len) self.encoder_layers = nn.ModuleList( [EncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)] ) self.decoder_layers = nn.ModuleList( [DecoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)] ) self.dropout = nn.Dropout(dropout) self.output_proj = nn.Linear(d_model, tgt_vocab_size) def forward(self, src, tgt): seq_len = tgt.size(1) tgt_mask = torch.tril( torch.ones(seq_len, seq_len, device=tgt.device) ).view(1, 1, seq_len, seq_len) src_emb = self.dropout(self.pos_encoding(self.src_embedding(src))) memory = src_emb for layer in self.encoder_layers: memory = layer(memory) tgt_emb = self.dropout(self.pos_encoding(self.tgt_embedding(tgt))) out = tgt_emb for layer in self.decoder_layers: out = layer(out, memory, tgt_mask) logits = self.output_proj(out) return logits

这里有几个细节要注意。

Embedding 层通常不会对向量做额外缩放，但原版论文中会把嵌入向量乘以根号 d_model，让嵌入值和位置编码的数值范围更接近。本文为了简单直接相加，不影响理解。

nn.ModuleList不会自动追踪模块内部的参数吗？它会的，ModuleList会把子模块注册到主模块中，让model.to(device)和model.parameters()正常工作。只是它不负责 forward 的调用，需要手动遍历。

关于 mask：torch.tril生成下三角矩阵，形状是[seq_len, seq_len]，通过view扩展成[1, 1, seq_len, seq_len]。因为在多头注意力中 scores 形状是[batch, heads, seq_len, seq_len]，mask 从左边开始广播，batch 和 heads 维度都自动扩展到实际大小。

上面的 forward 中我没有额外传 padding mask，是因为演示数据没有 padding。实际项目中，如果序列长度不齐，需要额外构造 src_key_padding_mask，在注意力分数中把 padding 位置设为负无穷。

5. 实战：训练一个倒序复制模型

5.1 任务定义

为了快速验证模型实现是否正确，我们构造一个简单的序列到序列任务：输入一串随机整数，输出这串整数的倒序。

例如输入[3, 7, 1, 5, 8, 2]，目标输出[2, 8, 5, 1, 7, 3]。

这个任务虽然简单，但要求模型能够利用全局信息，理解输入序列每个位置并完成重新排列，正好可以检验 Transformer 的序列建模能力。

5.2 数据生成与训练代码

import torch import torch.nn as nn from model import Transformer VOCAB_SIZE = 20 SEQ_LEN = 6 BATCH_SIZE = 64 EPOCHS = 40 LR = 1e-3 def make_batch(batch_size, seq_len, vocab_size): src = torch.randint(1, vocab_size, (batch_size, seq_len)) tgt = torch.flip(src, dims=[1]) bos = torch.zeros(batch_size, 1, dtype=torch.long) tgt_in = torch.cat([bos, tgt[:, :-1]], dim=1) return src, tgt_in, tgt def main(): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = Transformer( src_vocab_size=VOCAB_SIZE, tgt_vocab_size=VOCAB_SIZE, d_model=32, n_heads=4, d_ff=64, n_layers=2, dropout=0.1, max_len=SEQ_LEN, ).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=LR) criterion = nn.CrossEntropyLoss() model.train() for epoch in range(1, EPOCHS + 1): src, tgt_in, tgt = make_batch(BATCH_SIZE, SEQ_LEN, VOCAB_SIZE) src, tgt_in, tgt = src.to(device), tgt_in.to(device), tgt.to(device) logits = model(src, tgt_in) loss = criterion(logits.reshape(-1, VOCAB_SIZE), tgt.reshape(-1)) optimizer.zero_grad() loss.backward() optimizer.step() if epoch % 5 == 0: pred = logits.argmax(dim=-1) acc = (pred == tgt).float().mean().item() print(f"Epoch {epoch:02d} | loss: {loss.item():.4f} | acc: {acc:.3f}") if __name__ == "__main__": main()

这段代码有几个值得注意的地方。

倒序任务中tgt就是翻转后的序列，例如[2, 8, 5, 1, 7, 3]。我们构造解码器输入tgt_in = [BOS, 2, 8, 5, 1, 7]，也就是在目标序列前面补一个 0 作为开始符，然后去掉最后一个 token。这样每个位置的输出目标正好是tgt中对应位置的 token。

criterion使用CrossEntropyLoss，输入形状是[batch * seq_len, vocab_size]，目标形状是[batch * seq_len]，所以用reshape重新组织。

另外一个细节：make_batch每次训练迭代都重新生成随机数据，也就是说模型每次看到的都是一个全新批次。这样做可以保证训练数据非常充足，也方便观察模型是否能学出通用的倒序能力。

5.3 运行结果与分析

在 CPU 上运行这段代码，一般几十秒内就能看到结果。正常输出类似：

Epoch 05 | loss: 2.2631 | acc: 0.375 Epoch 10 | loss: 1.9854 | acc: 0.521 Epoch 15 | loss: 1.6880 | acc: 0.598 Epoch 20 | loss: 1.3861 | acc: 0.651 Epoch 25 | loss: 1.0242 | acc: 0.734 Epoch 30 | loss: 0.7031 | acc: 0.849 Epoch 35 | loss: 0.4262 | acc: 0.931 Epoch 40 | loss: 0.2371 | acc: 0.972

loss 明显下降，准确率不断提升，说明我们实现的 Transformer 结构是正确的，能够通过注意力机制完成倒序复制任务。

如果你跑到后期准确率卡在某个值上不去，可以尝试把模型 d_model 从 32 改成 64，或者n_layers从 2 改成 3，同时调小学习率，效果一般会有改善。

5.4 自回归生成验证

训练好的模型可以用来做推理。推理阶段不再使用 teacher forcing，而是逐 token 生成。

def greedy_decode(model, src, max_len=SEQ_LEN): model.eval() with torch.no_grad(): bos = torch.zeros(src.size(0), 1, dtype=torch.long, device=src.device) tgt = bos for _ in range(max_len): logits = model(src, tgt) next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True) tgt = torch.cat([tgt, next_token], dim=1) return tgt[:, 1:]

这里的关键是：生成第 i 个 token 时，模型输入是前 i-1 个已经生成的 token，然后取最后一个位置的预测概率分布。因为解码器内使用下三角 mask，所以当前位置不会看到未来 token，这是自回归生成最基本的形式。

对于倒序任务，贪心解码一般就能得到不错的效果。

6. 常见问题与排查思路

手写 Transformer 过程中，新手经常会遇到下面几类问题，我整理成表格说明。

问题现象	常见原因	排查与解决思路
程序报`d_model % n_heads != 0`	d_model 不能整除 n_heads	修改 d_model 或 n_heads，比如 d_model=32 配 n_heads=4
训练 loss 不下降	学习率过大或过小，代码存在 bug	先用小模型在简单任务上验证，检查 mask 和维度
准确率一直很低	mask 或损失函数使用错误	打印中间张量形状，检查 tgt 与 logits 是否对齐
显存不足	batch 太大或序列太长	降低 batch_size，缩小 d_model，使用梯度累积
推理时结果混乱	没有使用自回归生成	训练用 teacher forcing，推理必须逐 token 生成
mask 维度广播失败	mask 形状与 scores 不一致	确保 mask 是`[1, 1, seq_len, seq_len]`或`[batch, 1, seq_len, seq_len]`

排查维度问题时，一个很实用的建议是在 forward 中临时加打印：

print("scores shape:", scores.shape) print("mask shape:", mask.shape)

确认 scores 是[batch, heads, q_len, k_len]，mask 是[1, 1, q_len, k_len]或[batch, 1, q_len, k_len]，广播就不会出错。

另一个常见问题是忘记调用model.train()和model.eval()。Dropout 在训练和推理模式下行为不同，如果不切换，验证和推理结果会不稳定。

7. 最佳实践与工程建议

7.1 手写实现与官方 API 的选择

本文的手写实现用于理解原理。实际项目中，我更推荐使用 PyTorch 官方nn.Transformer或者 Hugging Face 的transformers库。原因很简单：官方实现经过大量优化和测试，支持更好的数值稳定性、梯度检查点、自动混合精度等功能。

transformer = nn.Transformer( d_model=512, nhead=8, num_encoder_layers=6, num_decoder_layers=6, dim_feedforward=2048, dropout=0.1, batch_first=True, )

但要注意，官方 API 的 mask 参数和维度约定和我们手写版本略有不同。使用前最好先了解tgt_mask、memory_mask、src_key_padding_mask这几个参数的区别。

7.2 模型参数与训练技巧

Transformer 对超参数比较敏感，实际工程中我建议先从小模型开始调参，再逐步放大。

小模型：d_model=64, n_heads=4, d_ff=128, n_layers=2。
中模型：d_model=256, n_heads=8, d_ff=512, n_layers=4。
大模型：d_model=512 或 768，n_heads=8 或 12，d_ff=2048，n_layers=6。

训练时优先使用 Adam 优化器，初始学习率可以考虑5e-4左右。如果训练不稳定，配合 warm up 策略：前几千步学习率线性上升，之后按步数衰减。

7.3 调试与可维护性建议

先把代码模块化，每个类尽量只做一件事，方便单独调试。比如MultiHeadAttention可以单独写一个测试用例，输入随机张量，检查输出形状是否为[batch, seq_len, d_model]，这样能快速定位是哪一层出问题。

数据方面，建议一开始不要用太大太复杂的数据集，先用几个样本做 overfit 测试。如果模型能在一个 batch 上达到接近 100% 准确率，说明结构实现没大问题，然后再扩展到完整数据集。

7.4 生产环境中的注意事项

生产环境中，Transformer 服务于线上推理时，需要额外关注：

性能优化：使用 bfloat16 或 float16 混合精度，减少显存占用。
批量推理：使用动态 batch，避免每个请求单独推理带来的显存浪费。
模型压缩：大模型可以使用量化、剪枝、蒸馏等方式。
安全边界：对输入长度做限制，防止超长序列导致显存爆炸。

另外，如果涉及模型部署和外部用户输入，必须做好输入校验，避免用户构造超长文本或畸形 token 序列影响服务稳定性。

8. 总结与下一步

到这里，我们已经用 PyTorch 从零实现了 Transformer 的全部核心组件，并在一个真实的倒序复制任务上完成了训练和验证。你亲手敲过一遍之后，再回头看架构图，会发现每个方块都有了具体的代码对应：Q、K、V 来自线性层，多头注意力的维度变换是固定的三板斧，编码器和解码器通过 memory 对接，mask 控制信息可见范围。

进一步学习的方向有几个：一是尝试用官方nn.Transformer做机器翻译任务；二是学习 BERT 和 GPT 如何基于 Transformer 改造，一个只使用编码器，一个只使用解码器；三是研究长序列优化方案，比如稀疏注意力、FlashAttention 等。Transformer 的体系非常庞大，但核心机制你已经掌握了，后续学习会顺畅很多。