""" ============================================================================= 第 6 讲 · 代码 6：奖励模型集成 —— 一个实用的"不确定性" ============================================================================= 对应讲义：lectures/lecture06_bayes_finetuning.md 第 7 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 代码 01~04 的贝叶斯方法在几个、几十个参数时很好用；神经网络奖励模型有几十亿个参数，MCMC 跑不动。 实践中最常用的"不确定性"是**集成**：用不同的随机种子（和 / 或不同的数据子集）训练几个奖励模型， 看它们的意见是否一致。意见不一致的地方，就是"不确定"的地方。 演示 1：三种集成（从零训练 / 共享同一个预训练 / 各自预训练）：成员之间的分歧能不能识别出 "没见过的输入"？能不能预测"哪些对会判错"？ 演示 2：Best-of-n（第 1、8 讲）：候选越多，单个奖励模型越容易挑中它"高估"的回答； 集成平均、"平均分 − 标准差"（保守）能不能缓解？ ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 06_reward_ensemble.py # 约 15 秒（CPU） 需要 numpy 和 PyTorch。依赖本目录的 text_world.py。 """ import numpy as np import torch from bayes_tools import pad, print_header from text_world import (COMMON, RARE, RewardModel, TextWorld, pad_batch, pretrain_embeddings, train_reward_model, true_reward) MEMBERS = 5 EPOCHS = 10 LR = 0.01 def auc_score(score, label): """ AUC：随机取一个正例（label = 1）和一个负例，正例的 score 更高的概率（第 5 讲）。 【用在这里】score = 集成的分歧，label = "集成判错了"。AUC = 0.5 表示分歧完全预测不了错误，1 表示完美预测。 """ score = np.asarray(score, dtype=float) label = np.asarray(label, dtype=bool) pos, neg = score[label], score[~label] if len(pos) == 0 or len(neg) == 0: return float("nan") # 比较所有 (正例, 负例) 对；相等算一半 diff = pos[:, None] - neg[None, :] return float(np.mean((diff > 0) + 0.5 * (diff == 0))) def build_ensemble(kind, world, chosen, rejected, shared_emb): """ 训练一个由 MEMBERS 个奖励模型组成的集成。 【参数】kind： "从零训练" ：每个成员随机初始化（不同种子），在自助抽样的数据上训练 "共享预训练 + 微调" ：所有成员从**同一份**预训练词向量出发，"先探针再微调"，种子和自助样本不同 "各自预训练 + 微调" ：每个成员用**自己的**预训练词向量（预训练的随机种子不同），再"先探针再微调" 【返回】成员（RewardModel）的列表 【为什么用自助抽样】让每个成员看到略有不同的数据，成员之间的分歧就包含了"数据换一批会怎样" 的不确定性（和第 5 讲的自助法同一个思想）。 """ members = [] for k in range(MEMBERS): torch.manual_seed(100 + k) if kind == "从零训练": m = RewardModel() train_reward_model(m, chosen, rejected, epochs=EPOCHS, lr=LR, seed=k, bootstrap=True) else: emb = shared_emb if kind == "共享预训练 + 微调" else pretrain_embeddings(world, seed=10 + k, epochs=5) m = RewardModel(init_emb=emb, freeze_emb=True) train_reward_model(m, chosen, rejected, epochs=EPOCHS, lr=LR, seed=k, bootstrap=True) m.emb.weight.requires_grad_(True) train_reward_model(m, chosen, rejected, epochs=EPOCHS, lr=LR, seed=k + 1000, bootstrap=True) members.append(m) return members @torch.no_grad() def member_rewards(members, responses): """每个成员给每个回答打分。返回形状 (成员数, 回答数) 的数组。""" tok, mask = pad_batch(responses) return np.stack([m(tok, mask).numpy() for m in members]) def pair_uncertainty(members, chosen, rejected): """ 对一组比较，算出集成的平均判断和几种"不确定性"。 【返回】dict（每个都是长度 = 对数的数组）： margins ：形状 (成员数, 对数)，每个成员给出的"分数差" r(被选中的) − r(被拒绝的) mean ：成员平均的分数差（集成的判断：> 0 表示认为被选中的更好） sd ：成员之间分数差的标准差（原始分歧） rel ：相对分歧 = sd / |mean|。分歧要和"判断有多强"放在一起看：平均差 3、分歧 0.5 的对很稳； 平均差 0.1、分歧 0.5 的对，换一个成员就可能判反。 """ margins = member_rewards(members, chosen) - member_rewards(members, rejected) mean = margins.mean(axis=0) sd = margins.std(axis=0) return dict(margins=margins, mean=mean, sd=sd, rel=sd / (np.abs(mean) + 1e-9)) # ============================================================================= # 演示 1 # ============================================================================= def demo_disagreement(world, ensembles): print_header("演示 1：集成成员之间的分歧 —— 能不能发现\"没见过的输入\"和\"会判错的对\"？") test_id = world.preference_pairs(2000, COMMON) test_ood = world.preference_pairs(2000, RARE) chosen = test_id[0] + test_ood[0] rejected = test_id[1] + test_ood[1] truly = np.concatenate([test_id[2], test_ood[2]]) > 0.5 is_ood = np.r_[np.zeros(2000), np.ones(2000)].astype(bool) print(f"每个集成 {MEMBERS} 个成员，各自在 400 对训练数据的自助样本上训练。测试集：2000 对同分布 + 2000 对换了说法，混在一起。") print("“判错” = 集成平均的判断和真实偏好方向相反。后四列都是 AUC（0.5 = 毫无用处，1 = 完美）。\n") print(f" {pad('集成', 20)} {pad('判错率', 15, 'right')} {pad('相对分歧', 9, 'right')} " f"{pad('预测判错：', 11, 'right')} {pad('预测判错：', 11, 'right')} {pad('预测判错：', 11, 'right')}") print(f" {pad('', 20)} {pad('同分布/换说法', 15, 'right')} {pad('→换了说法', 9, 'right')} " f"{pad('单模型信心', 11, 'right')} {pad('相对分歧', 11, 'right')} {pad('原始标准差', 11, 'right')}") for kind, members in ensembles.items(): u = pair_uncertainty(members, chosen, rejected) wrong = (u["mean"] > 0) != truly wrong0 = (u["margins"][0] > 0) != truly # 第一个成员自己判错了没有 rates = f"{wrong[~is_ood].mean():.3f} / {wrong[is_ood].mean():.3f}" print(f" {pad(kind, 20)} {rates:>15s} {auc_score(u['rel'], is_ood):>9.3f} " f"{auc_score(-np.abs(u['margins'][0]), wrong0):>11.3f} {auc_score(u['rel'], wrong):>11.3f} " f"{auc_score(u['sd'], wrong):>11.3f}") print(""" （“单模型信心”：只用一个成员，它给出的分数差越接近 0，就越可能判错 —— 不需要集成也能算。） 观察： · 从零训练的集成：“换了说法”的对判错率很高（接近瞎猜），但相对分歧能很好地把这些对挑出来（第 3 列）—— 没见过的词，每个成员的随机初始化都不一样，意见自然不一致。集成**知道**自己在这里不可靠； 单个模型做不到这一点（它对没见过的输入也会给出一个看起来正常的分数差），所以集成预测判错的能力明显更强。 · 共享同一个预训练的集成：成员们在“换了说法”上判断得好得多（判错率只略高于同分布）， 但分歧完全识别不出“换了说法”的输入（AUC ≈ 0.5）—— 它们共享同一份预训练表示，看法高度一致。 这里恰好没出大问题；可如果共享的表示本身有盲点，集成也看不见。 · 各自预训练的集成：成员更多样，分歧又能识别出“换了说法”的输入了，同时判错率也低。 Eisenstein 等（2023）在真实的奖励模型上发现：预训练种子不同的集成，比只有微调种子不同的集成更能缓解奖励黑客。 · 一个容易踩的坑（最后一列）：直接用分数差的**原始**标准差预测判错，AUC 甚至低于 0.5。 原因：两个回答差距很大时，各成员的分数差都很大，绝对的分歧也大，但这些对几乎不会判错。 分歧要相对于判断的强度来看（相对分歧），或者看“有几个成员意见相反”。""") # ============================================================================= # 演示 2 # ============================================================================= def demo_best_of_n(world, ensembles, n_pools=100): print_header("演示 2：Best-of-n —— 候选越多，越会挑中单个模型\"高估\"的回答") print(f"每个候选池 n 个回答（一半常用说法、一半换了说法），按奖励模型的分数挑最高的一个，报告挑中回答的") print(f"**真实奖励**（{n_pools} 个池的平均，越高越好）。三种挑法：只用第一个成员 / 集成平均分 / 平均分 − 标准差（保守）。\n") print(f" {pad('集成', 20)} {pad('n', 4, 'right')} {pad('单个成员', 9, 'right')} {pad('平均分', 8, 'right')} " f"{pad('保守', 7, 'right')} {pad('最好的', 7, 'right')} {pad('挑中换了说法的比例（单个/平均）', 32, 'right')}") # 先生成好所有候选池，让三种集成在**同一批**回答上比较 pools = {n: [[world.response(COMMON) for _ in range(n // 2)] + [world.response(RARE) for _ in range(n // 2)] for _ in range(n_pools)] for n in [4, 16, 64, 256]} for kind, members in ensembles.items(): for n, pool_list in pools.items(): vals = {"single": [], "mean": [], "lcb": [], "best": []} rare = {"single": [], "mean": []} for pool in pool_list: tr = np.array([true_reward(r) for r in pool]) S = member_rewards(members, pool) for name, sc in [("single", S[0]), ("mean", S.mean(axis=0)), ("lcb", S.mean(axis=0) - S.std(axis=0))]: k = int(np.argmax(sc)) vals[name].append(tr[k]) if name in rare: rare[name].append(k >= n // 2) # 后一半是"换了说法"的回答 vals["best"].append(tr.max()) frac = f"{np.mean(rare['single']):.2f} / {np.mean(rare['mean']):.2f}" print(f" {pad(kind, 20)} {n:>4d} {np.mean(vals['single']):>9.3f} {np.mean(vals['mean']):>8.3f} " f"{np.mean(vals['lcb']):>7.3f} {np.mean(vals['best']):>7.3f} {frac:>32s}") print() print("""观察： · 各自预训练的集成：单个成员对“换了说法”的回答有自己的系统性误差（它的预训练表示和别人不一样）。 n 越大，“挑分数最高的”越会挑中它**高估**的那些回答（挑中换了说法的比例约三分之一），挑中回答的真实奖励 和“最好的”差距越拉越大。集成平均把各成员独立的误差平均掉，几乎不再挑中被高估的回答；保守挑法再好一点点。 这就是**奖励过度优化**的雏形：优化压力越大（n 越大），越会去利用奖励模型的误差（第 8 讲）。 · 共享预训练的集成：三种挑法几乎一样 —— 成员们看法高度一致，平均之后和单个成员差不多。 这里它们本来就不太错，所以结果不差；但如果它们一起错了，平均也平均不掉。 · 从零训练的集成：没见过的词让回答的分数趋于“平庸”（随机词向量平均下来接近 0），很少被挑中， 所以这里看不到过度优化。【边界】这是玩具世界的特点；真实的奖励模型对分布外的输入可能给出极端的分数。 · Coste 等（2023）在 RLHF 里用奖励模型集成的平均或“平均 − 若干倍标准差”当奖励，减轻策略对奖励模型漏洞的利用。""") def main(): world = TextWorld(seed=3) chosen, rejected, _ = world.preference_pairs(400, COMMON) shared = pretrain_embeddings(world, seed=0, epochs=5) ensembles = {kind: build_ensemble(kind, world, chosen, rejected, shared) for kind in ["从零训练", "共享预训练 + 微调", "各自预训练 + 微调"]} demo_disagreement(world, ensembles) demo_best_of_n(world, ensembles) if __name__ == "__main__": main()