""" ============================================================================= 第 6 讲 · 代码 5：在"预训练"的表示上训练奖励模型 —— 从零 / 线性探针 / 全量微调 / 先探针再微调 ============================================================================= 对应讲义：lectures/lecture06_bayes_finetuning.md 第 6 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 真实的奖励模型都是在预训练语言模型上微调出来的。这个文件在 text_world.py 的玩具世界里回答： · 预训练到底帮了什么忙？（偏好数据里没见过的说法，也能打对分） · 只训练打分头（冻结预训练的表示）还是连表示一起训练？各有什么代价？ · 为什么奖励模型通常只训练 1 个 epoch？ 演示 1：预训练学到了什么 —— 同义词的词向量靠在一起（预训练语料里没有任何"好坏"信息）。 演示 2：四种训练方式 × 三种数据量，在"同分布"和"换了说法"两种测试集上的表现。 演示 3：多训练几遍会怎样 —— 训练损失一直降，验证损失先降后升（过拟合）。 演示 4：准确率的天花板 —— 标注本身有噪声，"和标注一致的比例"到不了 100%。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 05_reward_model_finetune.py # 约 15 秒（CPU） 需要 numpy 和 PyTorch。依赖本目录的 text_world.py。 """ import numpy as np import torch import torch.nn.functional as F from bayes_tools import pad, print_header from text_world import (COMMON, RARE, SYNONYMS, VOCAB, RewardModel, TextWorld, evaluate, pretrain_embeddings, synonym_similarity, token_name, train_reward_model) EPOCHS = 10 LR = 0.01 def make_model(method, emb, seed, chosen, rejected): """ 按指定的方式训练一个奖励模型。 【参数】method： "从零训练" ：随机初始化嵌入表，嵌入表和打分头一起训练 "线性探针" ：用预训练的嵌入表，冻结它，只训练打分头（16 + 1 个参数） "全量微调" ：用预训练的嵌入表初始化，嵌入表和打分头一起训练 "先探针再微调" ：先像线性探针那样只训练打分头，再解冻嵌入表一起训练（LP-FT，Kumar 等 2022） 【返回】训练好的 RewardModel 【为什么要"先探针再微调"】全量微调一开始，打分头是随机的，它传回来的梯度也是"随机方向"的， 会把预训练的词向量往乱七八糟的方向推（只推训练数据里出现过的词！）。先把打分头训练好， 再一起微调，词向量需要改动的就少得多，没见过的词和见过的词之间的关系也就保住了。 """ torch.manual_seed(seed) if method == "从零训练": m = RewardModel() return train_reward_model(m, chosen, rejected, epochs=EPOCHS, lr=LR, seed=seed) if method == "线性探针": m = RewardModel(init_emb=emb, freeze_emb=True) return train_reward_model(m, chosen, rejected, epochs=EPOCHS, lr=LR, seed=seed) if method == "全量微调": m = RewardModel(init_emb=emb) return train_reward_model(m, chosen, rejected, epochs=EPOCHS, lr=LR, seed=seed) if method == "先探针再微调": m = RewardModel(init_emb=emb, freeze_emb=True) train_reward_model(m, chosen, rejected, epochs=EPOCHS, lr=LR, seed=seed) m.emb.weight.requires_grad_(True) # 解冻 return train_reward_model(m, chosen, rejected, epochs=EPOCHS, lr=LR, seed=seed + 1000) raise ValueError(method) # ============================================================================= # 演示 1 # ============================================================================= def demo_pretraining(world): print_header("演示 1：预训练学到了什么？") print("在 20000 个无标注的句子上训练 CBOW（遮住一个词，用其余的词预测它），过 5 遍：") emb = pretrain_embeddings(world, seed=0, epochs=5, verbose=True) same, diff = synonym_similarity(emb) print(f"\n 预训练后：同义词之间的平均余弦相似度 {same:.3f}，不同概念的词之间 {diff:.3f}") E = F.normalize(emb[:VOCAB], dim=1) S = E @ E.T print("\n 几个“换了说法”的词（训练偏好数据里从没出现过），和它最像的 3 个词：") for t in [0 * SYNONYMS + 2, 9 * SYNONYMS + 3, 6 * SYNONYMS + 2]: sims = S[t].clone() sims[t] = -2 top = torch.topk(sims, 3).indices.tolist() print(f" {token_name(t)} → " + "、".join(f"{token_name(k)}（{S[t, k]:.2f}）" for k in top)) print(""" 观察： · 预训练只看过“哪些词经常一起出现”，从没被告诉过“准确³ 和 准确¹ 是同义词”，更不知道哪个概念好、哪个坏。 可是同义词的词向量自己靠到了一起（平均余弦相似度明显高于不同概念之间）。 · 这就是“分布假设”：意思相近的词出现在相似的上下文里。真实的语言模型从万亿个词里学到的， 是这种结构的极其丰富的版本。奖励模型要利用的，正是这些已经学好的表示。""") return emb # ============================================================================= # 演示 2 # ============================================================================= def demo_methods(world, emb, reps=5): print_header("演示 2：四种训练方式 × 三种数据量") test_id = world.preference_pairs(3000, COMMON) test_ood = world.preference_pairs(3000, RARE) methods = ["从零训练", "线性探针", "全量微调", "先探针再微调"] print(f"训练数据只用“常用说法”（每个概念的同义词 ¹²）；每种设定重复 {reps} 次（不同的训练数据和初始化）取平均。") print("指标：和**真实偏好方向**一致的比例（去掉了标注噪声）。“换了说法”的测试集只用同义词 ³⁴。\n") header = f" {pad('训练对数', 9)} {pad('方式', 14)} {pad('同分布', 8, 'right')} {pad('换了说法', 9, 'right')} " \ f"{pad('可训练参数', 11, 'right')}" print(header) for n in [100, 400, 1600]: sets = [world.preference_pairs(n, COMMON) for _ in range(reps)] for method in methods: res = [] for r, (ch, rj, _) in enumerate(sets): m = make_model(method, emb, seed=r, chosen=ch, rejected=rj) res.append((evaluate(m, *test_id)["agree"], evaluate(m, *test_ood)["agree"])) res = np.array(res).mean(axis=0) n_params = 17 if method == "线性探针" else (VOCAB + 1) * 16 + 17 print(f" {n:>9d} {pad(method, 14)} {res[0]:>8.3f} {res[1]:>9.3f} {n_params:>11d}") print() print("""观察： · 从零训练：数据够多时同分布表现很好，但在“换了说法”的测试集上和瞎猜（0.5）差不多 —— 训练数据里没出现过的词， 它的词向量还是随机初始化的样子，模型对它们一无所知。 · 线性探针（冻结预训练表示，只训练 17 个参数）：“换了说法”的表现远好于从零训练 —— 没见过的词和见过的 同义词有相近的词向量，打分头对它们的打分也相近；数据到 400 对以后，它在“换了说法”上和最好的方法一样好。 代价是同分布的上限较低：预训练的表示不是为“打分”准备的，只调打分头不够灵活。 · 全量微调：同分布最好，但“换了说法”比线性探针差。原因：微调只会移动训练数据里**出现过**的词的向量， 没出现过的同义词留在原地，两者之间的关系被破坏了（“微调扭曲了预训练的特征”，Kumar 等 2022）。 · 先探针再微调（LP-FT）：两边都不错 —— 同分布接近全量微调，换了说法接近线性探针；数据只有 100 对时 两项都是最好的。先把打分头训练好，再微调时词向量只需要小幅改动。 · 对奖励模型的启示：奖励模型在 RL 阶段要给**策略生成的新回答**打分，这些回答的分布和偏好数据不同 （第 8 讲）。在“换了说法”上的表现，比同分布的准确率更能说明它在 RLHF 里靠不靠得住。""") return test_id # ============================================================================= # 演示 3 # ============================================================================= def demo_epochs(world, emb): print_header("演示 3：多训练几遍会怎样？") ch, rj, _ = world.preference_pairs(200, COMMON) val = world.preference_pairs(1000, COMMON) print("200 对训练数据（每遍 7 步），1000 对验证数据（同分布）。全量微调，学习率 0.01。\n") print(f" {pad('训练遍数', 9)} {pad('训练损失', 9, 'right')} {pad('验证损失', 9, 'right')}") torch.manual_seed(0) m = RewardModel(init_emb=emb) hist = [] train_reward_model(m, ch, rj, epochs=60, lr=LR, val=(val[0], val[1]), history=hist) for e in [1, 2, 3, 5, 10, 15, 20, 30, 40, 60]: tr, va = hist[e - 1] print(f" {e:>9d} {tr:>9.3f} {va:>9.3f}") best = int(np.argmin([h[1] for h in hist])) + 1 p = val[2] floor = float(np.mean(-(p * np.log(p) + (1 - p) * np.log(1 - p)))) print(f""" 验证损失最小在第 {best} 遍。知道真实偏好概率时，验证损失的理论下限是 {floor:.3f}（标注本身的随机性）。 观察： · 训练损失一直在降，验证损失先降后升：后面的训练是在“背”训练数据里的噪声（第 5 讲的过拟合）。 · 训练损失甚至降到了理论下限以下 —— 这是过拟合最直接的信号：模型对训练数据的“确定”程度， 超过了标注员自己的一致程度。 · 真实的奖励模型有几十亿个参数，一遍就能把训练数据记住不少，所以 InstructGPT 和 Llama 2 都只训练 1 个 epoch，并且用验证集上的准确率 / 损失决定什么时候停（第 5 讲的提前停止）。""") # ============================================================================= # 演示 4 # ============================================================================= def demo_ceiling(world, emb, test_id): print_header("演示 4：准确率的天花板 —— 标注本身有噪声") ch, rj, _ = world.preference_pairs(1600, COMMON) m = make_model("先探针再微调", emb, seed=0, chosen=ch, rejected=rj) ev = evaluate(m, *test_id) p = test_id[2] # p_true[n] 是"被选中的那个"被选中的真实概率。完美模型总是预测"真实奖励更高的那个"， # 它和标注一致的概率是 max(p, 1 − p)；对所有对取平均就是准确率的天花板。 ceiling = float(np.mean(np.maximum(p, 1 - p))) print(f"用 1600 对训练的“先探针再微调”模型，在 3000 对同分布测试数据上：\n") print(f" 和真实偏好方向一致的比例 {ev['agree']:.3f}") print(f" 和标注一致的比例（准确率） {ev['acc']:.3f}") print(f" 完美模型和标注一致的期望比例 {ceiling:.3f} ← 准确率的天花板") print(f""" 观察： · 模型和真实偏好几乎完全一致，可是“准确率”只有 {ev['acc']:.0%} 左右，已经贴着天花板了 —— 因为标注员本身就有 {1 - ceiling:.0%} 左右的时候选了真实奖励更低的那个（两个回答差不多时更常见）。准确率的天花板由数据决定， 不由模型决定（第 5 讲第 8 节）。 · 真实世界的数字：InstructGPT 报告标注员之间的一致率约 73%，奖励模型在留出标注员上的准确率约 70%； Llama 2 的奖励模型在“两个回答差别很大”的对上准确率很高，在“差不多”的对上接近 50%。 看到“奖励模型准确率只有 70%”时，先问：标注员之间的一致率是多少？""") def main(): world = TextWorld(seed=1) emb = demo_pretraining(world) test_id = demo_methods(world, emb) demo_epochs(world, emb) demo_ceiling(world, emb, test_id) if __name__ == "__main__": main()