从零手写Transformer:PyTorch实现自注意力、位置编码与文本生成
2026/8/31 23:32:55 网站建设 项目流程

Transformer 是当前自然语言处理领域最核心的模型结构,但很多人学的时候只停留在调包阶段,真要用 PyTorch 从零写一遍,才发现自注意力、QKV、位置编码这些概念是断开的,每个模块都知道大概,合在一起就不知道怎么接。这篇文章不绕弯子,直接按一条完整实操路线走:从 token 化开始,把嵌入层、位置编码、自注意力、多头注意力、残差连接、前馈网络和训练循环全部手写,最后跑一个最小可用的文本生成 demo。

适合两类读者:一类是 PyTorch 用过但没完整写过 Transformer 的人,另一类是公式看了不少、但不知道张量在每个阶段长什么样的人。最值得关注的不是某一行代码,而是整条数据流怎么从 token 一路走到损失函数。下面按落地顺序拆一遍。

1. 先想清楚:手撕 Transformer 要撕的到底是什么

1.1 你至少需要什么基础

“手撕”不是从零造数学理论,也不是复刻 GPT-4。一个最小可用的 Transformer,核心就是把下面这几件事做出来:

  • 输入是一段 token 序列。
  • 模型先把每个 token 变成向量。
  • 通过自注意力让每个位置能看到其他位置的信息。
  • 用位置编码补上顺序关系。
  • 用前馈网络做非线性变换。
  • 最后输出下一个 token 的概率分布。

这段流程拆开后非常固定。我认为新手最怕的不是公式,而是不知道张量在每一步变成什么形状。所以这篇文章会一直强调形状变化,代码里也尽量把 shape 写清楚。

基础要求其实不高:会用 Python,会写函数和循环,用过 nn.Module 和 tensor 基本运算,知道矩阵乘法。线性代数不需要懂得很深,能把(batch, seq_len, d_model)这类形状在脑子里过一遍就够。

1.2 一条主线串起所有组件

先记住整条数据流:

原始文本 → token 序列 → embedding 向量 → 位置编码 → 自注意力(QKV) → 多头拼接 → 残差/层归一化 → 前馈网络 → 线性输出 → softmax

后面每个章节都会围绕这条主线展开。代码从零写,不调用 PyTorch 自带的 Transformer 封装,这样才能真正看清内部结构。

1.3 每个阶段张量长什么样

我先给出一张形状对照表,后面写代码时会反复对应回来:

张量形状含义
输入 x(batch, seq_len)token ID 序列
嵌入后(batch, seq_len, d_model)每个 token 的向量表示
Q / K / V(batch, seq_len, d_head)注意力的三个投影
注意力分数(batch, seq_len, seq_len)两两位置之间的相关程度
注意力输出(batch, seq_len, d_head)按权重聚合后的结果
模型最终输出(batch, seq_len, vocab_size)每个位置下一个 token 的分布

这张表建议收藏。排错的时候第一反应应该是:现在这步输出形状是什么,下一步想要的是什么形状。

2. 环境准备与最小数据管线:先让 Token 真正跑起来

2.1 PyTorch 环境怎么选

先确认机器有没有可用 GPU。我第一次跑这种手写 demo 时,其实是在 CPU 上完成的,完全能跑。先执行这一句:

python -c "import torch; print(torch.cuda.is_available())"

输出 True 说明 GPU 可用,False 也没关系。小 demo 用 CPU 就能跑,只是要把 d_model、层数、seq_len 调小一点。

安装 PyTorch 建议直接去官方站点选对应系统的命令,不要凭记忆敲。Windows 下最容易踩的坑是 CUDA 版本不匹配。判断标准:先看nvidia-smi里显示的 CUDA 版本,再选不超过它的 PyTorch CUDA 版本。CPU 环境只装 CPU 版就够了,硬装 CUDA 版反而会引入一堆 driver 报错。

2.2 从文本到 Token 序列

这篇文章的 demo 用字符级 tokenizer,目的就是少踩 OOV(未登录词)的坑。词表小、实现简单,几行代码就能完成。

import torch import torch.nn as nn import torch.nn.functional as F # 原始语料 text = "hello world this is a tiny transformer demo" # 建立字符级词表 chars = sorted(set(text)) stoi = {c: i for i, c in enumerate(chars)} # 字符 -> 索引 itos = {i: c for c, i in stoi.items()} # 索引 -> 字符 vocab_size = len(chars) # 整段文本变成 token 序列 ids = [stoi[c] for c in text] print(ids[:20]) print("vocab_size:", vocab_size)

为什么用字符级?因为词表通常只有几十个字符,训练时不会出现未知词。缺点也很明显:模型要同时学拼写和语法。但我们的目标是理解 Transformer 结构,不是拿这个模型去打榜,所以字符级是最合适的选择。

如果你以后换到真实任务,可以把这段替换成 BPE 或 SentencePiece,但外层结构完全不用动。

2.3 训练样本怎么切

Transformer 训练不能把整段文本一次丢进去,要把长文本切成固定长度的输入块。我一般切完后,让输入和标签错开一位,这样每个位置都在预测下一个 token。

def get_batch(ids, batch_size=16, seq_len=16): starts = torch.randint(0, len(ids) - seq_len - 1, (batch_size,)) x = torch.stack([torch.tensor(ids[i:i+seq_len]) for i in starts]) y = torch.stack([torch.tensor(ids[i+1:i+seq_len+1]) for i in starts]) return x, y

这里 x 是输入,y 是目标输出。x 的第 t 个 token 预测的是 y 的第 t 个 token,也就是原文本中的下一个字符。这个“错位一位”的操作是整个语言模型训练的核心逻辑。

3. 自注意力与 QKV:别背公式,直接看张量形状

3.1 Q、K、V 从哪里来

自注意力的输入是同一个序列 x,但会通过三个独立的线性层投影成 Q、K、V。为什么不能直接用 x 算相似度?因为模型需要根据任务学出“用哪些信息去查询、按哪些特征匹配”。如果直接用 x,每个位置的查询和键完全一样,模型就没有自由度了。

class SelfAttention(nn.Module): def __init__(self, d_model, d_head, dropout=0.1): super().__init__() self.d_head = d_head self.q = nn.Linear(d_model, d_head) self.k = nn.Linear(d_model, d_head) self.v = nn.Linear(d_model, d_head) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): Q = self.q(x) # (batch, seq_len, d_head) K = self.k(x) # (batch, seq_len, d_head) V = self.v(x) # (batch, seq_len, d_head) scores = Q @ K.transpose(-2, -1) / (self.d_head ** 0.5) # (batch, seq_len, seq_len) if mask is not None: scores = scores.masked_fill(mask == 0, float("-inf")) weights = torch.softmax(scores, dim=-1) weights = self.dropout(weights) out = weights @ V # (batch, seq_len, d_head) return out

这 30 行不到的代码就是单头自注意力的全部。关键动作就两个:

  1. Q 和 K 做点积,得到两两位置的相关度。
  2. softmax 把相关度变成权重,再用权重去加权 V。

Q 决定“我在找什么”,K 决定“我是什么内容”,V 决定“我最终提供什么信息”。三者分工明确。

3.2 为什么除根号 d_head

如果不除以根号 d_head,d_head 越大,点积的数值会越大,softmax 之后容易进入饱和区,梯度变得非常小。除一个根号 d_head 是为了让分数方差保持稳定。

这行代码很小的优化,但对训练稳定性影响很大。很多新手直接照抄公式,从不调这一项,结果模型训练时 loss 不稳定,然后到处怀疑数据问题。

3.3 因果掩码:为什么训练 decoder 要遮住未来

语言模型生成时只能看当前时刻之前的内容,不能偷看未来 token。训练阶段也必须模拟这个限制,否则模型会学到“作弊”模式。实现方式叫因果掩码。

def causal_mask(seq_len): return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool))

这个函数生成一个下三角矩阵:当前位置允许看到自己和之前的位置,不能看到之后的位置。

之后在注意力分数计算时,把这个掩码传进去:

mask = causal_mask(seq_len).unsqueeze(0) # (1, seq_len, seq_len) scores = scores.masked_fill(mask == 0, float("-inf"))

把未来位置对应的分数变成负无穷,softmax 之后权重就趋近于 0。这样每个位置只能聚合到当前位置和之前位置的信息。

这里要特别注意掩码的形状。如果 batch size 不为 1,通常要扩展成(batch, 1, seq_len, seq_len)或者(1, seq_len, seq_len)让它广播,千万不要直接拿(seq_len, seq_len)去和(batch, seq_len, seq_len)的 scores 做 masked_fill,会报形状错误。

4. 位置编码与多头注意力:补上顺序和多视角

4.1 为什么注意力本身没有顺序概念

自注意力是集合运算。把 token 序列乱序打乱,如果不加位置信息,输出排列也不会按语义变化。这对语言任务来说是致命的:"I love you""you love I"在集合意义上看起来一样,但语义完全不同。

所以必须在 embedding 阶段把位置信息加进去。最直接的方式是准备一个同样维度的位置向量,和 token embedding 相加。

4.2 正弦编码还是可学习编码

两种方案都常见:

  • 可学习位置编码:用nn.Embedding(max_len, d_model),实现简单,短序列任务表现稳定。
  • 正弦位置编码:用固定公式生成,不依赖训练,理论上对任意长度都能给出位置向量,外推能力更强。

小 demo 我建议直接用可学习位置编码,代码少,不容易错。如果以后要做长度外推或者长文本场景,再看正弦编码。

def sinusoidal_encoding(max_len, d_model): pe = torch.zeros(max_len, d_model) pos = torch.arange(max_len).unsqueeze(1).float() div = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model)) pe[:, 0::2] = torch.sin(pos * div) pe[:, 1::2] = torch.cos(pos * div) return pe.unsqueeze(0) # (1, max_len, d_model)

这个公式本身不需要背,但要理解它的作用:偶数维用 sin,奇数维用 cos,不同频率组合起来,让模型可以区分不同位置。

4.3 多头注意力:多个子空间并行

单头注意力只能关注一种关系。多头注意力的思路是:用多个独立的注意力头,每个头在不同的特征子空间里计算相关性,最后把所有头的结果拼起来再做一次线性变换。

class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super().__init__() assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除" self.num_heads = num_heads self.d_head = d_model // num_heads self.heads = nn.ModuleList( [SelfAttention(d_model, self.d_head) for _ in range(num_heads)] ) self.proj = nn.Linear(d_model, d_model) def forward(self, x, mask=None): heads = [h(x, mask) for h in self.heads] out = torch.cat(heads, dim=-1) # (batch, seq_len, d_model) return self.proj(out)

为什么需要多个头?因为一个头只能学一种权重分布。有些头可能负责相邻词搭配,有些头负责长距离指代,有些头负责某个具体句式。多个头相当于多个视角同时看问题。

实际实现里,主流做法是把 Q、K、V 一次投影到整个 d_model,再 reshape 拆成多个头。我这里为了可读性,用了 ModuleList,每个头独立投影。代码直观一些,不过计算效率不如高端写法。理解阶段不需要过度优化。

4.4 常见错误:d_model 不能被 num_heads 整除

每个头的维度是 d_model / num_heads,所以 d_model 必须能被 num_heads 整除。新手最容易这里报错。我一般开头加一行 assert,至少报错信息清晰。

5. 完整 Block:残差、层归一化、前馈网络

5.1 一个 Transformer Block 的标准结构

到了这一节,前面写好的模块就开始组合了。一个 decoder-only 的 Transformer Block 通常包含:

  1. LayerNorm 归一化。
  2. 多头自注意力。
  3. 残差连接。
  4. LayerNorm 归一化。
  5. 前馈网络(FFN)。
  6. 残差连接。
class TransformerBlock(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout=0.1): super().__init__() self.attn = MultiHeadAttention(d_model, num_heads) self.ffn = nn.Sequential( nn.Linear(d_model, d_ff), nn.ReLU(), nn.Linear(d_ff, d_model), ) self.ln1 = nn.LayerNorm(d_model) self.ln2 = nn.LayerNorm(d_model) self.dropout = nn.Dropout(dropout) def forward(self, x, mask=None): x = x + self.dropout(self.attn(self.ln1(x), mask)) x = x + self.dropout(self.ffn(self.ln2(x))) return x

这段代码几乎是所有 Transformer 大模型的骨架。GPT 系列、Llama 系列的 block 虽然细节有改动,但主结构都能看到这个影子。

5.2 为什么先 LayerNorm 再进注意力

我写的是 Pre-LN:先归一化,再计算注意力。这和原始 Transformer 论文里的 Post-LN 顺序不同。原始论文是先注意力、后归一化,训练起来对学习率更敏感,调参稍不注意就发散。

Pre-LN 的好处是训练更稳定,适合新手。你可以简单理解成:每层输入先被拉回一个正常分布,再进注意力或前馈网络,这样梯度不容易爆炸。

5.3 前馈网络不是可有可无

注意力本质是 token 之间的信息交换和加权求和,整体计算偏线性。前馈网络提供逐位置的非线性变换,让模型表达能力更强。

d_ff 通常取 d_model 的 2 到 4 倍。我这个小 demo 里 d_model=64,d_ff=128,已经够用。你如果调大 d_ff,模型容量会上升,但训练时间和显存占用也会明显增加。

6. 从零训练一个小模型:文本生成 Demo

6.1 完整模型定义

把前面的模块拼起来,就是一个极小的 decoder-only Transformer:

class TinyTransformer(nn.Module): def __init__(self, vocab_size, d_model=64, num_heads=4, d_ff=128, num_layers=2, max_len=128, dropout=0.1): super().__init__() self.token_embedding = nn.Embedding(vocab_size, d_model) self.position_embedding = nn.Embedding(max_len, d_model) self.blocks = nn.ModuleList([ TransformerBlock(d_model, num_heads, d_ff, dropout) for _ in range(num_layers) ]) self.ln = nn.LayerNorm(d_model) self.head = nn.Linear(d_model, vocab_size) def forward(self, x, mask=None): B, T = x.shape pos = torch.arange(T, device=x.device).unsqueeze(0) x = self.token_embedding(x) + self.position_embedding(pos) for block in self.blocks: x = block(x, mask) return self.head(self.ln(x))

这里的位置编码用nn.Embedding做的可学习位置编码。每次 forward 时,根据当前序列长度 T 生成位置索引。

6.2 训练循环怎么写

训练用交叉熵损失。注意 logits 和目标形状要对应:logits 是(batch, seq_len, vocab_size),目标 y 是(batch, seq_len)。CrossEntropyLoss 默认输入是(N, C),所以要把前两维展平。

model = TinyTransformer(vocab_size) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) loss_fn = nn.CrossEntropyLoss() for step in range(500): x, y = get_batch(ids, batch_size=16, seq_len=16) mask = causal_mask(16) logits = model(x, mask) loss = loss_fn(logits.reshape(-1, vocab_size), y.reshape(-1)) optimizer.zero_grad() loss.backward() optimizer.step() if step % 100 == 0: print(f"step {step}: loss {loss.item():.4f}")

这个循环里新手最常忽略的是optimizer.zero_grad()。如果不每次清零,梯度会在多个 step 之间累加,loss 会表现得非常诡异。

6.3 学习率怎么调

1e-3 对这个 5 万参数级别的小模型通常能跑。如果 loss 发散,优先降到 3e-4 或 1e-4。不要一上来就加 batch size,先把学习率降下来往往更有效。

这里也提醒一句:小数据任务不要开大步长训练,否则模型很快过拟合到语料上,loss 降了但生成时会一直重复。

6.4 生成阶段:采样而不是贪心

训练完成后,生成的基本思路是:把已经生成的 token 放回模型,预测下一个 token,然后拼回去继续预测。

def generate(model, start_text, max_new=50, temperature=0.8): model.eval() idx = torch.tensor([[stoi[c] for c in start_text]]) with torch.no_grad(): for _ in range(max_new): x = idx[:, -32:] mask = causal_mask(x.size(1)) logits = model(x, mask)[0, -1] probs = torch.softmax(logits / temperature, dim=-1) next_id = torch.multinomial(probs, 1).item() idx = torch.cat([idx, torch.tensor([[next_id]])], dim=1) return ''.join(itos[i] for i in idx[0].tolist())

temperature 控制随机性。温度越低,输出越确定,但容易重复;温度越高,输出越随机。小模型上我一般先试 0.6 到 0.9。直接选概率最大的 token(贪心)也不是不行,但输出会非常单调。

print(generate(model, start_text="hello"))

6.5 怎样算训练成功

不要指望这个迷你模型生成通顺长句。判断标准很简单:loss 明显下降,不振荡过大;生成的字符串里开始出现训练语料的高频字符组合。比如语料里有 hello,模型能输出 h、e、l 这类常见顺序,就说明已经学到东西了。

我的建议是:先把单条训练 pass 跑通,再去看生成质量。输出不对的时候,不要急着改模型结构,先确认 loss 有没有降、采样代码有没有问题。

7. 常见报错与排查顺序:新手最容易卡在哪

7.1 形状不匹配:先打印 shape

最典型的报错是mat1 and mat2 shapes cannot be multiplied。原因通常是某个 Linear 层输入维度不是期望的 d_model。

排查顺序很固定:

  1. 先看报错发生在哪一行。
  2. 在那一层之前加一行print(x.shape)
  3. 对照第三节的张量形状表,确认当前输出和下一层期望输入是否一致。

我几乎每次都靠打印 shape 解决问题。模型结构越复杂,越不要靠肉眼猜。

常见报错对应关系:

现象常见原因排查方向
mat1 and mat2 形状不匹配某层输入维度不对打印每层入口 shape
mask 广播失败掩码形状不是 (1, T, T)打印 mask.shape,必要时加 batch 维度
device mismatch模型在 GPU、输入在 CPU模型和输入统一 .to(device)
loss 变成 nan学习率太大或输入异常降低学习率,检查 tokenizer

7.2 训练不收敛:先看数据,再改参数

loss 一直不降,不要急着认为模型写得有问题。优先看三件事:

  • 输入 token 序列是否为空,词表是否建对。
  • 输入和标签是否错位正确。如果 x 和 y 完全相同,模型就是在学复制粘贴,loss 不会按预期下降。
  • 学习率是不是太大,导致梯度震荡。

loss 变成 nan,最常见原因是学习率太大。小模型直接用 1e-3 一般没事,但如果数据量很小,还是降到 3e-4 更稳。

7.3 显存、内存和任务规模怎么匹配

低配置能跑通 demo,不代表适合批量训练。不同任务规模的参数选择,可以参考这张表:

任务类型d_modellayers数据规模预期
学习 demo32-641-2几 KBCPU 分钟级
真实小语料128-2564-6MB 级CPU 较慢,建议 GPU
大规模训练512 以上6-12GB 级需要多卡和分布式优化

如果显存不够,不要无脑加 batch size。优先缩小 seq_len,或者减小 d_model。很多时候 seq_len 对显存的影响比 batch size 更明显,因为注意力矩阵是 seq_len 的平方。

7.4 生成重复:调整温度,而不是改模型

小模型生成时容易陷入重复循环。这时候不要急着增大网络,先试以下几种方式:

  • 降低 temperature,让采样更接近高概率 token。
  • 生成时限制最大长度,避免无限循环。
  • 检查训练 loss 是不是已经降到很低,如果过拟合,会加剧重复。

如果你希望后续做更深入的调试,还可以在注意力层把 weights 返回出来,画一个注意力热力图,看看模型到底在关注哪些位置。新手阶段不用急,先把主链路跑通。


如果把这套代码完整跑一遍,你会发现 Transformer 真正难的不是某个公式,而是数据、形状、训练和生成之间的衔接。我个人更建议先把单条训练 pass 跑稳,再调参数。每一步都确认张量形状和 loss 变化,比一次性把所有组件堆完要快得多。下次再有人问“能不能手撕 Transformer”,你可以直接把这篇文章发给他,让 TA 从 token 开始,跑一遍自己的小模型。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询