IT小霸王读论文的时候，我一度觉得自己把Transformer吃透了：自注意力、多头机制、位置编码……每一个概念都能说几句。但真到了要自己动手写代码，才发现脑子里全是模糊的“大概”。直到我逼着自己用 PyTorch 从零实现了一遍完整的 Transformer，那些论文里轻描淡写的细节——mask 该怎么传、残差和 LayerNorm 的顺序为什么这么摆、训练时为什么必须用学习率预热——才真正长在了自己身上。这篇文章就是我完整手撕 Transformer 的记录，从模块拆解到代码实现，再到训练一个极简任务并填平所有坑。适合那些看完理论但还没动过手的朋友，也适合想查漏补缺的人。
1. 说动手就动手：动手前的设计与整体拆解
1.1 为什么值得手写一遍
Transformer 的代码网上到处都是，PyTorch 官方也有 nn.Transformer 可以直接调用。那为什么还要从零手写？因为“调用 API”和“会实现”之间有巨大的鸿沟。用了封装好的模块，你不需要关心 Q、K、V 是怎么分头的，不需要关心 mask 是作用在哪个维度上，更不会知道为什么训练时 Loss 会突然变成 NaN。这些细节，只有在手写的时候才会一一暴露出来。
我自己很深的体会是：手写一遍，比读十篇文章都管用。尤其是当你亲自动手把一个一个 nn.Module 拼起来、把维度 align 上、看着 loss 从 2.0 降到 0.3 的时候，你对"Attention is All You Need"这篇文章的理解会真正上一个台阶。这也是这篇文章想传达的核心：理论看得再多，不如自己写一次。
1.2 整体架构与模块划分
在敲第一行代码之前，先得把整个模型的骨架在脑子里立起来。标准的 Transformer 包含两大块：Encoder 和 Decoder，每一块又由若干个相同的 Layer 堆叠而成。
Encoder 端的职责是：把输入序列的 token 映射成一组上下文相关的表示，每个 token 都能看到整句话的全部信息。Decoder 端的职责则是：在 Encoder 输出的基础上，自回归地生成目标序列，每个位置只能看到当前位置之前的输出，所以需要 masked self-attention，同时通过 cross-attention 去”查询”Encoder 端的信息。
具体到代码层面，我打算拆成这几个模块：
PositionalEncoding：给 token embedding 加上位置信息MultiHeadAttention：多头自注意力/交叉注意力的核心PositionwiseFeedForward：每个 token 独立经过的两层全连接EncoderLayer/DecoderLayer：单层 Encoder/Decoder，包含 attn + ffn + 残差 + LayerNormEncoder/Decoder：多层堆叠 + embedding + 位置编码Transformer：把 Encoder 和 Decoder 组装在一起，加上输出投影
这么拆的好处是：每个模块的职责单一，出问题的时候能快速定位。我见过很多初学者喜欢把所有逻辑写在一个超级大的类里，结果一个维度写错，调试起来非常痛苦。模块化之后，每个模块输入输出的 shape 都清晰可控，bug 定位难度会下降一个量级。
需要模型API调用？ 免费领10W Token，多模型网关一键接入 Claude、DeepSeek 等主流模型。
2. 自注意力机制：Transformer 的灵魂
2.1 Scaled Dot-Product Attention 的数学原理
自注意力要做的事情，通俗地讲就是：对序列里的每一个 token，让它去”看”序列里其他所有 token，然后根据它们之间的相关性，把其他 token 的信息加权汇总到这个 token 的表示里。
数学上，把输入 x 分别乘上三个权重矩阵 W_Q、W_K、W_V，得到 Query、Key、Value。Query 代表“我想找什么信息”，Key 代表“我能提供什么信息”，Value 代表“我真正的内容”。然后计算 Query 和所有 Key 的点积，得到相似度分数，经过 softmax 归一化成权重，最后用这个权重去加权求和 Value。
原论文里给了一个很重要的细节：点积之后要除以 sqrt(d_k)。我一开始没太在意这个系数，觉得反正 softmax 会归一化，除不除无所谓。后来动手训练才发现，如果不除以 sqrt(d_k)，当 d_k 比较大的时候，点积的结果会变得很大，softmax 的梯度会很平，训练直接推不动。这是从数值角度解释“Scaled”的意义。
除了缩放，另一个关键点是 mask。在 Decoder 的自注意力里，为了让模型不偷看未来信息，我们要把当前位置之后的所有位置的注意力分数设成一个极小的负数（而不是 0），这样 softmax 之后它们的权重会趋近于 0。为什么要填极小的负数而不是直接填 0？因为注意力分数要经过 softmax，填 0 意味着 exp(0)=1，它的概率并不会变成 0，信息照样会“泄漏”出去。
2.2 用 PyTorch 实现多头注意力
多头注意力的思路是：不只用一组 W_Q/W_K/W_V，而是把 d_model 维度切成 n_heads 份，每一份独立地做注意力计算，最后拼回去。这样做的意义在于：不同的头可以关注不同的信息——有的头关注语法关系，有的头关注相邻词，有的头关注长距离依赖。
下面是我实现这个模块的完整代码，注释写得很详细。
python复制import torch
import torch.nn as nn
import math
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.d_k = d_model // n_heads
# 注意：这里用了 d_model -> 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.w_o = nn.Linear(d_model, d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, q, k, v, mask=None):
batch_size = q.size(0)
# 投影后拆成多头：[batch, seq_len, n_heads, d_k] -> [batch, n_heads, seq_len, d_k]
Q = self.w_q(q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
K = self.w_k(k).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
V = self.w_v(v).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
# 注意力分数：[batch, n_heads, q_len, k_len]
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attn = torch.softmax(scores, dim=-1)
attn = self.dropout(attn)
# 加权求和，然后合并多头
context = torch.matmul(attn, V) # [batch, n_heads, q_len, d_k]
context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
return self.w_o(context)
2.3 mask 为什么这么写
上面的代码里，mask 的形状我认为是初学者最容易搞混的地方。我在这里说一下我在实践中的约定：
- 对于 encoder 的 padding mask，形状是
[batch, 1, 1, src_len]，也就是标记哪些位置是 padding token。 - 对于 decoder 的 attention mask（decoder 里同时需要 padding mask 和 causal mask），形状是
[batch, 1, tgt_len, tgt_len]，矩阵中因果部分为 1，其余为 0。
为什么这里维度里有两个 1？因为 scores 的形状是 [batch, n_heads, q_len, k_len]，mask 要能广播到这个形状。把 batch 维度保留，n_heads 维度用 1 广播，这样每个头共用一个 mask，就对了。
我第一次写的时候，直接把 mask 用成了 [batch, seq_len]，然后一运行维度直接报错。后来养成了一个习惯：每次写 forward 之前，先把所有张量的 shape 推导一遍再动手，这能省掉大量 debug 时间。
3. 位置编码与前馈网络：给模型注入顺序与非线性
3.1 位置编码的实现
注意力本身是“无序”的：你把输入序列任意打乱顺序，注意力计算的结果是一样的，因为它只做两两之间的加权，不关心谁在前谁在后。所以我们必须显式地把位置信息塞进去。原论文用的是三角函数式的固定位置编码：
PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
这里的想法是：用不同频率的正弦和余弦函数来表示相对位置。pos 表示 token 在序列中的位置，i 表示维度索引。这样设计的好处是，模型可以通过线性组合来感知相对位置关系。
代码实现如下：
python复制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)
pe = pe.unsqueeze(0) # [1, max_len, d_model]
self.register_buffer("pe", pe)
def forward(self, x):
return x + self.pe[:, : x.size(1)]
注意，我用了 register_buffer，这样位置编码会和模型一起被移动到 GPU，并且在保存/加载模型时不会被认为是模型参数。
这里有一个细节：torch.arange(0, d_model, 2) 的长度应该正好等于 d_model / 2。如果 d_model 是奇数（一般不会，但不排除），赋值的时候会出现 shape 不匹配。我们这里默认 d_model 是偶数。
3.2 前馈网络与 LayerNorm、残差连接的配合
除了注意力之外，每个 Transformer Block 里还包含一个前馈网络（FFN）。它的作用是对每个 token 的表示做一次非线性的变换。论文里的设置是：线性层先把维度从 d_model 升到 d_ff（通常是 2048），经过 ReLU，再降回 d_model。
python复制class PositionwiseFeedForward(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.relu = nn.ReLU()
self.dropout = nn.Dropout(dropout)
def forward(self, x):
return self.linear2(self.dropout(self.relu(self.linear1(x))))
关于残差连接和 LayerNorm，我在这里插一句我自己踩过的坑。原论文用的是 Post-Norm，也就是“残差之后再 LayerNorm”。但后来很多工作（比如 GPT 系列）用的是 Pre-Norm，也就是“先 LayerNorm 再进入子层”。两者的区别在训练稳定性上很明显。Pre-Norm 往往更稳，训练可以开更大的学习率；而 Post-Norm 更贴近原始论文，但如果 dropout 和 learning rate 设置不当，梯度会很不稳定。
我这次实现按原论文走 Post-Norm，用残差 + dropout + LayerNorm 的组合。这个顺序有讲究：先算子层输出，加残差，再 dropout，最后 LayerNorm。换句话说，每个子层的输出是：
code复制x = LayerNorm(x + Dropout(sublayer(x)))
这么做可以让梯度在残差路径上直接回传，缓解深层网络里的梯度消失问题，同时 LayerNorm 对激活值做归一化，让下一层拿到的是一个均值为 0、方差为 1 的分布，训练更稳定。
4. 搭起 Encoder 与 Decoder
4.1 EncoderLayer 的组装
有了注意力、前馈网络和位置编码这些积木，就可以开始组装 EncoderLayer。
每一个 EncoderLayer 包含两个子层：多头自注意力子层和 FFN 子层。每个子层都按“残差 + Dropout + LayerNorm”的方式连接。我实现的时候会把 LayerNorm、Dropout 都定义好，避免在 forward 里临时创建导致重复实例化。
python复制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.feed_forward = PositionwiseFeedForward(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.feed_forward(x)
x = self.norm2(x + self.dropout2(ffn_out))
return x
然后是多层堆叠的 Encoder：
python复制class Encoder(nn.Module):
def __init__(self, vocab_size, d_model, n_heads, d_ff, n_layers, max_len, dropout=0.1):
super().__init__()
self.embedding = nn.Embedding(vocab_size, d_model)
self.positional_encoding = PositionalEncoding(d_model, max_len)
self.layers = nn.ModuleList([
EncoderLayer(d_model, n_heads, d_ff, dropout)
for _ in range(n_layers)
])
self.dropout = nn.Dropout(dropout)
def forward(self, src, mask=None):
x = self.embedding(src)
x = self.positional_encoding(x)
x = self.dropout(x)
for layer in self.layers:
x = layer(x, mask)
return x
embedding 之后先加位置编码再到 dropout，这个顺序是从原论文来的。值得一提的小细节是：embedding 层一般还会乘一个 sqrt(d_model)，因为 embedding 的方差和位置编码的量级要匹配上，不然相加时位置编码容易被淹没。如果你发现训练不太收敛，可以检查一下这一步。
4.2 DecoderLayer 的 Cross-Attention
Decoder 比 Encoder 多了一个 Cross-Attention 子层。这个子层的 Q 来自 Decoder 侧前一个子层的输出，K 和 V 来自 Encoder 的输出。也就是说，Decoder 在生成每个 token 的时候，都可以去“看”Encoder 编码出来的完整输入序列。
Decoder 第一层自注意力需要用 mask 把当前位置之后的信息遮掉，否则模型在训练时偷懒——直接复制未来时刻的 target 就完事了，学不到真正的生成能力。
python复制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.feed_forward = PositionwiseFeedForward(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_output, src_mask=None, tgt_mask=None):
# Masked Self-Attention
attn_out = self.self_attn(x, x, x, tgt_mask)
x = self.norm1(x + self.dropout1(attn_out))
# Cross-Attention：Q 来自解码器，K/V 来自编码器
attn_out = self.cross_attn(x, enc_output, enc_output, src_mask)
x = self.norm2(x + self.dropout2(attn_out))
# FFN
ffn_out = self.feed_forward(x)
x = self.norm3(x + self.dropout3(ffn_out))
return x
这里有一个容易漏掉的地方：cross-attention 也要传 src_mask。因为 Encoder 的输出里，padding 位置是一些没有意义的向量，如果 Decoder 在计算注意力时把这些位置加权进来，会把噪声带入生成过程。所以我们需要用 padding mask 把那些位置遮掉。
4.3 组装完整的 Transformer
接下来就是把上面这些模块拼成最终模型。还需要一个输出投影层，把 Decoder 输出的 d_model 维向量映射到词表大小上去，用来计算每个位置的词概率。这个投影层的权重，很多实现会共享 embedding 层和最后的输出层，可以省参数。
python复制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, max_len=512, dropout=0.1):
super().__init__()
self.encoder = Encoder(src_vocab_size, d_model, n_heads, d_ff, n_layers, max_len, dropout)
self.decoder = Decoder(tgt_vocab_size, d_model, n_heads, d_ff, n_layers, max_len, dropout)
self.output_proj = nn.Linear(d_model, tgt_vocab_size)
def forward(self, src, tgt, src_mask=None, tgt_mask=None):
enc_output = self.encoder(src, src_mask)
dec_output = self.decoder(tgt, enc_output, src_mask, tgt_mask)
return self.output_proj(dec_output)
在实际搭建的时候，我建议你用一个小配置先把 forward 跑通：d_model=64, n_heads=4, d_ff=128, n_layers=2，输入 batch 的 [batch, seq_len]，确认输出 shape 正确，再慢慢把维度调大。千万别一开始就上原论文的 512/2048/8/6，那样一旦出错，定位问题的成本会非常高。
5. 训练一个真正能跑的任务
5.1 准备一个极简任务
模型搭好了，必须跑起来看效果。我选了一个最简单又能验证模型是否真正学会的任务：序列逆序（reversal task）。输入一个数字序列，让模型输出它的逆序序列。比如输入 [1, 3, 5, 7, 9]，期望输出 [9, 7, 5, 3, 1]。
这个任务虽然简单，但能很好地检验模型是否理解了序列位置关系和自回归生成逻辑。如果你手写的是带 mask 的自注意力，逆序任务会让 mask 的作用非常直观地体现出来。为了处理序列的起止，我在每个目标序列前面加 SOS，末尾加 EOS。
数据集的构造逻辑大致是：
python复制import random
class ReverseDataset:
def __init__(self, vocab_size=50, min_len=3, max_len=8, size=10000):
self.vocab_size = vocab_size
self.min_len = min_len
self.max_len = max_len
self.size = size
def __len__(self):
return self.size
def __getitem__(self, idx):
length = random.randint(self.min_len, self.max_len)
src = [random.randint(2, self.vocab_size - 1) for _ in range(length)]
tgt = src[::-1]
return {
"src": [2] + src, # 2 表示 SOS
"tgt_input": [2] + tgt, # decoder 输入
"tgt_output": tgt + [3], # 解码器目标输出，3 表示 EOS
}
这里的小细节是 tgt_input 和 tgt_output 错开一位：Decoder 在位置 i 的输入，预测的是位置 i+1 的目标。这样模型每一步都基于“已经生成的 token”来预测下一个 token，与推理时的自回归行为保持一致。
5.2 训练循环与学习率调度
训练循环本身不复杂，但有两个点我想专门强调。
第一个是 loss 的计算。因为序列里不同样本长度不一样，我们需要按 batch 内最大长度做 padding，然后在计算交叉熵时把 padding 位置的 loss 忽略掉。PyTorch 的 CrossEntropyLoss 支持 ignore_index，我们把 padding 位置设为 0，然后在计算时传进去就行。
第二个是学习率。原论文用了所谓的 Noam 学习率调度：先线性预热，再按步数的倒数平方根衰减。这个设计在 Transformer 训练里非常管用，尤其是你不太确定最优学习率是多少的时候。我自己实验下来的体感是：没有预热阶段，一开始的 loss 会非常不稳定，甚至直接飞掉；加上预热之后，训练会顺滑很多。
python复制class NoamSchedule:
def __init__(self, optimizer, d_model, warmup_steps=4000):
self.optimizer = optimizer
self.d_model = d_model
self.warmup_steps = warmup_steps
self.step_num = 0
def on_step(self):
self.step_num += 1
lr = self.d_model ** (-0.5) * min(
self.step_num ** (-0.5),
self.step_num * self.warmup_steps ** (-1.5)
)
for param_group in self.optimizer.param_groups:
param_group["lr"] = lr
训练循环的大致结构：
python复制model = Transformer(
src_vocab_size=50,
tgt_vocab_size=50,
d_model=64,
n_heads=4,
d_ff=128,
n_layers=2,
max_len=16,
dropout=0.1,
)
optimizer = torch.optim.Adam(model.parameters(), betas=(0.9, 0.98), eps=1e-9)
scheduler = NoamSchedule(optimizer, d_model=64, warmup_steps=2000)
criterion = nn.CrossEntropyLoss(ignore_index=0)
for epoch in range(30):
total_loss = 0
for batch in dataloader:
src = batch["src"]
tgt_input = batch["tgt_input"]
tgt_output = batch["tgt_output"]
src_mask = (src != 0).unsqueeze(1).unsqueeze(2)
tgt_pad_mask = (tgt_input != 0).unsqueeze(1).unsqueeze(2)
seq_len = tgt_input.size(1)
causal_mask = torch.tril(torch.ones(seq_len, seq_len)).bool()
tgt_mask = causal_mask & tgt_pad_mask
tgt_mask = tgt_mask.unsqueeze(0) # [1, 1, seq_len, seq_len]
logits = model(src, tgt_input, src_mask, tgt_mask)
loss = criterion(logits.reshape(-1, logits.size(-1)), tgt_output.reshape(-1))
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
scheduler.on_step()
total_loss += loss.item()
这里 torch.nn.utils.clip_grad_norm_ 要特别提一下。Transformer 的训练对梯度裁剪几乎是刚需。如果不裁剪，偶尔某一步梯度特别大，模型参数一下被冲出安全区，loss 直接飞到几百甚至变成 NaN。我在一开始偷懒没加这行，结果每隔一阵子训练就崩一次，加完之后再也没有这个问题了。
6. 训练中踩过的坑与排查思路
6.1 维度相关的坑
我估计 90% 的初学者在写 Transformer 时，第一个报错都是多维张量的 shape 不匹配。
最常见的几个：
view之后忘了.transpose(1, 2)，导致多头维度变成了 seq 维度transpose之后直接.view()，报 "view size is not compatible"，需要先.contiguous()- mask 的维度少了一维，广播不到 scores 上
- Decoder 的 causal mask 和 padding mask 是
&关系，不是两个 mask 分别作用
排查这类问题，我强烈建议在关键位置打印 shape，或者直接 pdb 进去断点调试。不要靠猜。
6.2 训练不收敛的排查方向
如果你发现 loss 降不下去或者降得很慢，按以下顺序排查：
- 先检查数据是不是有问题。拿一个样本出来，打印 src 和 tgt，人工看看对应关系对不对。
- 再检查 loss 在过拟合一个 batch 时能不能降低。我习惯在完整训练之前，先拿一个 batch 把模型过拟合到 loss 接近 0，如果连一个 batch 都拟合不了，说明模型结构有 bug。
- 接着看学习率。Transformer 对学习率非常敏感，不要在没做 warmup 的情况下直接用固定的 0.001 之类的值。
6.3 数值不稳定怎么办
训练过程中 loss 变成 NaN，是另一个高频事故。可能的原因包括：
- 学习率太大，参数一步跳出可行域
- 梯度爆炸，没有梯度裁剪
- 隐藏状态出现了极端值，后面 softmax 或 layernorm 都救不回来
解决思路首先是加梯度裁剪。如果加了还是崩，就把学习率调低，或者把 warmup steps 调大。另外，Post-Norm 的 Transformer 对初始化也比较敏感，如果你用的是比较大的 d_model，可以试试把 FFN 的输出初始化调小一些。
我自己遇到的最邪门的一次，是 masked_fill 时填的 -1e9 不够小，导致在注意力很稀疏的情况下 softmax 后的分布还是有一定梯度流动。后来我查了别人的实现，用 float('-inf') 也可以，但要注意在混合精度训练下 -inf 可能会带来一些奇怪的行为，所以实践中我一般用 -1e9，并且在 FP16 训练下，这个值可以保持安全。
6.4 推理时的一个隐藏坑
训练跑通之后，你肯定想试试模型的生成效果。这里有个容易忽略的点：训练时 Decoder 是“教师强制”（teacher forcing）模式，每一步都喂真实的目标 token。但推理的时候，模型只能拿到自己前一步生成的 token，一旦某一步生成错了，后续可能会一路错下去。
所以在验证模型效果的时候，不能只依赖训练时的 teacher forcing loss，一定要写一个自回归的生成循环：先把 SOS 作为第一个输入，然后反复取 logits[:, -1, :] 的 argmax（或抽样）作为下一步输入，直到生成 EOS 或达到最大长度。我用逆序任务做完这个自回归验证之后，才真正确信自己的模型没有白写。
python复制def greedy_decode(model, src, max_len=20):
model.eval()
src = src.unsqueeze(0)
src_mask = (src != 0).unsqueeze(1).unsqueeze(2)
with torch.no_grad():
enc_output = model.encoder(src, src_mask)
tgt = torch.tensor([[2]]).to(src.device) # SOS
for _ in range(max_len):
tgt_mask = torch.tril(torch.ones(tgt.size(1), tgt.size(1))).bool().unsqueeze(0).to(src.device)
logits = model.decoder(tgt, enc_output, src_mask, tgt_mask)
logits = model.output_proj(logits)
next_token = logits[:, -1, :].argmax(dim=-1)
tgt = torch.cat([tgt, next_token.unsqueeze(0)], dim=1)
if next_token.item() == 3: # EOS
break
return tgt.squeeze(0).tolist()
拿我自己跑的结果举例：用 vocab_size=50、序列长度 5-8 的逆序任务，我用了 d_model=64, n_heads=4, d_ff=128, n_layers=2, warmup_steps=2000，基本上 10 个 epoch 之后，自回归生成的结果就已经 100% 正确了。这个任务规模很小，所以跑起来很快，非常适合用来验证手写实现的正确性。
最后再分享一个小技巧：如果你准备拿这套手写代码去做更复杂的任务，千万不要直接把 d_model=64 这种小配置沿用到 d_model=512。维度变大之后，dropout、warmup 步数、梯度裁剪的阈值，甚至初始化方式，都需要重新调整。经验值是：模型越大，dropout 可以适当调小，warmup 步数要相应增加。这些参数之间是联动的，不要孤立地调某一个。
手写 Transformer 的过程，确实像标题说的那样——“纸上得来终觉浅”。真正动手之后你才会发现，论文里每一句话背后都可能有无数个实现细节的坑在等着你。但把这些坑一个一个填平之后，你对 Transformer 的理解，会变得完全不一样。