VLA学习之路
十七、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 个致命问题:
| 问题 | RNN | Transformer |
|---|---|---|
| 长距离依赖 | 梯度沿 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_model | num_heads | d_k |
|---|---|---|---|
| BERT-base | 768 | 12 | 64 |
| GPT-3 | 12288 | 96 | 128 |
| Llama-7B | 4096 | 32 | 128 |
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 的对比:
| BatchNorm | LayerNorm | |
|---|---|---|
| 归一化范围 | 同一 batch 里所有样本的同一维度 | 同一样本内部所有维度 |
| 对 batch 依赖 | 强(小 batch 就崩) | 无 |
| 常用场景 | CNN | Transformer / 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 扩展规律
核心: 靠层数 × 单层参数量扩展。
| 模型 | 层数 | 单层参数 | 总参数 |
|---|---|---|---|
| 你手写的 | 1 | 198K | 198K |
| BERT-base | 12 | ~7M | 110M |
| GPT-2 medium | 24 | ~15M | 345M |
| GPT-3 | 96 | ~1.8B | 175B |
从 200K → 175B,规模扩大 100 万倍。但每一层的核心逻辑,和你写的这段代码本质相同。
17.11 Transformer 阶段容易踩的坑
| 坑 | 症状 | 修复 |
|---|---|---|
d_k = d_model / num_heads | reshape 报 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 条方法论:
- 先讲思想再写代码 — Day 1-2 理解"为什么",Day 3-6 才写代码
- 模块独立测试 — 每个组件写完先跑独立测试,验证形状和参数量,再组合
- 健康检查数字 — 每个模块都有可验证的关键数字(参数量、初始输出、weights 加起来 = 1)
- 循环复用 — 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
DAMO开发者矩阵,由阿里巴巴达摩院和中国互联网协会联合发起,致力于探讨最前沿的技术趋势与应用成果,搭建高质量的交流与分享平台,推动技术创新与产业应用链接,围绕“人工智能与新型计算”构建开放共享的开发者生态。
更多推荐
所有评论(0)