从零开始用 PyTorch 构建 Transformer
我完全从零开始训练了自己的Transformer模型——没有HuggingFace，没有捷径。这里是完整的技术拆解：多头注意力、位置编码、训练循环，以及那些真正教会我这些机制如何运作的错误。
为什么要从零构建？
使用Transformer与理解Transformer之间存在巨大的鸿沟。我花了数月时间通过API微调预训练模型——Groq、Gemini、Mistral——服务于Manshverse项目。但我不断撞上一堵墙：当某些东西出错时，我无法推理其背后的原因。我把模型当作一个黑箱来对待。
于是我决定从零构建一个。不是教程中的玩具实现，而是一个真正可训练的Transformer架构，用纯PyTorch编写，可以在自己的数据上训练并研究其损失曲线。目标不是与GPT竞争——而是理解这台机器。
如果你不能从零构建它，你就没有真正理解它。——理查德·费曼（意译）
从RNN开始
在接触Transformer之前，我先从循环神经网络入手。这是有意为之——只有体会过顺序处理之痛，你才能真正欣赏注意力机制。
RNN每次处理一个词元。第t个词元的隐藏状态传递给第t+1个词元。这在理论上很优雅，但在实践中却是灾难性的：
梯度消失：随着序列变长，早期词元的信息呈指数级衰减。到第50个词元时，模型几乎已经记不清第1个词元了。
顺序瓶颈：训练无法并行化。每个词元都依赖于前一个词元。在我的本地机器上，这导致训练极其缓慢。
长距离依赖：模型根本无法学习相隔超过约20个位置的词元之间的关系。
我在莎士比亚文本上训练了一个字符级RNN。它在几百个epoch内学会了基本的单词结构，但输出超越单个短语后就变得语无伦次。损失在1.8附近停滞，拒绝下降。这是大多数人放弃并转向某个库的地方——而我转向了注意力机制。
自注意力：核心机制
Transformer论文（《Attention Is All You Need》）的关键洞见是用注意力完全取代循环。不再顺序处理词元，而是每个词元可以并行地关注所有其他词元。
这是数学原理。对于嵌入为向量的词元序列，我们计算三个投影：
Q、K、V 投影
Q = X @ self.W_q # 我在寻找什么？
K = X @ self.W_k # 我包含什么？
V = X @ self.W_v # 我携带什么信息？
缩放点积注意力
scores = (Q @ K.transpose(-2, -1)) / math.sqrt(d_k)
attn_weights = F.softmax(scores, dim=-1)
output = attn_weights @ V
Q @ K.transpose()运算计算每对词元之间的兼容性分数。softmax将其归一化为概率。然后attn_weights @ V产生所有词元表示的加权组合。
除以√d_k至关重要，而我最初忘记了这一步。如果没有缩放，点积在维度较大时会变得很大，将softmax推入梯度几乎为零的区域。在我加上这一行之前，我的模型完全不收敛。
关键洞见
自注意力在序列长度上具有O(n²)的复杂度——每个词元都关注所有其他词元。这就是Transformer上下文窗口有硬性限制的原因。这也是它们如此强大的原因：它们可以在单次前向传播中建模整个序列中的关系。
多头注意力机制
单个注意力头只能学习一种类型的关系。多头注意力机制并行运行多个注意力头，每个头拥有自己学习得到的投影，然后将结果拼接并投影：
class MultiHeadAttention(nn.Module):
def init(self, d_model, n_heads):
super().init()
self.n_heads = n_heads
self.d_k = 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)
def forward(self, x, mask=None):
B, T, C = x.shape
# Project and reshape for multiple heads
q = self.W_q(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
k = self.W_k(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
v = self.W_v(x).view(B, T, self.n_heads, self.d_k).transpose(1, 2)
# Attention scores
scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
attn = F.softmax(scores, dim=-1)
out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C)
return self.W_o(out)
我使用了8个注意力头，d_model=512，每个头的维度为64。在实践中，不同的头会学习关注不同的模式——有些关注相邻词元（局部语法），有些关注远距离依赖（语义关系）。
位置编码
注意力是置换不变的——它不知道词元的顺序。没有位置信息，从模型的角度来看，“the cat sat on the mat”与“mat the on sat cat the”是一样的。
原论文使用正弦位置编码：
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)
self.register_buffer('pe', pe.unsqueeze(0))
def forward(self, x):
return x + self.pe[:, :x.size(1)]
正弦编码非常巧妙：每个维度以不同的频率振荡，从而为每个位置生成独特的位置“指纹”。而且由于 sin(a+b) 可以表示为 sin(a) 和 cos(a) 的线性组合，模型可以通过线性变换学习相对位置。
完整架构
完整的Transformer块将多头注意力与前馈网络、层归一化和残差连接堆叠在一起：
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)
self.ff = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(),
nn.Linear(d_ff, d_model),
nn.Dropout(dropout)
)
self.ln1 = nn.LayerNorm(d_model)
self.ln2 = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x, mask=None):
# Pre-norm architecture (more stable training)
x = x + self.dropout(self.attn(self.ln1(x), mask))
x = x + self.ff(self.ln2(x))
return x
我使用的是预归一化（子层前的层归一化）而非后归一化（原论文的方法）。预归一化在训练过程中显著更稳定——梯度通过残差连接更均匀地流动，而且我无需使用学习率预热来防止早期发散。
我的最终模型：6层，8个头，d_model=512，d_ff=2048。大约4500万参数。
训练循环
训练transformer是理论遭遇残酷计算现实的地方。我最初的尝试是在一台没有GPU的本地机器上——纯CPU训练。我的数据集上一个epoch耗时超过40分钟。
让训练成功的关键决策：
AdamW优化器，权重衰减设为0.01。常规Adam导致几个epoch内就过拟合。
余弦退火学习率调度，从3e-4开始。预热阶段至关重要——直接跳到峰值学习率会立即导致损失发散。
梯度裁剪，max_norm=1.0。没有它，每50-100步就会因偶发的梯度爆炸而出现NaN损失。
迁移到GPU后使用混合精度训练（fp16）——速度提升近2倍，且质量无损。
调试技巧
如果你的transformer损失提前停滞且无法下降：检查你的注意力掩码。我花了整整两天调试一个“学习停滞”问题，结果发现是因果掩码搞错了——模型在训练时关注了未来的token，所以它学会了作弊而非预测。
扩展到云GPU
在CPU上进行本地训练虽有益处但不实用。我迁移到了Colab上的Google T4 GPU，后来又用上了云实例。T4的16GB显存能轻松容纳我的4500万参数模型，批量大小最高可达32。
速度差异惊人。之前在CPU上每个epoch需要40多分钟的训练，在T4上降到了不到3分钟。我现在可以快速迭代——尝试不同的超参数、可视化注意力模式、实时观察损失曲线。
经过约50个epoch的细致超参数调整，我在文本数据集上达到了0.85的训练损失。生成的输出在句子层面连贯，并显示出对语法和基本语义的清晰理解——远超我之前的RNN尝试。
我学到的
从头构建transformer教会了我任何教程或API封装都无法传授的东西：
注意力机制并非魔法，而是一种习得的加权平均。所谓“魔法”，源于在数百万次梯度更新中，学会该平均什么以及赋予每个词元多少权重。
规模即一切。我的4500万参数模型能够生成连贯文本，而GPT-3拥有1750亿参数。架构完全相同——差异纯粹在于规模、数据和算力。
调试机器学习与常规开发有着根本区别。在Web开发中，一个bug产生错误输出；在机器学习中，一个bug产生略微不那么正确的输出，而你无法判断这是bug还是仅仅训练不足。学会区分这两者本身就是一项技能。
损失曲线能告诉你一切。尖锐的尖峰意味着梯度爆炸；初始下降后的平台期意味着学习率过低或架构存在瓶颈；震荡意味着批大小过小。我学会了像医生解读心电图一样阅读损失曲线。
理解机器会改变你使用它的方式。在构建了自己的Transformer之后，我在提示工程、微调和调试生产级AI系统方面的能力显著提升。当Manshverse的Groq集成表现异常时，我现在能够从第一性原理出发推理其原因。
如果你认真对待AI工程，我认为从零构建一个Transformer是你能做的最有价值的学习练习。不是因为你能构建出比OpenAI更好的模型——而是因为理解机器会让你成为更好的构建者，仅此而已。
——
zhirenhun
一个热爱技术的程序员，喜欢分享前沿AI知识和开发经验。
