十七、Transformer 阶段完整总结(Day 3-6)

17.1 你走过的路

[✅] Day 1: Transformer 背景故事 — RNN 为什么不行
[✅] Day 2: Q/K/V 直觉 — 图书馆检索比喻
[✅] Day 3: 手写 Self-Attention          → self_attention.py (SelfAttention)
[✅] Day 4: 手写 Multi-Head Attention    → self_attention.py (MultiHeadSelfAttention)
[✅] Day 5: 手写 Position Encoding       → position_encoding.py
[✅] Day 6: 手写 FFN + EncoderLayer      → transfomer_encoder.py

4 个可运行代码,约 100 行核心逻辑,就是 GPT/Claude/ViT/VLA 的心脏。

17.2 为什么需要 Transformer

RNN 处理序列有 3 个致命问题:

问题RNNTransformer
长距离依赖梯度沿 100 层衰减 → 学不到每个词直接看所有词,距离无关
并行必须顺序处理 → 慢一次矩阵乘算所有位置 → 快
信息瓶颈全塞进最后一个 h_t每个词都保留自己的向量

核心洞察: 不按顺序读,让每个词直接"看到"其他所有词。


17.3 Self-Attention 核心公式

              Q · K^T
Attention = softmax(───────) · V
                    √d_k

图书馆比喻:

  • Q(Query): 你的查询 — “我在找什么”
  • K(Key): 每本书的封皮 — “我能被谁找到”
  • V(Value): 每本书的内容 — “我能提供什么”

5 步流程:

Q = x @ W_Q                              # (batch, seq, d) → (batch, seq, d)
K = x @ W_K
V = x @ W_V

scores = Q @ K.transpose(-2, -1)         # (batch, seq, seq)
scores = scores / math.sqrt(d_k)         # 缩放,防止 softmax 梯度饱和
weights = F.softmax(scores, dim=-1)      # 每行加起来 = 1
output = weights @ V                     # (batch, seq, d)

关键 shape:

  • 输入 (batch, seq_len, d_model)
  • 输出 (batch, seq_len, d_model) — 形状不变
  • Attention weights (batch, seq_len, seq_len) — “每个词对每个词的关注度”

17.4 Multi-Head Attention

为什么单头不够: 一个 attention 只能给一组权重,但一个词往往有多种关系(主谓 / 动宾 / 位置 / 修饰)。

多头思想: 把 d_model 拆成 num_heads × d_k,每个 head 独立学一种关系。


d_model = 128, num_heads = 8 → d_k = 16

注意: 用整除 //,不是 /(/ 得到浮点数,reshape 会报错)。

关键 shape 流转:

输入:              (batch, seq_len, 128)
Q/K/V 生成:        (batch, seq_len, 128)
reshape+transpose: (batch, 8, seq_len, 16)  ← 每头独立
scores:            (batch, 8, seq_len, seq_len)
weights:           (batch, 8, seq_len, seq_len)  ← 8 组权重!
output:            (batch, 8, seq_len, 16)
transpose+reshape: (batch, seq_len, 128)
W_O 投影:          (batch, seq_len, 128)

注意: 除以 √d_k(每头维度),不是 √d_model。

参数量: 4 个 Linear(d_model, d_model)(比单头多一个 W_O):

  • 单头: 3 × 16512 = 49,536
  • 多头: 4 × 16512 = 66,048(只多 33%)

行业惯例: d_k 保持在 64-128 之间,先定 d_k 再算 num_heads = d_model // d_k。

模型d_modelnum_headsd_k
BERT-base7681264
GPT-31228896128
Llama-7B409632128

17.5 Position Encoding

问题: Self-Attention 天然没有位置概念。“猫追狗” 和 “狗追猫” 经过 attention 后结果一样。

Sinusoidal PE 公式:

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

做法: 生成 (max_len, d_model) 的固定矩阵,加到 embedding 上:


x = x + self.pe[:, :seq_len, :]  # 截取前 seq_len 个位置

用 register_buffer 注册:不是可学习参数,但跟模型走(.to("cuda"))。

验证正确的最强检查: 位置 0 的编码前 8 维必须是 [0, 1, 0, 1, 0, 1, 0, 1](因为 sin(0)=0, cos(0)=1)。

Sinusoidal PE vs Learned PE:

类型优点缺点
Sinusoidal(原论文)不用学,能外推到更长序列表达可能不够灵活
Learned(BERT/GPT/ViT)更灵活,能学到最适合任务的位置超过 max_len 就不行

关键坑: 广播加法。self.pe 是 (1, max_len, d_model),输入是 (batch, seq_len, d_model),seq_len < max_len 时必须切片,不能直接 x + self.pe。


17.6 FFN(Feed-Forward Network)

为什么需要: Attention 是"线性聚合",没有真正的非线性变换。FFN 补上这个能力。

结构: 2 层 MLP + GELU 激活。

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)     # 扩宽
        self.linear2 = nn.Linear(d_ff, d_model)     # 收窄

    def forward(self, x):
        return self.linear2(F.gelu(self.linear1(x)))

关键设计:

  • d_ff = 4 × d_model(Transformer 论文经验值)
  • 用 GELU 不用 ReLU(GELU 是 ReLU 的平滑版,训练更稳)

参数量占比:一层 EncoderLayer 里,FFN 占 2/3,Attention 只占 1/3。大模型的大部分知识存储在 FFN 里。


17.7 LayerNorm

跟 BatchNorm 的对比:

BatchNormLayerNorm
归一化范围同一 batch 里所有样本的同一维度同一样本内部所有维度
对 batch 依赖强(小 batch 就崩)无
常用场景CNNTransformer / RNN

用法极简:


self.norm = nn.LayerNorm(d_model)   # 参数就是最后一维的大小
x = self.norm(x)                    # 输入输出 shape 完全一样

参数量: 2 × d_model(每个维度学一个 scale 和 shift)。极少。

Pre-LN vs Post-LN:

Post-LN(原论文):    x → sublayer → +residual → LayerNorm
Pre-LN(现代主流):   x → LayerNorm → sublayer → +residual

现代模型全用 Pre-LN,训练更稳。


17.8 Residual Connection

跟 ResNet 一样的 skip connection。Transformer 里每个子层都用:


x = x + sublayer(norm(x))

为什么每个子层都要: 让梯度能沿"高速公路"传下去,深层 Transformer(100+ 层)才能训。


17.9 完整 TransformerEncoderLayer


class TransformerEncoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model)
        self.attn = MultiHeadSelfAttention(d_model, num_heads)
        self.norm2 = nn.LayerNorm(d_model)
        self.ffn = FeedForward(d_model, d_ff)

    def forward(self, x):
        x = x + self.attn(self.norm1(x))    # 子层 1: Attention + residual
        x = x + self.ffn(self.norm2(x))     # 子层 2: FFN + residual
        return x
        

数据流:

输入 x                                    (batch, seq_len, d_model)
  │
  ├─→ norm1 → attn → +x                    [attention 分支]
  │
  ├─→ norm2 → ffn → +x                     [FFN 分支]
  │
输出                                      (batch, seq_len, d_model)

特点: 输入输出形状永远一样,可以堆几十层。

参数量分布(d_model=128, num_heads=8, d_ff=512):

Multi-Head Attention:  66,048   (33%)
FeedForward:          131,712   (66%)  ← 大头!
LayerNorm × 2:            512   (0.3%)
────────────────────
Total:                198,272

17.10 Transformer 扩展规律

核心: 靠层数 × 单层参数量扩展。

模型层数单层参数总参数
你手写的1198K198K
BERT-base12~7M110M
GPT-2 medium24~15M345M
GPT-396~1.8B175B

从 200K → 175B,规模扩大 100 万倍。但每一层的核心逻辑,和你写的这段代码本质相同。


17.11 Transformer 阶段容易踩的坑

坑症状修复
d_k = d_model / num_headsreshape 报 float 错误用 // 整除
V = K.reshape(...) 复制粘贴 typo不报错,但结果全错逐字符检查
scores = Q @ K(少 transpose)维度不匹配报错Q @ K.transpose(-2, -1)
torch.mm(Q, K)只支持 2D用 @ 或 torch.matmul
x + self.pe 直接加广播维度不匹配切片 self.pe[:, :seq_len, :]
Attention 后忘了 FFN只做线性聚合,表达弱每层都要 attn + ffn
忘了 residual深层梯度消失每个子层都 x + sublayer(x)
忘了 super().__init__()属性赋值报错nn.Module 子类必写

17.12 Transformer 学习方法论

你走过的路,总结成 4 条方法论:

  1. 先讲思想再写代码 — Day 1-2 理解"为什么",Day 3-6 才写代码
  2. 模块独立测试 — 每个组件写完先跑独立测试,验证形状和参数量,再组合
  3. 健康检查数字 — 每个模块都有可验证的关键数字(参数量、初始输出、weights 加起来 = 1)
  4. 循环复用 — Day 6 直接 import Day 4-5 的模块,不重复实现

这些方法论适用于任何复杂系统的构建。


17.13 你现在能读懂的代码

回头看你最初的 elk_transformer_pretrain.py,现在你能读懂的部分:

# ✅ 完全看懂 —— 你自己也写过一份
class PositionalEncoding(nn.Module):
    ...

# ✅ 完全看懂 —— 这就是 6 层堆叠 EncoderLayer + Pre-LN
encoder_layer = nn.TransformerEncoderLayer(
    d_model=cfg.d_model,          # 你知道:特征维度
    nhead=cfg.nhead,              # 你知道:多头数
    dim_feedforward=cfg.dim_feedforward,  # 你知道:FFN 中间维度
    activation="gelu",            # 你知道:FFN 用 GELU
    batch_first=True,             # 你知道:(batch, seq, dim) 顺序
    norm_first=True,              # 你知道:Pre-LN
)
self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=cfg.num_layers)

# ✅ 完全看懂 —— CLS token 就是 BERT/你 Day 7 要学的分类头准备
self.cls_token = nn.Parameter(torch.zeros(1, 1, cfg.d_model))

80% 的代码你都能讲清楚每一行在做什么。剩下 20% 是 MTM(Masked Time-step Modeling)预训练细节 — 那是应用层,以后学。

从"看不懂复制的代码"到"每一行都能讲清楚",你走了大约 2 周 + 亲手写 100 行核心 Transformer。


17.14 下一站

  • Day 7: 拼装 TransformerClassifier(堆叠 + CLS + 分类头)
  • Day 8-10: 用 Transformer 训 MNIST
  • Day 11+: Vision Transformer(ViT)训 CIFAR-10
  • 未来: CLIP → VLM → VLA

Logo

DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐