""" ============================================================================= CS486 · 第 1 讲共享模块：hop_world.py —— "多跳指针追踪"推理世界（标准库） ============================================================================= 【这个模块要解决什么问题】 研究"LLM 会不会推理"有一个根本困难：真实模型在海量网页上训练过，我们很难知道它答对一道题 是"推理出来的"还是"背下来的"、"凑巧蒙对的"。所以推理研究里常用一种办法： **造一个规则完全已知、答案可以机械验证、难度可以精确控制的合成任务**， 在上面从零训练小模型，观察它学会了什么、没学会什么（"Physics of Language Models"、 "Grokked Transformers are Implicit Reasoners"、"Two-Hop Curse" 等工作都是这样做的）。 本模块就是这样一个任务：**多跳指针追踪（multi-hop pointer chasing）**。 题目： d>a;a>c;b>d;c>b;3?a 含义： "d 指向 a；a 指向 c；b 指向 d；c 指向 b。从 a 出发走 3 步，到哪里？" 求解： a → c → b → d，答案是 d 几个符号： "x>y" 一条陈述（statement）：x 指向 y。陈述之间用 ";" 分隔。 "3?a" 问题（query）：从 a 出发走 3 跳（hop）。跳数 k 是一个数字 1–9。 陈述的顺序默认是随机打乱的——模型必须"按名字去找"下一条陈述，而不能"按位置往后读"。 【为什么选这个任务（直觉）】 1. 它是**天生串行**（inherently serial）的：要知道第 2 跳去哪，必须先知道第 1 跳到了哪。 这正是"多步推理"最核心的结构——数学题的中间结果、多跳问答 （"《X》导演的母亲是谁？"先找导演、再找母亲）、执行一段程序时追踪变量…… 2. **每一跳本身很简单**：只是"在上下文里找到以当前字母开头的那条陈述，抄下它指向的字母"。 这恰好是 Transformer 最擅长的一种操作（induction head 式的"查找 + 复制"）。 所以难点只在"把很多跳串起来"，不在单步计算——这让我们能干净地研究"串行深度"。 3. 答案可以机械验证（solve 函数），难度由 k（跳数）、n（节点数）、干扰陈述数精确控制， 题目无穷多，训练数据和测试数据可以严格分开采样。 【输出格式（三种"作答方式"，对应讲义 §3 的三个实验条件）】 同一道题可以要求模型用三种方式作答（format_example）： direct（直接回答）：输入 "...;3?a"，要求紧接着输出 "d$" ——一次前向计算（forward pass）就要给出答案。 pause（停顿 / 填充 token）：输入 "...;3?a___"（k 个下划线 "_" 作为"思考空位"）， 要求接着输出 "d$"。填充 token 给了模型更多的"位置"去计算， 但这些位置上**没有写下任何中间结果**（Pfau et al. 2024 "Let's think dot by dot"）。 cot（思维链）：输入 "...;3?a"，要求输出 "cbd|d$"——先把每一跳到达的字母写出来， 再在 "|" 后面给答案。每写一个字母，它就成为下一步的输入。 "$" 是结束符。三种方式的唯一区别是"答案前面有什么"，题目完全相同，这样比较才公平。 【这个世界简化了什么（务必记住）】 - 没有自然语言、没有世界知识、没有歧义；真实 LLM 的推理要难得多也乱得多。 - 节点只有 16 个字母，陈述只有一种形式。我们研究的是**机制**（串行深度、捷径、长度泛化）， 不是量级。讲义会说明哪些结论能推广到真实 LLM、哪些不能。 【运行】 python3 hop_world.py # 打印几道示例题和三种格式，并做自检（求解器 vs 构造答案） 【依赖】只用 Python 标准库。 """ import random from dataclasses import dataclass # --------------------------------------------------------------------------- # 词表（vocabulary） # --------------------------------------------------------------------------- # 【思路】我们用"字符级"词表：每个字符就是一个 token。真实 LLM 用 BPE 等子词分词， # 但对这个任务来说字符级最简单，也让"一个 token = 一个字母 = 一跳"的对应关系一目了然。 LETTERS = "abcdefghijklmnop" # 16 个节点名 DIGITS = "0123456789" # 跳数 k（只用 1–9） SYMBOLS = ">;?|$_ " # 指向、分隔、提问、答案标记、结束、填充（pause）、补齐（padding） VOCAB = LETTERS + DIGITS + SYMBOLS # 【Python 小知识】字典推导式 {c: i for i, c in enumerate(VOCAB)} 把每个字符映射到它的编号， # enumerate 同时给出下标 i 和元素 c。STOI = "string to index"。 STOI = {c: i for i, c in enumerate(VOCAB)} PAD = " " # 补齐用的字符：一批序列长短不一时，短的在右边补空格 PAD_ID = STOI[PAD] FILLER = "_" # pause 模式的填充 token END = "$" # 结束符 ANSWER_MARK = "|" # cot 模式里"推理结束、下面是答案"的标记 MODES = ("direct", "pause", "cot") ORDERS = ("shuffle", "sorted", "reversed") def encode(s): """ 把字符串变成 token 编号列表。 【参数】s：只含 VOCAB 里字符的字符串 【返回】list[int] 【边界】遇到词表外的字符会抛 KeyError——这是故意的：静默地丢掉未知字符会让 bug 很难发现。 """ return [STOI[c] for c in s] def decode(ids): """把 token 编号列表变回字符串（encode 的逆操作）。""" return "".join(VOCAB[i] for i in ids) # --------------------------------------------------------------------------- # 一道题 # --------------------------------------------------------------------------- @dataclass class HopProblem: """ 一道多跳指针追踪题。 【字段】 statements ：陈述列表，按"显示给模型的顺序"排列，例如 ["d>a", "a>c", "b>d", "c>b"] start ：起点字母 k ：跳数（1–9） walk ：真实的行走路径，walk[0] = start，walk[i] = 走 i 跳后的位置；长度 k + 1 n_nodes ：起点所在的环有几个节点 n_distract ：干扰陈述的条数（它们组成另一个与答案无关的环） 【Python 小知识】@dataclass 自动生成 __init__ 等方法，下面的 @property 让方法像属性一样访问： 写 prob.answer 而不是 prob.answer()。 """ statements: list start: str k: int walk: list n_nodes: int n_distract: int @property def answer(self): """正确答案：走 k 跳后到达的字母。""" return self.walk[-1] @property def prompt(self): """ 给模型看的题面，例如 "d>a;a>c;b>d;c>b;3?a"。 【为什么问题放在最后、起点字母放在最末尾】 模型从左往右读，读到最后一个字符时就要开始作答。把起点字母放在题面的**最后一个位置**， direct 模式的"第 1 跳"和 cot 模式的"第 1 个推理步"就是完全相同的计算： "在当前位置，找到以当前字母开头的陈述，输出它指向的字母"。 这样比较 direct 和 cot 时，差别只在"能不能把中间结果写下来再接着算"。 （我们在原型里先试过把起点放在 "?a3=" 中间的版本，direct 模式连 1 跳都学不会—— 见代码 README 的"关于实验结果"。格式细节对小模型的可学习性影响很大。） """ return ";".join(self.statements) + f";{self.k}?{self.start}" @property def trace(self): """思维链：依次写出每一跳到达的字母，例如 "cbd"（不含起点）。""" return "".join(self.walk[1:]) @property def n_statements(self): """陈述总条数（= 环的节点数 + 干扰陈述数）；决定题面长度。""" return len(self.statements) def make_problem(rng, n, k, order="shuffle", n_distract=0): """ 随机生成一道题。 【思路】 1. 从 16 个字母里不重复地抽 n + n_distract 个。 2. 前 n 个组成一个**环**：x0 → x1 → … → x(n-1) → x0，起点是 x0。 为什么用环？这样每个字母都恰好有一条出边，不存在"没有出边的终点"—— 否则模型可以用一个捷径（"找那个从来不出现在箭头左边的字母"）不经推理就猜出终点。 k 可以大于等于 n：那就绕着环走（k = n 时回到起点）。 3. 后 n_distract 个字母组成**另一个独立的环**（干扰陈述）：它们和答案无关， 只是让题面更长、信息更杂——类似 GSM-NoOp 往数学题里加无关的句子。 4. 按 order 排列陈述： "shuffle" ：全部打乱（默认，训练用） "sorted" ：按行走顺序排列（x0>x1; x1>x2; …），干扰陈述放在最后 "reversed" ：按行走顺序的逆序排列，干扰陈述放在最后 【参数】 rng ：random.Random 实例（用独立的随机数生成器，结果才可复现） n ：环的节点数，2 ≤ n，且 n + n_distract ≤ 16 k ：跳数，1 ≤ k ≤ 9（题面里只用一个数字表示 k） order ：陈述顺序，见上 n_distract ：干扰环的节点数；0 表示没有干扰；不能等于 1（1 个节点的环 "x>x" 是自环， 也允许，但我们约定干扰环至少 2 个节点，免得出现退化的题） 【返回】HopProblem 【边界】参数不合法时抛 ValueError（而不是悄悄生成一道错题）。 【Python 小知识】rng.sample(seq, m) 从序列里不重复地随机抽 m 个；rng.shuffle(lst) 原地打乱列表。 (i + 1) % n 是取余：i = n - 1 时回到 0，这就是"环"。 """ if not (2 <= n and n + n_distract <= len(LETTERS)): raise ValueError(f"n={n}, n_distract={n_distract} 超出字母表大小 {len(LETTERS)}") if not (1 <= k <= 9): raise ValueError(f"k={k} 必须在 1–9 之间（题面只用一个数字表示跳数）") if n_distract == 1: raise ValueError("干扰环至少要 2 个节点") if order not in ORDERS: raise ValueError(f"order 必须是 {ORDERS} 之一，收到 {order!r}") names = rng.sample(LETTERS, n + n_distract) cycle, distract = names[:n], names[n:] chain_statements = [f"{cycle[i]}>{cycle[(i + 1) % n]}" for i in range(n)] distract_statements = [f"{distract[i]}>{distract[(i + 1) % n_distract]}" for i in range(n_distract)] if order == "shuffle": statements = chain_statements + distract_statements rng.shuffle(statements) elif order == "sorted": statements = chain_statements + distract_statements else: # "reversed" statements = chain_statements[::-1] + distract_statements # [::-1] 是把列表倒过来 walk = [cycle[i % n] for i in range(k + 1)] return HopProblem(statements=statements, start=cycle[0], k=k, walk=walk, n_nodes=n, n_distract=n_distract) def solve(prompt): """ 机械求解器：只看题面字符串，按规则一步步走，返回答案和路径。 【思路】这是"标准答案"的独立来源。make_problem 构造题目时已经知道答案， 但我们再写一个只看题面的求解器来核对——如果两者不一致，说明题面的生成或格式有 bug。 研究中这叫"对数据做单元测试"：合成数据的 bug 会直接变成"模型学不会"的假象。 【参数】prompt：形如 "d>a;a>c;...;3?a" 的题面 【返回】(answer, walk)：答案字母和完整路径（含起点） 【边界】题面格式不对（缺 "?"、陈述不是 "x>y"、某个字母没有出边）时抛 ValueError。 同一个字母有两条出边时以最后一条为准（我们的生成器不会产生这种题）。 """ if "?" not in prompt: raise ValueError("题面里没有 '?'") body, query = prompt.rsplit(";", 1) # rsplit 从右边切一刀：最后一段是问题 k_str, start = query.split("?") succ = {} for st in body.split(";"): if len(st) != 3 or st[1] != ">": raise ValueError(f"陈述格式不对：{st!r}") succ[st[0]] = st[2] walk = [start] for _ in range(int(k_str)): if walk[-1] not in succ: raise ValueError(f"字母 {walk[-1]} 没有出边") walk.append(succ[walk[-1]]) return walk[-1], walk # --------------------------------------------------------------------------- # 三种作答格式 # --------------------------------------------------------------------------- def format_example(prob, mode): """ 把一道题变成 (模型输入, 期望输出) 两段字符串。 【参数】prob：HopProblem；mode："direct" / "pause" / "cot" 【返回】(input_str, target_str) direct：("…;3?a", "d$") pause ：("…;3?a___", "d$") —— k 个填充 token 由我们提供，模型只需在最后作答 cot ：("…;3?a", "cbd|d$") —— 推理链由模型自己生成 【为什么 pause 的填充 token 数 = k】 这样 pause 和 cot 在"答案前有多少个位置可以用来计算"上是一样的（计算量匹配，compute-matched）。 如果 cot 赢了 pause，就不能说"只是因为 cot 多算了几个位置"——这是一个**对照实验** （control）的设计：只改变一个因素（中间位置上有没有写下中间结果）。 【边界】mode 不合法时抛 ValueError。 """ if mode == "direct": return prob.prompt, prob.answer + END if mode == "pause": return prob.prompt + FILLER * prob.k, prob.answer + END if mode == "cot": return prob.prompt, prob.trace + ANSWER_MARK + prob.answer + END raise ValueError(f"mode 必须是 {MODES} 之一，收到 {mode!r}") def max_new_tokens(mode, k): """ 评测时最多让模型生成几个 token。 【思路】direct / pause 只需要"答案 + 结束符"2 个 token；cot 需要 k 个推理字母 + "|" + 答案 + "$"。 我们多给 2 个余量，让模型"多写"的错误也能被看到（而不是被截断掩盖）。 """ return 2 if mode in ("direct", "pause") else k + 5 def parse_answer(generated, mode): """ 从模型生成的字符串里取出答案字母。 【参数】generated：模型在输入之后生成的字符串（可能包含 "$" 之后的垃圾）；mode：作答格式 【返回】答案字母（str），解析失败时返回 None 【边界】 direct / pause：取第一个字符；如果它不是字母，返回 None。 cot：取 "|" 后面的第一个字符；没有 "|"（模型没写完推理就停了，或者写了太多）返回 None。 注意：cot 模式我们**只看 "|" 后面的答案**，不检查推理链本身对不对—— "推理链错了但答案对了"和"推理链对了但答案抄错了"都可能发生（讲义 §4.4；第 6 讲讨论忠实性）。 """ if mode in ("direct", "pause"): c = generated[:1] return c if c and c in LETTERS else None if ANSWER_MARK not in generated: return None c = generated.split(ANSWER_MARK, 1)[1][:1] return c if c and c in LETTERS else None # --------------------------------------------------------------------------- # 训练分布与评测集 # --------------------------------------------------------------------------- def sample_train_problem(rng, n_range=(4, 8), k_fixed=None, order="shuffle"): """ 从训练分布里抽一道题。 【训练分布（全讲统一）】 n（环的节点数）在 n_range 里均匀取整数；默认 4–8。 k（跳数）：k_fixed 为 None 时在 1 … min(n-1, 9) 里均匀取（不绕环；题面只能写一位数的跳数， 所以 n ≥ 11 时最多 9 跳）；否则固定为 k_fixed。 没有干扰陈述；顺序由 order 决定（默认打乱）。 【参数】rng：random.Random；n_range：(最小, 最大) 节点数；k_fixed：None 或固定跳数；order：陈述顺序 【返回】HopProblem 【边界】k_fixed ≥ n 时会绕环（合法，但改变了题目的性质）；调用方自己负责选合理的组合。 """ n = rng.randint(n_range[0], n_range[1]) # randint 两端都包含 k = k_fixed if k_fixed is not None else rng.randint(1, min(n - 1, 9)) return make_problem(rng, n, k, order=order) def make_eval_set(n, k, count, order="shuffle", n_distract=0, seed=0): """ 生成一个固定的评测集：count 道 (n, k, order, n_distract) 相同的题。 【思路】评测集用**独立的种子**生成，并且种子由参数决定（同样的参数永远得到同样的题）， 这样不同模型、不同脚本在"同一张考卷"上比较——这是配对比较（paired comparison）的前提（代码 04）。 训练数据来自另一个随机数生成器，二者"统计上独立"。 小题目（例如 n = 4、k = 1）的可能题目总数有限，训练时很可能见过某些评测题—— 但模型见过的是"同一道题"而不是"答案的规律"，对指针追踪这种每道题答案都不同的任务影响很小； 第 2 讲讨论真实 LLM 的数据污染（contamination）为什么严重得多。 【参数】n、k、order、n_distract：同 make_problem；count：题数；seed：评测集种子 【返回】list[HopProblem] 【Python 小知识】hash 字符串在不同运行之间会变（PYTHONHASHSEED），所以种子我们用普通的算术拼出来， 不用 hash()。 """ order_id = ORDERS.index(order) rng = random.Random(1_000_003 * seed + 10_007 * n + 101 * k + 11 * n_distract + order_id) return [make_problem(rng, n, k, order=order, n_distract=n_distract) for _ in range(count)] # --------------------------------------------------------------------------- # 自检 # --------------------------------------------------------------------------- def _self_test(): """ 自检：随机生成 2000 道各种参数的题，用独立的求解器核对答案、路径和格式。 """ rng = random.Random(123) for _ in range(2000): n = rng.randint(2, 10) nd = rng.choice([0, 0, 2, 3, 4]) if n + nd > len(LETTERS): nd = 0 k = rng.randint(1, 9) order = rng.choice(ORDERS) p = make_problem(rng, n, k, order=order, n_distract=nd) ans, walk = solve(p.prompt) assert ans == p.answer and walk == p.walk, (p, ans, walk) assert p.n_statements == n + nd for mode in MODES: inp, tgt = format_example(p, mode) encode(inp + tgt) # 所有字符都在词表里 assert parse_answer(tgt, mode) == p.answer # 期望输出能被正确解析 # 评测集可复现：同样的参数 → 同样的题 a = make_eval_set(8, 3, 5, seed=7) b = make_eval_set(8, 3, 5, seed=7) assert [x.prompt for x in a] == [x.prompt for x in b] # 解析器的边界情况 assert parse_answer("", "direct") is None assert parse_answer("cb", "cot") is None # 没有 "|" assert parse_answer("cbd|", "cot") is None # "|" 后面是空的 print("自检通过：2000 道随机题的构造答案 = 求解器答案；三种格式都能被正确解析；评测集可复现。") if __name__ == "__main__": rng = random.Random(0) print(f"词表大小 = {len(VOCAB)}：{VOCAB!r}\n") p = make_problem(rng, n=5, k=3) print("一道题（打乱顺序，n = 5，k = 3）：") print(f" 题面 prompt : {p.prompt}") print(f" 路径 walk : {' → '.join(p.walk)}") print(f" 答案 answer : {p.answer}\n") for mode in MODES: inp, tgt = format_example(p, mode) print(f" {mode:<6} 输入 = {inp:<28} 期望输出 = {tgt}") print() for order in ORDERS: q = make_problem(random.Random(1), n=5, k=2, order=order, n_distract=3) print(f" 顺序 {order:<8}（含 3 条干扰陈述）：{q.prompt} 答案 {q.answer}") print() _self_test()