""" objectives.py —— 第 2 讲 §6：四种预训练目标怎样把"一篇文档"变成训练样本 ===================================================================== 【这个文件要回答的问题】 预训练不需要人工标注：答案就藏在文本里。不同的"自监督目标"只是用不同的方式 把一段文本拆成"输入"和"要预测的答案"： FLM 完整语言建模（full LM / CLM，GPT）：每个位置预测下一个 token，100% 的 token 都有损失。 PLM 前缀语言建模（prefix LM）：随机选一个切分点，前缀双向可见、只当条件； 只对前缀之后的 token 计算损失（平均约 50%）。 MLM 掩码语言建模（BERT）：随机选 15% 的位置，其中 80% 换成 、10% 换成随机词、 10% 保持不变；只对这 15% 计算损失。 SC span corruption（T5）：随机选约 15% 的 token，组成平均长 3 的若干段（span）， 每段在输入里换成一个哨兵 ；目标序列是 " 第 0 段 第 1 段 ... "。 每个函数只处理一篇文档，返回一个字典；*_batches 函数把一批文档组装成张量；*_loss 函数计算损失。 【怎样使用】 from objectives import mlm_example, mlm_batches, mlm_loss ex = mlm_example(ids, rng, vocab) # {"input": [...], "targets": {位置: 原 id}} 04 号脚本先把四种目标的样本打印出来，再分别训练四个模型做比较。 【LLM 联系】 Wang et al. (2022) 在 5B 参数、168B token 的规模上比较的正是 FLM、PLM 和 MLM（span corruption 版本）。 他们"按看过的 token 数对齐预算"：FLM 每看一个 token 就得到一个训练信号，PLM 只有一半，MLM 只有约 15%～18%。 """ import random import torch import torch.nn.functional as F from garden_lm import BOS, EOS, MASK, PAD, SENTINELS, pad_batch # ===================================================================== # 1. 单篇文档 → 一个样本 # ===================================================================== def flm_example(ids, bos_id): """ 完整语言建模：输入 = [] + ids[:-1]，目标 = ids（错开一位），每个位置都有损失。 【参数】ids：一篇文档的 id 列表（以 结尾）；bos_id： 的 id。 【返回】{"input": list, "target": list, "loss_mask": list[bool]}，三者等长。 """ return {"input": [bos_id] + ids[:-1], "target": list(ids), "loss_mask": [True] * len(ids)} def plm_example(ids, bos_id, rng): """ 前缀语言建模：在 [0, n-1] 里随机选切分点 s，前 s 个 token（连同 ）组成双向可见的前缀。 【返回】{"input", "target", "loss_mask", "prefix_len"}。 prefix_len = s + 1（ 也算在前缀里）；loss_mask 只在位置 s..n-1 为 True。 【为什么前缀里不能算损失】前缀内部是双向注意力：位置 t 能看见位置 t+1 的输入， 而位置 t 的目标恰好就是 t+1 的输入——"答案就在眼前"，算损失没有意义。 所以前缀只当条件，损失只算在前缀之后（那里仍是因果的）。 【边界】s = 0 时前缀只有 ，退化成完整语言建模；s = n-1 时只预测最后一个 token（）。 """ n = len(ids) s = rng.randrange(0, n) return {"input": [bos_id] + ids[:-1], "target": list(ids), "loss_mask": [t >= s for t in range(n)], "prefix_len": s + 1} def mlm_example(ids, rng, vocab, p=0.15, n_real=90): """ BERT 式掩码语言建模。序列 = [] + ids； 永远不遮。 【思路】随机选约 p 的位置作为"要预测的位置"。对每个选中的位置： 80% 换成 ，10% 换成一个随机的真实词，10% 保持原样。 【为什么不全换成 】预训练时输入里有 ，微调和使用时却没有；如果模型只在 处 才认真预测，它在没有 的输入上就学不到好的表示。混入随机词和原词， 逼模型对**每个**位置都保持警惕（Devlin et al. 2018 的解释）。 【参数】vocab：Vocab（要用 的 id）；p：遮盖比例；n_real：随机替换只从前 n_real 个 id （花园世界的真实 token）里抽，不会抽到 //哨兵。 【返回】{"input": list, "targets": {位置: 原 id}}。位置是在 [] + ids 里的下标。 【边界】至少选 1 个位置（很短的文档按 15% 算可能是 0 个）。 """ seq = [vocab.stoi[BOS]] + list(ids) cand = list(range(1, len(seq))) k = max(1, round(p * len(cand))) chosen = sorted(rng.sample(cand, k)) inp = list(seq) for pos in chosen: r = rng.random() if r < 0.8: inp[pos] = vocab.stoi[MASK] elif r < 0.9: inp[pos] = rng.randrange(3, n_real) # 跳过 return {"input": inp, "targets": {pos: seq[pos] for pos in chosen}} def _random_segmentation(n_items, n_segments, rng): """把 n_items 个元素随机切成 n_segments 段，每段至少 1 个。【返回】各段长度的列表。""" cuts = sorted(rng.sample(range(1, n_items), n_segments - 1)) if n_segments > 1 else [] bounds = [0] + cuts + [n_items] return [bounds[i + 1] - bounds[i] for i in range(n_segments)] def span_corruption_example(ids, rng, vocab, noise=0.15, mean_span=3.0): """ T5 的 span corruption（Raffel et al. 2020）。 【思路】 1. 要遮 n_noise = round(noise × n) 个 token，分成 n_spans = round(n_noise / mean_span) 段； 2. 把 n_noise 个"噪声 token"随机切成 n_spans 段（每段至少 1 个），把保留的 token 随机切成 n_spans + 1 段（首尾两段可以为空，中间每段至少 1 个），交替排列： 保留段 0、噪声段 0、保留段 1、噪声段 1、……、噪声段 k-1、保留段 k； 3. 输入：保留段原样保留，每个噪声段换成一个哨兵 ； 目标： 噪声段 0 噪声段 1 ... 。 【为什么用哨兵和"只输出被遮的部分"】目标序列很短（约 n × 15% + 段数 + 1），decoder 的计算量小； 哨兵告诉模型"这一段该填到哪里"。和 BERT 相比，一段可以有好几个 token， 模型得学会"填多长"，更接近生成。 【和 T5 官方实现的一处不同】T5 的实现总是以一段噪声**结尾**（保留段和噪声段各切 k 段、严格交替）。 T5 的样本是 512 个 token 的长块，一块里有约 26 段，"最后一段在末尾"无关紧要； 而花园世界的文档很短，一篇只有 1～2 段——照搬的话，被遮的那段**永远在文档末尾**， 模型会学到"空总在最后"，在句子中间挖空时完全不会填（本讲 README 记录了这次翻车）。 所以这里让第一段和最后一段保留段可以为空，噪声段可以出现在任何位置。 【参数】ids：文档 id 列表（以 结尾， 不参与遮盖）；vocab：Vocab；noise、mean_span 同 T5 默认值。 【返回】{"input": list, "target": list, "spans": [(起点, 长度), ...]}（起点是在去掉 的文档里的下标）。 【边界】 - 只有 8 个哨兵：长文档的段数被截到 8，平均段长相应变长（花园世界里只有很少的长文档受影响）。 - 至少遮 1 个 token、分 1 段；段数不超过"保留 token 数 + 1"（保证中间的保留段非空）。 """ body = list(ids[:-1]) # 不遮 n = len(body) n_noise = min(max(1, round(noise * n)), max(1, n - 1)) n_keep = n - n_noise n_spans = max(1, min(len(SENTINELS), round(n_noise / mean_span), n_keep + 1)) noise_lens = _random_segmentation(n_noise, n_spans, rng) # 保留段：把 n_keep + 2 切成 n_spans + 1 个正整数，再把首尾两段各减 1（于是首尾可以为 0） keep_lens = _random_segmentation(n_keep + 2, n_spans + 1, rng) keep_lens[0] -= 1 keep_lens[-1] -= 1 inp, tgt, spans, pos = [], [], [], 0 for i in range(n_spans): inp += body[pos:pos + keep_lens[i]] pos += keep_lens[i] sent = vocab.stoi[SENTINELS[i]] inp.append(sent) tgt.append(sent) tgt += body[pos:pos + noise_lens[i]] spans.append((pos, noise_lens[i])) pos += noise_lens[i] inp += body[pos:] + [vocab.stoi[EOS]] tgt.append(vocab.stoi[EOS]) return {"input": inp, "target": tgt, "spans": spans} # ===================================================================== # 2. 一批文档 → 张量；损失函数 # ===================================================================== def plm_batches(docs, vocab, batch_size, seed): """ 前缀 LM 的批次（每次调用重新随机选切分点）。 【返回】list[(ids, pad, loss_mask, prefix_len)]：ids (B, L)，pad (B, L)，loss_mask (B, L)，prefix_len (B,)。 """ rng = random.Random(seed) bos = vocab.stoi[BOS] exs = sorted((plm_example(vocab.encode(d), bos, rng) for d in docs), key=lambda e: len(e["input"])) out = [] for i in range(0, len(exs), batch_size): chunk = exs[i:i + batch_size] ids, pad = pad_batch([e["input"] for e in chunk], vocab.stoi[PAD]) tgt, _ = pad_batch([e["target"] for e in chunk], vocab.stoi[PAD]) lm = torch.zeros_like(pad) for r, e in enumerate(chunk): lm[r, :len(e["loss_mask"])] = torch.tensor(e["loss_mask"]) out.append((ids, pad, tgt, lm, torch.tensor([e["prefix_len"] for e in chunk]))) return out def plm_loss(model, batch): ids, pad, tgt, lm, prefix_len = batch logits = model(ids, pad=pad, prefix_len=prefix_len) nll = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), tgt.reshape(-1), reduction="none") keep = lm.reshape(-1) return (nll * keep).sum(), keep.sum() def mlm_batches(docs, vocab, batch_size, seed): """ MLM 的批次（每次调用重新随机遮盖 = 动态掩码）。 【返回】list[(ids, pad, tgt, loss_mask)]：tgt 在被选中的位置是原 id，其余位置随便填（不计损失）。 """ rng = random.Random(seed) items = [] for d in sorted(docs, key=len): ex = mlm_example(vocab.encode(d), rng, vocab) tgt = list(ex["input"]) mask = [False] * len(tgt) for pos, orig in ex["targets"].items(): tgt[pos] = orig mask[pos] = True items.append((ex["input"], tgt, mask)) out = [] for i in range(0, len(items), batch_size): chunk = items[i:i + batch_size] ids, pad = pad_batch([c[0] for c in chunk], vocab.stoi[PAD]) tgt, _ = pad_batch([c[1] for c in chunk], vocab.stoi[PAD]) lm = torch.zeros_like(pad) for r, c in enumerate(chunk): lm[r, :len(c[2])] = torch.tensor(c[2]) out.append((ids, pad, tgt, lm)) return out def mlm_loss(model, batch): ids, pad, tgt, lm = batch logits = model(ids, pad=pad) nll = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), tgt.reshape(-1), reduction="none") keep = lm.reshape(-1) return (nll * keep).sum(), keep.sum() def sc_batches(docs, vocab, batch_size, seed): """ span corruption 的批次（每次调用重新随机选段）。 【返回】list[(src, src_pad, tgt, tgt_pad)]：decoder 的输入是 tgt[:, :-1]，目标是 tgt[:, 1:]。 目标序列前面补一个 作为 decoder 的起始符。 """ rng = random.Random(seed) bos = vocab.stoi[BOS] exs = [span_corruption_example(vocab.encode(d), rng, vocab) for d in sorted(docs, key=len)] out = [] for i in range(0, len(exs), batch_size): chunk = exs[i:i + batch_size] src, src_pad = pad_batch([e["input"] for e in chunk], vocab.stoi[PAD]) tgt, tgt_pad = pad_batch([[bos] + e["target"] for e in chunk], vocab.stoi[PAD]) out.append((src, src_pad, tgt, tgt_pad)) return out def sc_loss(model, batch): src, src_pad, tgt, tgt_pad = batch logits = model(src, tgt[:, :-1], src_pad=src_pad, tgt_pad=tgt_pad[:, :-1]) nll = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), tgt[:, 1:].reshape(-1), reduction="none") keep = (~tgt_pad[:, 1:]).reshape(-1) return (nll * keep).sum(), keep.sum()