""" ============================================================================= 第 5 讲 · 代码 2：上下文学习 = 隐式的贝叶斯推断？——一个能精确算出"理想答案"的小世界 ============================================================================= 对应讲义：lectures/lecture05_emergence_icl.md 第 3、4 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- 第 1 讲的 base model 不会 few-shot：给 2 个没见过格式的示例，准确率仍是 0.00。那么 in-context learning（ICL） 到底是怎么来的？Xie et al. 2021 的回答是：预训练文档各自有一个隐藏的"概念"，要预测好下一个 token， 模型必须从前文推断概念——提示里的示例只是推断的证据。这个文件在 icl_world.py 的小世界里检验这个说法： 演示 1：两个预训练世界。世界 A：8 个固定的"任务"反复出现；世界 B：每篇文档一个新的随机任务。 在两个世界上各预训练一个小 Transformer（2 层、64 维），只用"预测下一个字符"。 演示 2：ICL 曲线——示例越多越准吗？和**贝叶斯最优**的答案（完美学会预训练分布的模型）比一比； 再和"只看词频、不看配对"的贝叶斯比一比：模型有没有用上"哪个输入配哪个标签"？ 演示 3：Min et al. 2022 的消融：标签随机、输入分布错、标签空间错、格式错——哪个伤害最大？ 贝叶斯答案能不能预测这些现象？ 演示 4：Wei et al. 2023 的问题：示例里的标签被**翻转**、或者是一个**没见过的映射**时，模型跟着示例走， 还是跟着预训练学到的先验走？两个世界的"理想答案"截然不同。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 02_icl_as_bayes.py # 第一次约 1 分钟（预训练 2 个小模型并缓存）；之后约 15 秒 需要 PyTorch。依赖本目录的 icl_world.py、seq_tasks.py，以及 code/lecture01 的 tiny_gpt.py、console.py。 【Python 小知识】 - lambda g: doc_batch(TOK, world, 32, g)：一个"没有名字的小函数"，等价于 def f(g): return doc_batch(TOK, world, 32, g) 训练循环只认识"给我一个随机数生成器、我还你一批数据"的函数，lambda 把其余参数提前填好。 - {x for x, _ in it.demos}：集合推导式，得到示例里出现过的字母的集合；"in 集合"的判断比"in 列表"快。 """ import os import sys HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) import random # noqa: E402 from console import pad, print_header # noqa: E402 from icl_world import (BLOCK_SIZE, CONCEPTS, CONDITIONS, INPUTS, LABEL_SETS, bayes_accuracy, doc_batch, # noqa: E402 input_dist, make_doc, make_items, make_tokenizer, model_accuracy) from seq_tasks import GPTConfig, config_dict, load_or_train, new_model, train_masked_lm # noqa: E402 TOK = make_tokenizer() SHAPE = dict(n_layer=2, n_head=4, d_model=64) STEPS = 2500 N_ITEMS = 500 # 每个 (条件, k) 的评测提示个数 KS = [0, 1, 2, 4, 8, 12] def get_model(world): """ 在世界 world 上预训练（或读取缓存的）小 Transformer。 【思路】和第 1 讲的预训练一模一样：只有"预测下一个字符"，没有任何任务标签。每 500 步记录一次 "8 个正确示例时的 ICL 准确率"，看 ICL 是在训练的什么时候出现的。 【返回】(model, info)；info["curve"] = list[(step, 准确率)] """ cfg = GPTConfig(vocab_size=len(TOK.alphabet), block_size=BLOCK_SIZE, **SHAPE) meta = dict(config=config_dict(cfg), task="icl_world", world=world, steps=STEPS, batch_size=32, lr=3e-3, model_seed=0, data_seed=1, concepts=[[t, [m[x] for x in INPUTS]] for t, m in CONCEPTS]) def build_and_train(): model = new_model(cfg, seed=0) probe = make_items("gold", 8, 300, seed=7) curve = [] train_masked_lm(model, lambda g: doc_batch(TOK, world, 32, g), steps=STEPS, seed=1, callback=lambda step, m: curve.append((step, model_accuracy(m, TOK, probe))), callback_every=500) return model, dict(curve=curve) model, info, _ = load_or_train(f"icl_world{world}", meta, build_and_train) return model, info # ============================================================================= # 演示 1：两个预训练世界 # ============================================================================= def demo_worlds(models): print_header("演示 1：两个预训练世界——同样的格式，任务的\"多样性\"不同") rng = random.Random(3) for world, desc in (("A", "8 个固定任务（4 个主题 × 2 种映射）反复出现"), ("B", "每篇文档一个新的随机映射（任务无穷多）")): print(f"世界 {world}：{desc}。两篇预训练文档：") for _ in range(2): print(" " + make_doc(rng, world)) print("\n主题 t 的常用字母 / 标签空间：" + "；".join( f"{t}: {INPUTS[2 * t]}{INPUTS[2 * t + 1]} / {{{','.join(LABEL_SETS[t])}}}" for t in range(len(LABEL_SETS)))) print(f"（常用字母各占 {input_dist(0)['a']:.1f}，其余 6 个字母各约 {input_dist(0)['c']:.3f}；标签有 10% 的概率是 0–9 的随机数字）") print("\n世界 A 的 8 个概念（输入 a–h 各自对应的标签）：") for t, m in CONCEPTS: print(f" 主题 {t}：" + "".join(m[x] for x in INPUTS)) print("\n训练过程中，8 个正确示例时的 ICL 准确率（每 500 步，300 个提示）：") for world, (_, info) in models.items(): print(f" 世界 {world}：" + " ".join(f"{s}:{a:.2f}" for s, a in info["curve"])) print(""" 观察 Observations： · 两个世界的文档格式完全一样，差别只在"任务从哪里来"：世界 A 的映射只有 8 种，模型可以把它们 **记在权重里**，示例只需要帮它认出是哪一种；世界 B 每篇文档的映射都是新的，想预测对，只能**从前文照抄**。 · 训练曲线：世界 A 的模型在第 500 步就有了很强的 ICL（0.91），之后升到 0.99–1.00； 世界 B 的模型在第 500 步是 0.72，之后一直停在 0.68–0.70（演示 2、4 会说明这 0.7 是怎么来的）。""") # ============================================================================= # 演示 2：ICL 曲线 vs 贝叶斯最优 # ============================================================================= def demo_icl_curve(models): print_header("演示 2：示例越多越准吗？—— 模型 vs 贝叶斯最优（完整证据）vs 贝叶斯（只看词频、不看配对）") print(f"每个 k：{N_ITEMS} 个提示，示例标签正确（gold）。准确率 = 模型下一个字符的 argmax 是否等于正确标签。\n") head = pad("k", 5) for world in ("A", "B"): head += pad(f"{world}:模型", 9, ">") + pad(f"{world}:贝叶斯", 10, ">") + pad(f"{world}:只看词频", 12, ">") + " " print(head) for k in KS: items = make_items("gold", k, N_ITEMS, seed=1000 + k) row = pad(k, 5) for world in ("A", "B"): row += (pad(f"{model_accuracy(models[world][0], TOK, items):.3f}", 9, ">") + pad(f"{bayes_accuracy(world, items, 'full'):.3f}", 10, ">") + pad(f"{bayes_accuracy(world, items, 'bag'):.3f}", 12, ">") + " ") print(row) print(""" （"贝叶斯"：一个完美学会了该世界预训练分布的模型；"只看词频"：同样完美，但只知道示例里出现了哪些输入、 哪些标签，不知道哪个输入配哪个标签。两者都由 icl_world.bayes_predict 精确算出，不是训练出来的。） 观察 Observations： · 世界 A：模型的 ICL 曲线几乎和贝叶斯最优重合——k = 0 时 0.520 vs 0.516（只靠查询字母猜主题）， 1 个示例 0.786 vs 0.784，2 个 0.920 vs 0.914，4 个 0.984 vs 0.982。一个只学过"预测下一个字符"的 10 万参数模型，在这个世界里的 ICL 和"对隐藏概念做贝叶斯推断"在数值上几乎一样—— 这就是 Xie et al. 说的"隐式贝叶斯推断"。 · 模型明显好于"只看词频"的贝叶斯（k = 2：0.920 vs 0.840；k = 4：0.984 vs 0.868）： 它确实用上了"哪个输入配哪个标签"，用来在同一主题的两种映射之间做选择。 · 世界 B：贝叶斯最优随 k 稳步上升（k = 8 时 0.871，k = 12 时 0.920）——它的办法是 "查询字母在示例里出现过，就照抄它的标签"。模型却停在 0.7 左右，和"只看词频"的贝叶斯几乎一样 （k = 8：0.678 vs 0.686）：它学会了从示例里认出主题和标签空间，但没有学会照抄。演示 4 再看这件事。""") # ============================================================================= # 演示 3：Min et al. 的消融 # ============================================================================= def demo_min_ablations(models, k=8): print_header(f"演示 3：示例的哪一部分重要？（Min et al. 2022 的消融，k = {k} 个示例）") rows = [("zero-shot（没有示例）", "gold", 0), (CONDITIONS["gold"], "gold", k), (CONDITIONS["random"], "random", k), (CONDITIONS["random_oodin"], "random_oodin", k), (CONDITIONS["random_space"], "random_space", k), (CONDITIONS["inputs_only"], "inputs_only", k)] print(pad("示例", 46) + pad("A:模型", 9, ">") + pad("A:贝叶斯", 10, ">") + pad("B:模型", 9, ">") + pad("B:贝叶斯", 10, ">")) for name, cond, kk in rows: items = make_items(cond, kk, N_ITEMS, seed=2000 + kk) line = pad(name, 46) for world in ("A", "B"): acc = model_accuracy(models[world][0], TOK, items) bayes = "—" if cond == "inputs_only" else f"{bayes_accuracy(world, items, 'full'):.3f}" line += pad(f"{acc:.3f}", 9, ">") + pad(bayes, 10, ">") print(line) print(""" （"—"：这种提示在预训练分布下概率为 0——预训练文档里字母后面一定跟着数字——贝叶斯答案没有定义， 模型在这里的行为是"外推"，不是推断。） 观察 Observations（先看世界 A，它是"模型 ≈ 贝叶斯"的那个世界）： · 标签随机：从 1.000 掉到 0.722（贝叶斯 0.705），但仍明显高于 zero-shot（0.504）。随机标签仍然告诉模型 "标签空间是哪两个数字"，输入仍然带着"主题的常用字母"，被破坏的只是"哪个输入配哪个标签"这一条证据。 Min et al. 在真实模型上看到的"标签对不对没那么重要"，在贝叶斯视角下就是：**示例里有好几条证据， 标签的对错只是其中一条**。 · 但在这个世界里，标签对错比在真实 LLM 里重要得多：随机标签只保住了相对 zero-shot 提升的约 44% （(0.722 − 0.504) / (1.000 − 0.504)），而 Min et al. 报告的平均下降只有 0–5 个百分点。贝叶斯答案给出了原因： 这里每个主题有两种映射，只有配对证据能区分它们；真实任务里，"great → positive"这类对应关系 模型在预训练时早已学会（先验很强），示例几乎只需要提供格式、标签空间和输入分布。 · 再把输入换成"不按主题分布"的字母：贝叶斯答案从 0.705 降到 0.617（输入分布这条证据也没了）， 模型却只从 0.722 降到 0.708——在这一项上模型**比理想的学习者更不在乎输入分布**。这是本表里模型和贝叶斯 差得最多的一处，如实记下。 · 标签换成标签空间之外的数字 4–9：模型 0.606、贝叶斯 0.691，都低于"随机但在标签空间里"的 0.72 左右—— 标签空间这条证据很重要。 · 只有输入、没有标签（格式错）：0.522，回到 zero-shot（0.504）的水平——格式一坏，示例几乎完全没用。 第 1 讲 §4.2 里"没见过的问法，2-shot 也是 0.00"是同一类现象。 · 世界 B：正确标签 0.704（贝叶斯 0.878），随机标签 0.510（贝叶斯 0.506）。在一个"映射每篇都不同"的世界里， 标签是映射的**唯一**证据，随机标签让贝叶斯答案也只剩"认出标签空间"的那一点（zero-shot 0.372 → 0.506）。 **标签对错重不重要，取决于预训练分布有没有把映射"预先定好"**——这正是 Min et al. 和 Wei et al. 2023 两组看似矛盾的发现之间的桥梁（讲义 §3.3）。""") # ============================================================================= # 演示 4：翻转的标签、没见过的映射 # ============================================================================= def demo_flip_and_new(models): print_header("演示 4：标签翻转、没见过的映射——跟着示例走，还是跟着先验走？") ks = [0, 2, 4, 8, 12] for cond in ("flip", "new"): print(f"\n{CONDITIONS[cond]}（准确率 = 是否给出**示例所暗示的**标签）") print(pad("k", 5) + pad("A:模型", 9, ">") + pad("A:贝叶斯", 10, ">") + pad("B:模型", 9, ">") + pad("B:贝叶斯", 10, ">")) for k in ks: items = make_items(cond, k, N_ITEMS, seed=3000 + k) line = pad(k, 5) for world in ("A", "B"): line += (pad(f"{model_accuracy(models[world][0], TOK, items):.3f}", 9, ">") + pad(f"{bayes_accuracy(world, items, 'full'):.3f}", 10, ">")) print(line) # 世界 B：把 k = 8 的正确示例提示按"查询字母在示例里出现过没有"分成两组 items = make_items("gold", 8, N_ITEMS, seed=1008) seen = [it for it in items if it.xq in {x for x, _ in it.demos}] unseen = [it for it in items if it.xq not in {x for x, _ in it.demos}] print(f"\n世界 B，k = 8 个正确示例，按\"查询字母在示例里出现过没有\"分组（{len(seen)} / {len(unseen)} 个提示）：") print(pad("", 22) + pad("模型", 8, ">") + pad("贝叶斯", 8, ">") + pad("只看词频", 10, ">")) for name, group in (("查询字母出现过", seen), ("查询字母没出现过", unseen)): print(pad(name, 22) + pad(f"{model_accuracy(models['B'][0], TOK, group):.3f}", 8, ">") + pad(f"{bayes_accuracy('B', group, 'full'):.3f}", 8, ">") + pad(f"{bayes_accuracy('B', group, 'bag'):.3f}", 10, ">")) print(""" 观察 Observations： · 世界 A 的贝叶斯答案**跟不上**翻转的标签和没见过的映射：k = 8 时只有 0.544 和 0.650，k = 12 也没有更好。 它的先验里只有 8 个概念，示例再多，也只能在这 8 个里挑一个"最接近"的。模型同样如此，而且比贝叶斯更固执 （翻转 k = 8：0.444 vs 0.544）。Xie et al. 在他们的玩具世界里也发现：预训练时没见过的概念，ICL 学不会。 · 世界 B 的贝叶斯答案能跟上任何映射（翻转 k = 12：0.899；没见过的映射 k = 12：0.913）：它的先验是 "每篇文档的映射都是新的"，最优策略就是照抄示例——这正是一个"学习算法"。**同样是贝叶斯推断，先验不同， 表现出来的行为完全不同**：任务少时，ICL 更像"认出一个见过的任务"；任务足够多样时，最优的 ICL 才是 "从示例里学一个新任务"（Raventós et al. 2023 在线性回归上看到了这个转变）。 · 我们的世界 B 模型没有学会照抄。最后一张表是证据：查询字母在示例里出现过时，贝叶斯答案是 1.000， 模型只有 0.747——和"只看词频"的贝叶斯（0.747）一模一样；没出现过时三者都在 0.5 左右。 模型的全部 ICL 都可以用"认出主题和标签空间"来解释，没有用到"这个字母上次配的是哪个标签"。 · 照抄需要一个"找到前面同一个字母、读出它后面的标签"的电路——induction head（Olsson et al. 2022；第 12 讲）。 它通常在训练中经过一段平台期才突然形成；在 2500 步、约 380 万个 token 的预算里，这个 2 层、使用可学习 绝对位置编码的小模型没有形成它（README 记录了更长训练和其他配置的尝试，同样没有形成）。 · Wei et al. 2023 在真实模型上看到：小模型几乎不理会翻转的标签，只有很大的模型会跟着翻转的示例走。 这两个世界给出一个可能的解读（只是解读，不是证明）：能不能"覆盖先验"，取决于模型学到的是"认出任务" 还是"照抄 / 学习"的机制——后者需要预训练数据里有足够多样的任务，也需要模型真的学会那个电路。""") if __name__ == "__main__": models = {w: get_model(w) for w in ("A", "B")} demo_worlds(models) demo_icl_curve(models) demo_min_ablations(models) demo_flip_and_new(models)