""" ============================================================================= 第 3 讲 · 代码 5：ELECTRA——让每个位置都有训练信号，学得更快吗？ ============================================================================= 对应讲义：lectures/lecture03_architectures_objectives.md 第 5.4 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- BERT 每处理一个序列，只在被遮住的 ~15% 位置上学习（代码 03 演示 3）。ELECTRA（Clark et al. 2020）的想法： 1. 一个小的**生成器**（generator，本身就是一个小 masked LM）在被遮住的位置采样出"看起来合理"的词，填回去； 2. 一个**判别器**（discriminator）对**每一个**位置判断：这个词是原来的，还是被生成器换过的？ （replaced token detection，替换 token 检测） 3. 预训练完以后扔掉生成器，只用判别器（它就是一个 BERT 那样的双向编码器）。 演示 1：同样的数据、同样的判别器大小、同样的步数，比较三个双向编码器在训练过程中的表示质量 （线性探针：食物好不好，测试形容词从没在标注数据里出现过）： BERT（masked LM，15%） ELECTRA（判别器在 100% 的位置上算损失） ELECTRA-15%（判别器只在被遮住的那 15% 位置上算损失——论文表 5 的消融，用来区分 "收益来自所有位置都有信号"还是"来自别的地方"） 演示 2：生成器和判别器各自在做什么：生成器有多准（换了的词占多少）、判别器能不能抓出被换过的词。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 05_electra_efficiency.py # 第一次约 80 秒（训练 3 个模型并缓存）；之后约 1 秒 需要 PyTorch、numpy。缓存在 checkpoints/（l5_bert.pt 等）。 """ import os import random import sys import numpy as np import torch import torch.nn as nn import torch.nn.functional as F HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from arch_tools import (IGNORE, Corpus, WordTokenizer, count_params, few_label_accuracy, full_mask, # noqa: E402 hidden_at, load_or_train, make_masked_model, mlm_batch, train_objective) from console import pad, print_header # noqa: E402 from fm_world import (FOOD_NEG, FOOD_POS, FOODS, SERVICE_NEG, SERVICE_POS, STAFF, build_world, # noqa: E402 make_review) from tiny_gpt import GPTConfig, lr_at # noqa: E402 T, BATCH, STEPS, LR = 32, 32, 2500, 3e-3 MODEL_SEED, DATA_SEED = 0, 1 LAMBDA = 50.0 # 判别器损失的权重（ELECTRA 论文用 50：判别器的二分类损失比生成器的多分类损失小得多） CURVE_EVERY = 500 PROBE_K = 32 # ============================================================================= # 探针：训练过程中反复调用 # ============================================================================= class Probe: """ 固定的探针任务：600 条标注池（食物形容词只用前一半）、400 条测试（只用后一半），表示 = 各位置平均。 【思路】和代码 04 演示 3 完全相同的设定，这样两份代码的数字可以对照着看。 """ def __init__(self, tok): rng = random.Random(7) seen = (FOOD_POS[:4], FOOD_NEG[:4]) unseen = (FOOD_POS[4:], FOOD_NEG[4:]) pool = [make_review(rng, food_adjs=seen) for _ in range(600)] test = [make_review(rng, food_adjs=unseen) for _ in range(400)] self.pool_ids = torch.tensor([tok.encode("\n" + r[0]) for r in pool]) self.test_ids = torch.tensor([tok.encode("\n" + r[0]) for r in test]) self.yp = np.array([float(r[2]) for r in pool]) self.yt = np.array([float(r[2]) for r in test]) @torch.no_grad() def features(self, encoder, ids): """双向编码器（BERT 或 ELECTRA 判别器）在每条评论上的平均表示：(n, 64)。""" _, h = encoder(ids, full_mask(ids.shape[1]), return_hidden=True) return h.mean(1).double().numpy() def accuracy(self, encoder, k=PROBE_K, reps=10): """k 个标签的线性探针在测试集（没见过的形容词）上的准确率。""" Xp = self.features(encoder, self.pool_ids) Xt = self.features(encoder, self.test_ids) return few_label_accuracy(Xp, self.yp, Xt, self.yt, k, reps=reps, seed=k) # ============================================================================= # ELECTRA 的训练 # ============================================================================= def electra_step_batch(corpus, gen_model, B, g, rate=0.15): """ 造一个 ELECTRA 的 batch。 【思路】 1. 和 BERT 一样随机选 15% 的位置（换行符除外）；生成器的输入里，选中的位置**全部**换成 [M] （ELECTRA 的生成器不用 80/10/10）。 2. 生成器在选中的位置预测原词（这是它自己的 masked LM 损失），并从它的预测分布里**采样**一个词填回去。 3. 判别器的输入 = 填回去以后的序列；标签 = 这个位置的词和原文是否不同。 生成器碰巧采到原词时，标签算"原词"（论文也这样处理）。 【返回】(x, gen_in, gen_logits, sel, x_corrupt, labels) 【边界】采样不可导：判别器的损失不会传回生成器（论文同样如此；他们试过用强化学习训练对抗式生成器，效果更差）。 【PyTorch 小知识】torch.multinomial(probs, 1, generator=g) 按概率抽样；传入 generator 保证可复现。 """ tok = corpus.tok x = corpus.windows(B, T, g) sel = (torch.rand(B, T, generator=g) < rate) & (x != tok.nl) gen_in = torch.where(sel, torch.full_like(x, tok.mask_id), x) gen_logits = gen_model(gen_in, full_mask(T)) with torch.no_grad(): probs = F.softmax(gen_logits[..., :tok.n_regular].double(), dim=-1) # 只在普通 token 里采样 sampled = torch.multinomial(probs.view(-1, tok.n_regular), 1, generator=g).view(B, T) x_corrupt = torch.where(sel, sampled, x) labels = (x_corrupt != x).float() return x, gen_in, gen_logits, sel, x_corrupt, labels def train_electra(gen_model, disc, head, corpus, all_tokens, probe, callback_every=CURVE_EVERY): """ ELECTRA 的训练循环：生成器的 masked LM 损失 + λ × 判别器的二分类损失，一起用 AdamW 更新。 【参数】all_tokens：True = 判别器在所有位置上算损失（ELECTRA）；False = 只在被选中的 15% 位置上算（ELECTRA-15%） 【返回】dict：curve（探针准确率随步数）、gen_curve（被选中的位置里真的被换掉的比例）、history 【思路】学习率调度、梯度裁剪、batch 的取法都和代码 04 / arch_tools.train_objective 相同，保证可比。 """ tok = corpus.tok g = torch.Generator().manual_seed(DATA_SEED) params = list(gen_model.parameters()) + list(disc.parameters()) + list(head.parameters()) opt = torch.optim.AdamW(params, lr=LR, weight_decay=0.1, betas=(0.9, 0.95)) curve, gen_curve, history = [], [], [] replaced_frac = [] for step in range(STEPS + 1): if step % callback_every == 0 or step == STEPS: disc.eval() curve.append((step, probe.accuracy(disc))) if replaced_frac: gen_curve.append((step, float(np.mean(replaced_frac[-100:])))) disc.train() if step == STEPS: break for group in opt.param_groups: group["lr"] = lr_at(step, STEPS, LR) x, gen_in, gen_logits, sel, x_corrupt, labels = electra_step_batch(corpus, gen_model, BATCH, g) tgt = torch.where(sel, x, torch.full_like(x, IGNORE)) gen_loss = F.cross_entropy(gen_logits.reshape(-1, gen_logits.size(-1)), tgt.reshape(-1), ignore_index=IGNORE) _, h = disc(x_corrupt, full_mask(T), return_hidden=True) disc_logits = head(h).squeeze(-1) # (B, T)：每个位置"被换过"的打分 weight = torch.ones_like(labels) if all_tokens else sel.float() disc_loss = (F.binary_cross_entropy_with_logits(disc_logits, labels, reduction="none") * weight).sum() \ / weight.sum() loss = gen_loss + LAMBDA * disc_loss opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(params, 1.0) opt.step() replaced_frac.append((labels.sum() / sel.sum()).item()) if step % 50 == 0: history.append((step, gen_loss.item(), disc_loss.item())) for m in (gen_model, disc, head): m.eval() return dict(curve=curve, gen_curve=gen_curve, history=history) def disc_cfg(tok): """判别器（和 BERT 一样大）：2 层、64 维、4 个头。""" return GPTConfig(vocab_size=tok.vocab_size, block_size=T, n_layer=2, n_head=4, d_model=64) def gen_cfg(tok): """生成器：1 层、32 维、2 个头——比判别器小得多（论文建议生成器是判别器的 1/4 到 1/2 大小）。""" return GPTConfig(vocab_size=tok.vocab_size, block_size=T, n_layer=1, n_head=2, d_model=32) def load_models(tok, corpus, probe): """训练（或读缓存）三个模型：BERT、ELECTRA、ELECTRA-15%。""" base_meta = dict(script="05", T=T, batch=BATCH, steps=STEPS, lr=LR, model_seed=MODEL_SEED, data_seed=DATA_SEED, data="fm_world 默认世界的 review 行", probe_k=PROBE_K, version=1) out = {} def build_bert(): return {"disc": make_masked_model(disc_cfg(tok), seed=MODEL_SEED)} def train_bert(modules): curve = [] def track(step, model): curve.append((step, probe.accuracy(model))) hist = train_objective(modules["disc"], lambda B, g: mlm_batch(corpus, B, T, g), steps=STEPS, batch_size=BATCH, lr=LR, seed=DATA_SEED, callback=track, callback_every=CURVE_EVERY) return dict(curve=curve, history=hist) out["bert"] = load_or_train("l5_bert", dict(base_meta, kind="bert", cfg=vars(disc_cfg(tok))), build_bert, train_bert) for name, all_tokens in [("electra", True), ("electra15", False)]: def build(): gen_model = make_masked_model(gen_cfg(tok), seed=MODEL_SEED + 1) disc = make_masked_model(disc_cfg(tok), seed=MODEL_SEED) torch.manual_seed(MODEL_SEED + 2) return {"gen": gen_model, "disc": disc, "head": nn.Linear(64, 1)} def train(modules, all_tokens=all_tokens): return train_electra(modules["gen"], modules["disc"], modules["head"], corpus, all_tokens, probe) meta = dict(base_meta, kind=name, cfg=vars(disc_cfg(tok)), gen_cfg=vars(gen_cfg(tok)), lam=LAMBDA, all_tokens=all_tokens) out[name] = load_or_train(f"l5_{name}", meta, build, train) return out # ============================================================================= # 演示 1：学习曲线 # ============================================================================= def demo_curves(out, tok): print_header(f"演示 1：同样的判别器、同样的步数 —— 探针准确率（k = {PROBE_K}，测试形容词没见过）随训练步数的变化") names = {"bert": "BERT（masked LM 15%）", "electra": "ELECTRA（100% 位置）", "electra15": "ELECTRA-15%（只在遮住处）"} steps = [s for s, _ in out["bert"][1]["curve"]] print(pad("模型", 28) + "".join(pad(f"{s} 步", 9, ">") for s in steps)) for key, label in names.items(): accs = dict(out[key][1]["curve"]) print(pad(label, 28) + "".join(pad(f"{accs[s]:.3f}", 9, ">") for s in steps)) n_disc = count_params(out["bert"][0]["disc"]) n_gen = count_params(out["electra"][0]["gen"]) print(f"\n判别器（= BERT）{n_disc:,} 个参数；ELECTRA 另有生成器 {n_gen:,} 个参数（每步多约 {n_gen / n_disc:.0%} 的计算，" f"训练完就扔掉）和一个 65 个参数的判别头。") t = [f"{names[k]} {out[k][1]['train_seconds']:.0f} 秒" for k in names if not out[k][1]["from_cache"]] if t: print("（训练用时：" + "；".join(t) + "）") print(""" 观察 Observations： · BERT 的探针在前 1000 步停在 0.5 左右，之后上升，2000 步时约 0.9（2500 步回落到 0.8：探针只用 32 个标签， 表示本身也还在变化）。这和代码 04 演示 1 的"平台期之后突然学会"一致——两份代码里的 BERT 是同一个模型。 · ELECTRA 和 ELECTRA-15% 在整个训练过程中都停在 0.5 左右，也就是瞎猜。**在这个迷你世界、这个算力下， ELECTRA 没有表现出论文里的样本效率优势，反而不如 BERT。**原本想用 ELECTRA-15% 拆分"收益来自哪里"， 现在两者都没学到东西，这个对照也就无从谈起。演示 2 看看问题出在哪里。 · 论文里的数字（同样的算力）：ELECTRA-Small 在 GLUE 上 79.9，BERT-Small 75.1；消融实验（论文表 5）里 ELECTRA 85.0、只在 15% 位置上算判别损失的 ELECTRA-15% 82.4、BERT 82.2——收益主要来自"所有位置都有信号"。""") # ============================================================================= # 演示 2：生成器和判别器在做什么 # ============================================================================= @torch.no_grad() def demo_inside(out, tok, corpus): print_header("演示 2：生成器换掉了多少词？判别器抓得出来吗？") mods = out["electra"][0] gen_curve = out["electra"][1]["gen_curve"] print("被选中的 15% 位置里，生成器采样出的词和原词不同（真的\"换掉了\"）的比例（最近 100 步的平均）：") print(" " + " ".join(f"{s} 步 {f:.2f}" for s, f in gen_curve)) g = torch.Generator().manual_seed(123) x, gen_in, gen_logits, sel, x_corrupt, labels = electra_step_batch(corpus, mods["gen"], 256, g) _, h = mods["disc"](x_corrupt, full_mask(T), return_hidden=True) probs = torch.sigmoid(mods["head"](h).squeeze(-1)) pred = (probs > 0.5).float() rep = labels.bool() n_pos = x.numel() print(f"\n训练好的 ELECTRA，在 256 段新采样的输入上（{n_pos:,} 个位置）：被选中 {int(sel.sum())} 个，" f"真正被换掉 {int(rep.sum())} 个（占全部位置的 {rep.float().mean().item():.1%}）。") print(f" 判别器：被换掉的词里抓出 {pred[rep].mean().item():.2f}；原词里误报 {pred[~rep].mean().item():.3f}") # 把被换掉的词分成三类：换成同类、意思没变的词 / 换成意思相反的词 / 换成别的（多半语法上就不对） pos = set(FOOD_POS + SERVICE_POS) neg = set(FOOD_NEG + SERVICE_NEG) groups = [set(FOODS), set(STAFF), {"good", "ok", "bad"}] def category(o, n): if (o in pos and n in neg) or (o in neg and n in pos) or {o, n} == {"and", "but"} \ or ({o, n} <= groups[2]): return "意思变了（正负翻转 / and↔but / 评分）" if (o in pos and n in pos) or (o in neg and n in neg) or any(o in s and n in s for s in groups[:2]): return "同类替换，意思没变（同正负的形容词、食物、员工）" return "其他（多半是词性都不对的替换）" stats = {} for b, t in rep.nonzero().tolist(): c = category(tok.vocab[x[b, t]], tok.vocab[x_corrupt[b, t]]) n, caught = stats.get(c, (0, 0)) stats[c] = (n + 1, caught + int(pred[b, t].item())) print("\n 被换掉的词分三类：") for c in sorted(stats, key=lambda k: -stats[k][0]): n, caught = stats[c] print(f" {pad(c, 50)} {n:>4} 个（{n / rep.sum().item():.0%}），判别器抓出 {caught / n:.2f}") shown = 0 print("\n 几个例子（原词 → 换上的词：判别器认为\"被换过\"的概率；后面是换过之后的上下文）：") for b in range(x.shape[0]): for t in range(1, T - 1): if rep[b, t] and shown < 5: ctx = " ".join(tok.vocab[i].replace("\n", "⏎") for i in x_corrupt[b, max(0, t - 5):t + 6].tolist()) print(f" {tok.vocab[x[b, t]]} → {tok.vocab[x_corrupt[b, t]]}：{probs[b, t].item():.2f} …{ctx}…") shown += 1 print(""" 观察 Observations： · 生成器学得很快：被选中的位置里真被换掉的比例从 0.92（第 500 步）降到 0.29（第 2500 步）。 训练后期只有约 4% 的位置被换掉，判别器看到的绝大多数标签都是"原词"。 · 被换掉的词里，约一半是**同类替换**（waiter → cook、bread → pie、同为正面的两个形容词）——意思没变， 原则上就无法分辨；另一半**改变了意思**（例如句子里是 "sour and the waiter was polite. rating: bad"， polite 显然是被换过的），要发现它们，判别器得做"形容词正负 + and/but + 评分"的推理； 只有约 6% 是词性都不对的明显替换。 · 判别器只学会了抓明显的替换（约 0.8），对改变意思的替换几乎全漏，对原词从不误报：它停在"语法检查器"的阶段， 没学到探针需要的正负知识。真正有信息量的标签（改变了意思的替换）只占全部位置的约 2%——比 BERT 的约 14% 还少。 · 教训：ELECTRA 的"每个位置都有信号"，只有当"判断换没换"**需要**我们想让编码器学的知识时才有用； 它的效果取决于生成器犯什么样的错，而这又取决于数据。真实文本比我们的模板丰富得多，生成器的错误五花八门， 判别器能从中学到大量语言知识。所以这个玩具实验既不能证明、也不能推翻论文的结论—— 它说明的是一个机制上的前提条件。""") def summarize(out): print_header("小结：探针准确率（k = 32，测试形容词没见过；0.5 = 瞎猜）") names = {"bert": "BERT", "electra": "ELECTRA", "electra15": "ELECTRA-15%"} for key, label in names.items(): curve = out[key][1]["curve"] best_step, best = max(curve, key=lambda sa: sa[1]) print(f" {pad(label, 12)} 训练过程中最高 {best:.3f}（第 {best_step} 步）；训练结束时 {curve[-1][1]:.3f}") if __name__ == "__main__": world = build_world() tok = WordTokenizer(world.train_lines) review_lines = [l for l in world.train_lines if world.line_genres[l] == "review"] corpus = Corpus(review_lines, tok) probe = Probe(tok) print(f"训练数据：和代码 04 相同的 {len(review_lines):,} 条评论（{len(corpus.ids_list):,} 个词级 token）。") out = load_models(tok, corpus, probe) demo_curves(out, tok) demo_inside(out, tok, corpus) summarize(out)