""" ============================================================================= 第 6 讲 · 工具模块：一个"同义词"文本世界，用来演示"预训练 + 微调"的奖励模型 ============================================================================= 对应讲义：lectures/lecture06_bayes_finetuning.md 第 6、7 节 ----------------------------------------------------------------------------- 【这个模块解决什么问题？】 ----------------------------------------------------------------------------- 真实的奖励模型是在一个预训练好的大语言模型上加一个"打分头"，再用偏好数据微调出来的。 为什么不从零训练？因为偏好数据很贵、很少（几万到几十万对），而预训练模型已经从海量无标注文本里 学到了"哪些词意思相近"之类的知识。我们用一个极小的玩具世界把这件事演示出来： · 词表：10 个"概念"（准确、清晰、例子……错误），每个概念有 4 个同义词，共 40 个词。 每个概念有一个"人类的评价" v_c（准确 +1.2，错误 −1.5……）。 · 预训练语料：大量无标注的句子。同一个句子里的词来自一两个概念，所以同义词总是出现在相似的上下文里。 语料里**没有**任何关于"好坏"的信息。 · 回答：一串词。真实奖励 r*(回答) = 各个词所属概念的 v_c 的平均值。 · 偏好数据：两个回答，标注员按 BT 模型 P(a ≻ b) = σ(κ·(r*(a) − r*(b))) 选一个。 **训练用的偏好数据只用每个概念的前两个同义词（"常用说法"）**；测试时分两种： 同分布（常用说法）和换了说法（只用后两个同义词，训练时从没见过）。 奖励模型：r_θ(回答) = 打分头( 词向量的平均值 ) · 词向量（嵌入表，40 × 16）可以从零开始，也可以来自预训练； · 打分头是一个线性层（16 → 1）。 ----------------------------------------------------------------------------- 【PyTorch 小知识】（本课第一次用 PyTorch，这里集中说明） ----------------------------------------------------------------------------- torch.tensor(...) ：和 numpy 数组类似的"张量"；可以放在 GPU 上，可以自动求导 nn.Module ：所有神经网络层 / 模型的基类。在 __init__ 里定义参数和子层，在 forward 里写计算 nn.Embedding(V, D) ：一张 V × D 的表，输入词的编号，输出那一行（词向量） nn.Linear(D, 1) ：y = W x + b loss.backward() ：自动计算 loss 对所有参数的梯度（存在每个参数的 .grad 里） optimizer.step() ：按梯度更新参数；optimizer.zero_grad() 把上一步的梯度清零（否则会累加） torch.no_grad() ：在这个 with 块里不记录计算图（评估时用，省内存、更快） p.requires_grad_(False) ：冻结参数 p（不再计算它的梯度，优化器也不会更新它） F.logsigmoid(x) ：数值稳定的 log σ(x) torch.manual_seed(k) ：固定随机种子，让初始化和打乱顺序可复现 ----------------------------------------------------------------------------- 【用法】 ----------------------------------------------------------------------------- from text_world import TextWorld, pretrain_embeddings, RewardModel, train_reward_model, evaluate world = TextWorld(seed=0) emb = pretrain_embeddings(world, seed=0) ... 直接运行本文件会打印词表、几个例子，并检查预训练学到的词向量： python3 text_world.py """ import math import numpy as np import torch import torch.nn as nn import torch.nn.functional as F # 小模型在 CPU 上用单线程反而更快（多线程的调度开销比计算本身还大），也让结果更容易复现。 torch.set_num_threads(1) # ============================================================================= # 第一部分：词表与真实奖励 # ============================================================================= # (概念名, 人类对这个概念的评价 v_c)。评价只用于生成"真实奖励"，学习算法看不到。 CONCEPTS = [("准确", 1.2), ("清晰", 0.8), ("例子", 0.6), ("步骤", 0.5), ("礼貌", 0.1), ("华丽", -0.1), ("啰嗦", -0.6), ("含糊", -0.7), ("跑题", -0.9), ("错误", -1.5)] N_CONCEPTS = len(CONCEPTS) SYNONYMS = 4 # 每个概念 4 个同义词 VOCAB = N_CONCEPTS * SYNONYMS # 40 个词；第 c 个概念的第 k 个同义词编号为 c·4 + k PAD = VOCAB # 第 41 个编号留给"填充"（把不同长度的回答补齐成一样长） COMMON = (0, 1) # 常用说法：每个概念的前两个同义词（训练数据只用它们） RARE = (2, 3) # 换了说法：后两个同义词（只在预训练语料和"换了说法"的测试里出现） VALUES = np.array([v for _, v in CONCEPTS]) def token_name(t): """词的编号 → 名字，例如 0 → "准确¹"，5 → "清晰²"。""" sup = "¹²³⁴" return CONCEPTS[t // SYNONYMS][0] + sup[t % SYNONYMS] def true_reward(tokens): """ 真实奖励 r*(回答) = 各个词所属概念的评价 v_c 的平均值。 【参数】tokens：词编号的列表（不含 PAD） 【返回】float 【边界】用平均值而不是总和，所以"多说几句"本身不加分也不扣分（长度中性）。第 8 讲会故意引入长度偏差。 """ return float(np.mean([VALUES[t // SYNONYMS] for t in tokens])) class TextWorld: """ 生成预训练语料、回答和偏好数据。 【参数】seed：随机种子；kappa：标注员的"敏锐度" κ（越大，标注越接近确定性，噪声越小） """ def __init__(self, seed=0, kappa=4.0): self.rng = np.random.default_rng(seed) self.kappa = kappa # ---------------------------------------------------------------- 预训练语料 def pretraining_corpus(self, n_sentences=20000, length=8): """ 无标注的预训练语料：每个句子先选两个概念（话题），每个位置从这两个概念里随机选一个，再随机选一个同义词。 【返回】形状 (n_sentences, length) 的整数数组 【为什么这样造】真实文本里，意思相近的词出现在相似的上下文里（"分布假设"）。这里同一个概念的 4 个同义词 在语料里的用法完全一样，所以一个"从上下文预测词"的模型会给它们学出相近的向量 —— 即使它从没被告诉过"准确¹ 和 准确³ 是同义词"。 """ rng = self.rng topics = np.stack([rng.choice(N_CONCEPTS, size=2, replace=False) for _ in range(n_sentences)]) which = rng.integers(0, 2, size=(n_sentences, length)) concept = np.take_along_axis(topics, which, axis=1) syn = rng.integers(0, SYNONYMS, size=(n_sentences, length)) return concept * SYNONYMS + syn # ---------------------------------------------------------------- 回答与偏好 def response(self, synonyms): """ 生成一个回答：长度 4~10；先给这个回答抽一个"概念配比"（狄利克雷分布，让不同回答侧重不同的概念）， 再逐个位置抽概念和同义词。 【参数】synonyms：允许使用的同义词下标，例如 COMMON = (0, 1) 【返回】词编号的列表 """ rng = self.rng length = int(rng.integers(4, 11)) mix = rng.dirichlet(np.full(N_CONCEPTS, 0.5)) concepts = rng.choice(N_CONCEPTS, size=length, p=mix) syn = rng.choice(np.array(synonyms), size=length) return [int(c * SYNONYMS + s) for c, s in zip(concepts, syn)] def preference_pairs(self, n, synonyms=COMMON): """ 生成 n 对回答和标注员的选择。 【返回】(chosen, rejected, p_true)：两个回答列表（"被选中的"在前）和"真实的偏好概率" p_true[n] = σ(κ·(r*(chosen) − r*(rejected)))，只用于评估（例如计算准确率的上限） """ chosen, rejected, p_true = [], [], [] for _ in range(n): a, b = self.response(synonyms), self.response(synonyms) p = 1.0 / (1.0 + math.exp(-self.kappa * (true_reward(a) - true_reward(b)))) if self.rng.random() < p: chosen.append(a), rejected.append(b), p_true.append(p) else: chosen.append(b), rejected.append(a), p_true.append(1 - p) return chosen, rejected, np.array(p_true) def pad_batch(seqs): """ 把长度不同的回答补齐成一个矩阵。 【返回】(tokens, mask)：tokens 形状 (B, 最大长度)，不够长的位置填 PAD；mask 同形状，真实的词为 1、填充为 0 【Python 小知识】torch.full((B, L), PAD) 生成一个全是 PAD 的矩阵，再把每个回答写进它的那一行。 """ L = max(len(s) for s in seqs) tokens = torch.full((len(seqs), L), PAD, dtype=torch.long) mask = torch.zeros((len(seqs), L)) for b, s in enumerate(seqs): tokens[b, :len(s)] = torch.tensor(s, dtype=torch.long) mask[b, :len(s)] = 1.0 return tokens, mask # ============================================================================= # 第二部分：预训练（CBOW：从上下文的平均词向量预测中间的词） # ============================================================================= class CBOW(nn.Module): """ 一个极简的"语言模型"：用句子里其他词的词向量的平均值，预测被遮住的那个词。 【思路】这就是 word2vec 的 CBOW（Mikolov 等 2013）。它和 BERT 的"遮住一个词再预测"是同一个思想， 只是没有 Transformer：上下文的表示就是词向量的平均。训练完以后我们只要它的嵌入表 self.emb。 """ def __init__(self, dim=16): super().__init__() self.emb = nn.Embedding(VOCAB + 1, dim, padding_idx=PAD) # 输入词向量（我们要的东西） self.out = nn.Linear(dim, VOCAB) # 输出层：词向量 → 40 个词的打分 def forward(self, context, mask): """context：(B, L) 的上下文词编号；mask：(B, L)。返回 (B, 40) 的打分（logits）。""" h = (self.emb(context) * mask.unsqueeze(-1)).sum(1) / mask.sum(1, keepdim=True) return self.out(h) def pretrain_embeddings(world, dim=16, n_sentences=20000, epochs=3, seed=0, verbose=False): """ 在无标注语料上预训练词向量。 【思路】对语料里的每个句子、每个位置：把这个位置遮住，用其余 7 个词预测它（交叉熵损失）。 同义词出现在相同的上下文里、也被相同的上下文预测，它们的词向量会被推到相近的地方。 【参数】world：TextWorld；dim：词向量维数；epochs：把语料过几遍；seed：初始化的随机种子 【返回】形状 (41, dim) 的张量：预训练好的嵌入表（最后一行是 PAD，为 0） 【边界】预训练完全不知道"哪个概念好、哪个坏"（语料里没有这种信息）。它学到的只是"哪些词是一类"。 奖励模型要做的，是在这个表示之上学出"哪一类好"。 """ torch.manual_seed(seed) corpus = torch.tensor(world.pretraining_corpus(n_sentences), dtype=torch.long) n, L = corpus.shape model = CBOW(dim) opt = torch.optim.Adam(model.parameters(), lr=0.01) eye = torch.eye(L) for ep in range(epochs): perm = torch.randperm(n) total = 0.0 for start in range(0, n, 256): batch = corpus[perm[start:start + 256]] # 每个句子随机遮住一个位置：mask 在那个位置为 0 pos = torch.randint(0, L, (len(batch),)) mask = 1.0 - eye[pos] target = batch[torch.arange(len(batch)), pos] loss = F.cross_entropy(model(batch, mask), target) opt.zero_grad() loss.backward() opt.step() total += loss.item() * len(batch) if verbose: print(f" 预训练第 {ep + 1} 遍：平均交叉熵 {total / n:.3f}（随机猜是 log 40 = {math.log(40):.3f}；" f"只知道话题时的下限约 log 8 = {math.log(8):.3f}）") return model.emb.weight.detach().clone() def synonym_similarity(emb): """ 检查词向量：同一个概念的同义词之间、不同概念的词之间，余弦相似度平均是多少。 【返回】(同义词之间的平均相似度, 不同概念之间的平均相似度) """ E = F.normalize(emb[:VOCAB], dim=1) S = (E @ E.T).numpy() concept = np.arange(VOCAB) // SYNONYMS same = (concept[:, None] == concept[None, :]) & ~np.eye(VOCAB, dtype=bool) diff = concept[:, None] != concept[None, :] return float(S[same].mean()), float(S[diff].mean()) # ============================================================================= # 第三部分：奖励模型 # ============================================================================= class RewardModel(nn.Module): """ r_θ(回答) = 打分头( 回答里所有词的词向量的平均 )。 【参数】 dim ：词向量维数 init_emb ：None → 随机初始化（"从零训练"）；否则用这个张量初始化嵌入表（"预训练"） freeze_emb：True → 冻结嵌入表，只训练打分头（"线性探针"） 【偏好学习联系】真实的奖励模型（InstructGPT、Llama 2）：取一个预训练 / SFT 过的语言模型，去掉最后的 "预测下一个词"的输出层，换成一个输出一个数的线性层，用回答最后一个位置的隐藏状态打分。 这里的"词向量平均"相当于一个极简的编码器。 """ def __init__(self, dim=16, init_emb=None, freeze_emb=False): super().__init__() self.emb = nn.Embedding(VOCAB + 1, dim, padding_idx=PAD) if init_emb is not None: with torch.no_grad(): self.emb.weight.copy_(init_emb) if freeze_emb: self.emb.weight.requires_grad_(False) self.head = nn.Linear(dim, 1) def forward(self, tokens, mask): """tokens、mask：(B, L)。返回长度 B 的奖励。""" h = (self.emb(tokens) * mask.unsqueeze(-1)).sum(1) / mask.sum(1, keepdim=True) return self.head(h).squeeze(-1) def train_reward_model(model, chosen, rejected, epochs=10, lr=0.01, batch_size=32, seed=0, bootstrap=False, val=None, history=None): """ 用 BT 损失训练奖励模型：loss = −log σ(r(被选中的) − r(被拒绝的))（第 1、5 讲）。 【参数】 model ：RewardModel chosen, rejected ：两个回答列表 epochs ：把数据过几遍 lr ：Adam 的学习率（只更新 requires_grad 为 True 的参数） batch_size ：每一步用多少对 seed ：打乱顺序（以及自助法抽样）的随机种子 bootstrap ：True → 先有放回地重新抽一遍训练数据（集成时让成员看到不同的数据） val ：可选，(chosen, rejected) 验证集；配合 history 记录每遍之后的验证损失 history ：可选，空列表；每遍结束后追加 (训练损失, 验证损失) 【返回】model（就地训练好的） 【边界】InstructGPT 和 Llama 2 的奖励模型都只训练 1 个 epoch：多训练几遍很容易过拟合（代码 05 演示 3）。 """ g = torch.Generator().manual_seed(seed) n = len(chosen) idx = torch.randint(0, n, (n,), generator=g) if bootstrap else torch.arange(n) ct, cm = pad_batch([chosen[k] for k in idx.tolist()]) rt, rm = pad_batch([rejected[k] for k in idx.tolist()]) params = [p for p in model.parameters() if p.requires_grad] opt = torch.optim.Adam(params, lr=lr) for ep in range(epochs): perm = torch.randperm(n, generator=g) total = 0.0 for start in range(0, n, batch_size): b = perm[start:start + batch_size] margin = model(ct[b], cm[b]) - model(rt[b], rm[b]) loss = -F.logsigmoid(margin).mean() opt.zero_grad() loss.backward() opt.step() total += loss.item() * len(b) if history is not None: vloss = bt_loss(model, *val) if val is not None else float("nan") history.append((total / n, vloss)) return model @torch.no_grad() def reward_margins(model, chosen, rejected): """返回每一对的 r(被选中的) − r(被拒绝的)（numpy 数组）。""" ct, cm = pad_batch(chosen) rt, rm = pad_batch(rejected) return (model(ct, cm) - model(rt, rm)).numpy() def bt_loss(model, chosen, rejected): """平均 BT 损失（对数损失）。""" m = torch.tensor(reward_margins(model, chosen, rejected)) return float(-F.logsigmoid(m).mean()) def evaluate(model, chosen, rejected, p_true): """ 评估奖励模型。 【返回】dict： agree ：和**真实偏好方向**一致的比例（r_θ 的差值和 r* 的差值同号）—— 去掉了标注噪声，最能说明"学得对不对" acc ：和**标注**一致的比例（真实世界里只能算这个） loss ：对数损失 【思路】p_true > 0.5 表示"被选中的那个真实奖励更高"。 """ m = reward_margins(model, chosen, rejected) truly_better = p_true > 0.5 return dict(agree=float(np.mean((m > 0) == truly_better)), acc=float(np.mean(m > 0)), loss=float(np.mean(np.logaddexp(0.0, -m)))) if __name__ == "__main__": world = TextWorld(seed=0) print("词表（每行一个概念：评价 v_c 和 4 个同义词；¹² 是常用说法，³⁴ 是换了说法）：") for c, (name, v) in enumerate(CONCEPTS): print(f" {name} v = {v:+.1f} " + " ".join(token_name(c * SYNONYMS + k) for k in range(SYNONYMS))) ch, rj, p = world.preference_pairs(3) print("\n三对偏好数据（被选中的在前）：") for a, b, q in zip(ch, rj, p): print(f" {' '.join(token_name(t) for t in a)} (r* = {true_reward(a):+.2f})") print(f" ≻ {' '.join(token_name(t) for t in b)} (r* = {true_reward(b):+.2f}) 真实偏好概率 {q:.2f}") print("\n预训练词向量：") emb = pretrain_embeddings(world, seed=0, verbose=True) same, diff = synonym_similarity(emb) print(f" 同义词之间的平均余弦相似度 {same:.3f}，不同概念的词之间 {diff:.3f}") rand_same, rand_diff = synonym_similarity(nn.Embedding(VOCAB + 1, 16).weight.detach()) print(f" （随机初始化的词向量：同义词之间 {rand_same:.3f}，不同概念之间 {rand_diff:.3f}）")