""" ============================================================================= 第 7 讲 · 代码 3：多头注意力与 Transformer 块 ============================================================================= 对应讲义：lectures/lecture07_self_attention_transformer.md 第 6、7、8 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- transformer.py 实现了完整的 Transformer。这个文件逐个检查、解剖它的零件： 演示 1：多头注意力里张量形状的变化（每一步打印形状） 演示 2：我们的实现和 PyTorch 自带的 nn.MultiheadAttention 结果一致 演示 3：为什么要"多头"？—— 一个需要同时看三个位置的小任务：头数不够时做不到 演示 4：层归一化在算什么（手算 vs nn.LayerNorm） 演示 5：参数量公式：每个块 12d² + 13d；用同样的代码搭一个 GPT-2 small，参数恰好是 124,439,808 演示 6：残差连接的作用：8 层的网络去掉残差就完全学不动（表示"坍缩"）；Pre-LN vs Post-LN ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 03_multihead_and_block.py 需要 PyTorch，约 1~2 分钟（演示 3、6 要训练十几个小模型，演示 5 要创建一个 1.2 亿参数的模型）。 """ import os os.environ.setdefault("OMP_NUM_THREADS", "1") import math import time import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from transformer import (MultiHeadSelfAttention, TransformerBlock, TransformerLM, count_parameters) torch.set_num_threads(1) # 小模型单线程最快（见第 6 讲 04 号文件的说明） def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 演示 1：形状 # ============================================================================= def demo_shapes(): """把 MultiHeadSelfAttention.forward 的每一步重新走一遍，打印张量形状。""" print_header("演示 1：多头注意力里张量形状的变化（B=2 个句子，T=5 个词，D=16，H=4 个头）") torch.manual_seed(0) B, T, D, H = 2, 5, 16, 4 Dh = D // H mha = MultiHeadSelfAttention(D, H, causal=True) x = torch.randn(B, T, D) steps = [] qkv = mha.qkv(x) steps.append(("qkv = 线性层(x)", qkv.shape, "一次算出 Q、K、V，拼在最后一维")) q, k, v = qkv.split(D, dim=2) steps.append(("q, k, v = split", q.shape, "各自 (B, T, D)")) q = q.view(B, T, H, Dh) steps.append(("q.view(B, T, H, Dh)", q.shape, f"把 D={D} 拆成 H={H} 个头 × Dh={Dh}")) q = q.transpose(1, 2) k = k.view(B, T, H, Dh).transpose(1, 2) v = v.view(B, T, H, Dh).transpose(1, 2) steps.append(("q.transpose(1, 2)", q.shape, "把\"头\"挪到前面：之后每个头独立计算")) scores = q @ k.transpose(-2, -1) / math.sqrt(Dh) steps.append(("scores = q @ kᵀ / √Dh", scores.shape, "每个头一个 T×T 的打分矩阵")) attn = torch.softmax(scores.masked_fill(torch.triu(torch.ones(T, T, dtype=torch.bool), 1), -math.inf), -1) steps.append(("attn = softmax(掩码后的 scores)", attn.shape, "每一行是一个概率分布")) out = attn @ v steps.append(("out = attn @ v", out.shape, "每个头的输出")) out = out.transpose(1, 2).contiguous().view(B, T, D) steps.append(("拼接各个头", out.shape, "(B, T, H, Dh) → (B, T, D)")) out = mha.proj(out) steps.append(("proj（输出线性层 W^O）", out.shape, "把各个头的信息混合起来")) print() for name, shape, note in steps: print(f" {name:<32s} {str(tuple(shape)):<16s} {note}") assert torch.allclose(out, mha(x), atol=1e-6) print("\n 自我检查：手动走一遍的结果和 mha(x) 完全一致 ✓") # ============================================================================= # 演示 2：和 nn.MultiheadAttention 对比 # ============================================================================= def demo_vs_pytorch(): """把我们的参数复制到 nn.MultiheadAttention 里，比较输出和注意力权重。""" print_header("演示 2：和 PyTorch 自带的 nn.MultiheadAttention 对比") torch.manual_seed(0) B, T, D, H = 3, 7, 32, 4 x = torch.randn(B, T, D) for causal in [False, True]: ours = MultiHeadSelfAttention(D, H, causal=causal) ref = nn.MultiheadAttention(D, H, batch_first=True) with torch.no_grad(): # nn.MultiheadAttention 把 W_Q、W_K、W_V 存成一个 (3D, D) 的矩阵，顺序和我们的 qkv 层相同 ref.in_proj_weight.copy_(ours.qkv.weight) ref.in_proj_bias.copy_(ours.qkv.bias) ref.out_proj.weight.copy_(ours.proj.weight) ref.out_proj.bias.copy_(ours.proj.bias) # PyTorch 的 attn_mask：True 表示"不允许关注"（和我们的 future 矩阵含义相同） mask = torch.triu(torch.ones(T, T, dtype=torch.bool), 1) if causal else None y_ref, w_ref = ref(x, x, x, attn_mask=mask, need_weights=True, average_attn_weights=False) y_ours, w_ours = ours(x, return_attn=True) print(f"\n {'因果' if causal else '双向'}注意力：输出最大差 {(y_ref - y_ours).abs().max():.1e}，" f"注意力权重最大差 {(w_ref - w_ours).abs().max():.1e}") assert torch.allclose(y_ref, y_ours, atol=1e-5) and torch.allclose(w_ref, w_ours, atol=1e-5) print("\n 自我检查：两种实现完全一致（差别只是浮点误差）✓") # ============================================================================= # 演示 3：为什么要多头 # ============================================================================= class OneLayerAttention(nn.Module): """ 只有一层注意力的小模型（没有前馈网络），用来做演示 3 的"同时取回几个位置"任务。 【结构】词嵌入 + 可学习位置嵌入 → 一层因果多头注意力（带残差）→ 每个目标一个输出层 【参数】vocab（词表大小）, D（总维度）, H（头数）, T（序列长度）, n_out（目标个数） """ def __init__(self, vocab, D, H, T, n_out): super().__init__() self.tok = nn.Embedding(vocab, D) self.pos = nn.Embedding(T, D) self.attn = MultiHeadSelfAttention(D, H, causal=True) self.outs = nn.ModuleList([nn.Linear(D, vocab) for _ in range(n_out)]) def forward(self, X): h = self.tok(X) + self.pos(torch.arange(X.shape[1])) h = h + self.attn(h) return [out(h) for out in self.outs] def train_copy_task(H, offsets=(1, 2, 4), D=24, steps=1500, vocab=10, T=12, seed=0): """ 训练演示 3 的任务：在每个位置 t，同时输出 x[t-1]、x[t-2]、x[t-4]（offsets 指定往前看几个位置）。 【思路】 一个注意力头在每个位置只有**一个**概率分布。要同时取回几个不同位置的词， 单个头只能把权重分给这几个位置，得到它们的词向量的**加权平均** —— 平均之后就分不清哪个词在哪个位置了。 多个头各自看一个位置，再拼接起来，每个词的信息就分开保存了。 为了公平，所有模型的总维度 D 相同（头越多，每个头的维度越小），参数量完全相同。 【参数】 H : 头的个数（必须能整除 D） offsets : 要取回的相对位置 D, steps, vocab, T : 模型维度、训练步数、词表大小、序列长度 seed : 随机种子 【返回】(每个目标的准确率列表, 参数量)，在 2000 个新序列上评估（只算 t ≥ max(offsets) 的位置） """ torch.manual_seed(seed) model = OneLayerAttention(vocab, D, H, T, len(offsets)) opt = torch.optim.Adam(model.parameters(), lr=3e-3) s0 = max(offsets) for _ in range(steps): X = torch.randint(0, vocab, (64, T)) outs = model(X) # 第 t 个位置的目标是 X[:, t-k]：位置 s0..T-1 对应目标 X[:, s0-k .. T-1-k] loss = sum(F.cross_entropy(o[:, s0:].reshape(-1, vocab), X[:, s0 - k:T - k].reshape(-1)) for o, k in zip(outs, offsets)) opt.zero_grad() loss.backward() opt.step() with torch.no_grad(): X = torch.randint(0, vocab, (2000, T)) accs = [(o[:, s0:].argmax(-1) == X[:, s0 - k:T - k]).float().mean().item() for o, k in zip(model(X), offsets)] return accs, count_parameters(model) def demo_why_multihead(): print_header("演示 3：为什么要多头？—— 同时取回前 1、前 2、前 4 个位置的词") print(""" 任务：随机的词序列（词表 10 个词，瞎猜的准确率是 10%），在每个位置 t 同时预测 x[t-1]、x[t-2]、x[t-4]。 模型：只有一层注意力，总维度 24，比较 1、2、3、6 个头（每个头 24、12、8、4 维）。每种 3 个随机种子。""") t0 = time.perf_counter() print(f"\n {'头数':<6s}{'x[t-1]':>9s}{'x[t-2]':>9s}{'x[t-4]':>9s}{'最差的目标':>12s}{'参数量':>9s}") for H in [1, 2, 3, 6]: runs = [train_copy_task(H, seed=s) for s in range(3)] accs = np.array([r[0] for r in runs]) # (种子数, 目标数) worst = accs.min(axis=1).mean() # 每个种子最差的那个目标，再平均 print(f" {H:<8d}" + "".join(f"{a:>9.1%}" for a in accs.mean(axis=0)) + f"{worst:>12.1%}{runs[0][1]:>10,}") print(f"\n （用时 {time.perf_counter() - t0:.0f} 秒）") print(""" 观察：参数量完全相同。 - 1 个头：总有目标做不好（"最差的目标"接近瞎猜）。它在每个位置只有一个注意力分布： 要么只盯一个位置（其他目标只能靠猜），要么把权重分给几个位置，得到词向量的平均 —— 平均会丢掉"哪个词在哪个位置"的信息。 - 2 个头：三个目标都不错，但没有全对（两个头要分管三个位置）。 - 3 个头、6 个头：全对。头的个数够用之后，每个头负责一个位置，再把结果拼起来。 多个头 = 多个独立的注意力分布，每个头可以负责一种关系（比如一个头找主语、一个头找宾语）。 真实模型里确实能找到这样"分工"的头（讲义第 6.2 节；04 号文件演示 4）。""") # ============================================================================= # 演示 4：层归一化 # ============================================================================= def demo_layernorm(): print_header("演示 4：层归一化在算什么？") torch.manual_seed(0) x = torch.tensor([[2.0, 4.0, 6.0, 8.0], [100.0, 101.0, 102.0, 103.0]]) ln = nn.LayerNorm(4) with torch.no_grad(): ln.weight.copy_(torch.tensor([1.0, 1.0, 2.0, 2.0])) ln.bias.copy_(torch.tensor([0.0, 0.0, 0.0, 1.0])) mu = x.mean(dim=-1, keepdim=True) var = x.var(dim=-1, unbiased=False, keepdim=True) # 注意：用的是"除以 n"的方差 manual = (x - mu) / torch.sqrt(var + ln.eps) * ln.weight + ln.bias print(f"\n 输入 x：\n{x.numpy()}") print(f" 每行的均值 μ = {mu.squeeze().numpy()}，方差 σ² = {var.squeeze().numpy()}") print(f" 标准化 (x - μ)/σ：\n{((x - mu) / torch.sqrt(var + ln.eps)).numpy().round(3)}") print(f" 再乘 γ = {ln.weight.detach().numpy()}，加 β = {ln.bias.detach().numpy()}：\n{manual.detach().numpy().round(3)}") assert torch.allclose(manual, ln(x), atol=1e-5) print(""" 自我检查：手算和 nn.LayerNorm 一致 ✓ 观察： - 层归一化对**每个位置的向量**单独做（沿特征维度求均值和方差），和批里的其他句子、其他位置无关。 这和第 3 讲的批归一化（BatchNorm，沿批维度统计）不同：句子长短不一、批里的句子随时在变， 按批统计不稳定，所以 NLP 模型几乎都用层归一化。 - 两行输入的"量级"差了 25 倍，标准化之后完全一样 —— 后面的层看到的数值范围始终稳定。 - γ（缩放）和 β（偏移）是可学习的，模型需要的话可以把标准化"撤销"掉。""") # ============================================================================= # 演示 5：参数量 # ============================================================================= def demo_param_count(): print_header("演示 5：参数量公式") print(""" 一个块（Pre-LN）的参数： W_Q、W_K、W_V（合成一个 D→3D 层） 3D² + 3D 输出投影 W^O D² + D 前馈网络 D→4D→D (4D² + 4D) + (4D² + D) 两个层归一化（各有 γ、β） 4D 合计 12D² + 13D""") for D in [64, 128, 768]: block = TransformerBlock(D, n_heads=4) print(f" D = {D:>4d}：公式 {12 * D * D + 13 * D:>10,}，实际 {count_parameters(block):>10,}") assert count_parameters(block) == 12 * D * D + 13 * D print(""" 整个模型 = 词嵌入 V·D + 位置嵌入 max_len·D + L 个块 + 最后的层归一化 2D （输出层和词嵌入共用同一个矩阵，不再额外算参数）。 用同样的代码搭一个和 GPT-2 small 一样大的模型：V = 50257，max_len = 1024，D = 768，12 层，12 个头 ……""") V, ctx, D, L = 50257, 1024, 768, 12 formula = V * D + ctx * D + L * (12 * D * D + 13 * D) + 2 * D t0 = time.perf_counter() gpt2 = TransformerLM(V, d_model=D, n_heads=12, n_layers=L, max_len=ctx) actual = count_parameters(gpt2) print(f"\n 公式 {formula:,}，实际 {actual:,}（创建用时 {time.perf_counter() - t0:.1f} 秒）") print(f" 其中词嵌入占 {V * D / actual:.1%}，12 个块占 {L * (12 * D * D + 13 * D) / actual:.1%}") assert actual == formula == 124_439_808 del gpt2 print(""" 自我检查：恰好是 OpenAI 公布的 GPT-2 small 参数量 124,439,808 ✓ —— 我们的结构和 GPT-2 完全相同，只是大小不同。 观察：小模型里词嵌入占了很大比例；模型越大，块的参数（∝ L·D²）越占主导。 计算参数量的这套方法，第 9 讲估算训练成本时还会用到。""") # ============================================================================= # 演示 6：残差连接与 Pre-LN / Post-LN # ============================================================================= class NoResidualBlock(TransformerBlock): """去掉残差连接的块：x = Attention(LN(x))；x = FFN(LN(x))。只用于演示 6。""" def forward(self, x, pad_mask=None, return_attn=False): x = self.attn(self.ln1(x), pad_mask) return self.ffn(self.ln2(x)) class DeepModel(nn.Module): """ 演示 6 用的深层模型：词嵌入 + 位置嵌入 → n_layers 个块 → 层归一化 → 输出层。 【参数】 kind : "pre_ln"（Pre-LN，默认写法）、"post_ln"（原论文的 Post-LN）、"no_residual"（Pre-LN 去掉残差） n_layers : 块的个数 D, vocab, T : 维度、词表大小、序列长度 参数使用 PyTorch 的默认初始化。 """ def __init__(self, kind, n_layers, D=64, vocab=10, T=12): super().__init__() self.tok, self.pos = nn.Embedding(vocab, D), nn.Embedding(T, D) if kind == "no_residual": self.blocks = nn.ModuleList([NoResidualBlock(D, 4) for _ in range(n_layers)]) else: self.blocks = nn.ModuleList([TransformerBlock(D, 4, pre_ln=(kind == "pre_ln")) for _ in range(n_layers)]) self.ln, self.out = nn.LayerNorm(D), nn.Linear(D, vocab) def hidden(self, X): """返回最后一个块的输出 (B, T, D)。""" h = self.tok(X) + self.pos(torch.arange(X.shape[1])) for b in self.blocks: h = b(h) return h def forward(self, X): return self.out(self.ln(self.hidden(X))) def collapse_score(model, T=12): """ 初始化时，最后一层不同位置的表示之间的平均余弦相似度。 【思路】 如果网络把不同位置、不同的词都变成了几乎同一个向量（余弦相似度接近 1）， 后面的层就再也分不清"哪个位置是哪个词"了。Dong 等（2021）证明： 只有注意力、没有残差连接的网络，随着层数增加会以极快的速度"坍缩"到这种状态（rank collapse）。 【返回】float，-1~1，越接近 1 说明表示越"坍缩" """ with torch.no_grad(): h = F.normalize(model.hidden(torch.randint(0, 10, (16, T))), dim=-1) sim = h @ h.transpose(1, 2) # (16, T, T) 两两余弦相似度 return sim[:, ~torch.eye(T, dtype=torch.bool)].mean().item() # 去掉对角线（自己和自己） def train_deep(kind, n_layers=8, steps=300, lr=1e-3, seed=0): """ 在"预测 x[t-3]"的任务上训练一个深层模型，返回 (初始化时的坍缩程度, 前 50 步平均损失, 最后 50 步平均损失, 准确率)。 【边界】瞎猜的损失是 ln 10 ≈ 2.30，准确率 10%。 """ torch.manual_seed(seed) model = DeepModel(kind, n_layers) score = collapse_score(model) opt = torch.optim.Adam(model.parameters(), lr=lr) losses = [] for _ in range(steps): X = torch.randint(0, 10, (64, 12)) loss = F.cross_entropy(model(X)[:, 3:].reshape(-1, 10), X[:, :-3].reshape(-1)) opt.zero_grad() loss.backward() opt.step() losses.append(loss.item()) with torch.no_grad(): X = torch.randint(0, 10, (2000, 12)) acc = (model(X)[:, 3:].argmax(-1) == X[:, :-3]).float().mean().item() return score, np.mean(losses[:50]), np.mean(losses[-50:]), acc def demo_residual(): print_header("演示 6：残差连接的作用 —— 训练一个 8 层的网络") print(""" 任务：在每个位置 t 预测 x[t-3]（随机序列，词表 10 个词；瞎猜的损失 2.30、准确率 10%）。 模型：8 个块，D = 64，Adam，学习率 1e-3，训练 300 步。""") names = {"pre_ln": "Pre-LN（有残差）", "post_ln": "Post-LN（有残差）", "no_residual": "Pre-LN 去掉残差"} print(f"\n {'':<14s}{'初始化时的坍缩程度':>12s}{'前 50 步损失':>12s}{'最后 50 步损失':>13s}{'准确率':>8s}") t0 = time.perf_counter() for kind in ["pre_ln", "post_ln", "no_residual"]: score, first, last, acc = train_deep(kind) print(f" {names[kind]:<12s}{score:>14.3f}{first:>14.3f}{last:>14.3f}{acc:>11.1%}") print(f"\n （用时 {time.perf_counter() - t0:.0f} 秒；\"坍缩程度\" = 最后一层不同位置的表示之间的平均余弦相似度）") print(""" 观察： - 去掉残差连接，8 层的网络**完全学不动**：损失一直停在瞎猜的水平。 原因看"坍缩程度"：还没开始训练，最后一层所有位置的表示就已经几乎是同一个向量了（余弦相似度 ≈ 1）。 每一层注意力都在对各个位置做加权平均，一层层平均下去，不同位置的差别被抹平了（Dong 等 2021：rank collapse）。 输入是什么词、在哪个位置，到最后一层已经看不出来，梯度也就不知道该往哪个方向改。 - 有残差连接：x + f(x) 保留了一条"原样传过去"的通路，每个位置自己的信息不会被平均掉； 反向传播时 x + f(x) 对 x 的导数是 I + ∂f/∂x，梯度也有一条直达的路径（第 6 讲第 4.1 节）。 - 在这个小任务上 Pre-LN 和 Post-LN 都能训练。Post-LN（原始 Transformer 的写法）的问题在更深、更大的模型上才明显： Xiong 等（2020）证明它在初始化时靠近输出的层梯度偏大，必须用学习率预热，否则训练不稳定； Pre-LN 更稳，所以 GPT-2 以后的模型大多改用 Pre-LN（讲义第 7.4 节）。""") def main(): demo_shapes() demo_vs_pytorch() demo_why_multihead() demo_layernorm() demo_param_count() demo_residual() if __name__ == "__main__": main()