""" ============================================================================= CS525 · 第 10 讲模块：多领域的合成语料 + 可调大小的计数语言模型（mix_world.py） ============================================================================= 对应讲义：lectures/lecture10_data_mixing.md §6 ----------------------------------------------------------------------------- 【这个模块提供什么】 ----------------------------------------------------------------------------- MixWorld 一个"知道全部真相"的多领域世界： 6 种"成分"（通用、百科、对话、代码、数学、噪声）各有自己的"语法"； 每个数据来源（网页、维基、书、代码、数学、论坛……）是这些成分的一种混合， 有自己的数据量（独特 token 数）； 8 个"下游任务"也是成分的混合，用交叉熵（比特 / token）评测。 MixWorld.train_eval 给一个混合比例 p、训练预算 R、模型大小 width： 按比例从每个来源取 p_j · R 个 token（不够就重复），"训练"，返回 8 个任务的损失。 MixWorld.natural / temperature 两种最常见的启发式比例。 ----------------------------------------------------------------------------- 【为什么用一个"数数"的模型】 ----------------------------------------------------------------------------- 数据混合的研究需要训练**成百上千个**模型（OLMix 的一组实验就要几百次代理训练）。 真实做法是训练很多个小 Transformer；我们用一个几毫秒就能"训练"完的三元（trigram）计数模型代替： 训练 = 把选中的数据里每个"(前两个词) → 下一个词"出现的次数加起来； 预测 = 用 Witten–Bell 平滑：见过的多就相信计数，见过的少就退回到"只看前一个词"的二元模型。 它保留了数据混合问题里最重要的几个性质： (1) **预算有限**：总 token 数 R 固定，多给一个来源，就少给别的来源； (2) **来源之间有迁移，也有冲突**：几个成分共享一些"功能词"（编号 0–3），在这些上下文里， 不同来源的统计会混在一起——多放代码会让"讲代码"变好，也可能让"阅读"变差； (3) **数据量有限，重复的价值递减**：一个来源只有 N_j 个独特 token。取的比这多就得重复。 重复的 token 按 Muennighoff et al.（2023）的"有效数据"公式打折（见 _effective_epochs）； 数据稀疏时，重复还会让模型对少数见过的样本"过于自信"——平滑的回退变少； (4) **模型大小有意义**：width 控制模型能区分多少种"前两个词"的组合（见 _logp）。 小模型把很多上下文合并在一起，学不到细节，它对"哪种混合更好"的判断也可能和大模型不同。 它当然不是真的语言模型：没有泛化、没有优化过程、训练和数据顺序无关（这一点对"在线混合"很重要， 03 会讨论）。讲义 §6 列出了哪些结论依赖于这些简化。 ----------------------------------------------------------------------------- 【数据联系】 ----------------------------------------------------------------------------- * 第 6、7 讲：数据的"质量"和"数量"——这里每个来源既有质量（噪声成分占多少），也有数量（N_j）。 * 第 7 讲：重复几遍还有用（R*_D ≈ 15.4）——这里直接用了那个公式。 * 第 11 讲：合成数据——可以看作"新增一个来源"，同样要决定它占多少。 """ import numpy as np V = 40 # 词表大小：40 个 token（编号 0–39） COMPONENTS = ["通用", "百科", "对话", "代码", "数学", "噪声"] # 每个成分用哪些词。0–3 是所有成分共享的"功能词"（像 the、of、括号、空格），其余大多是各自的专有词。 # 代码和数学共享 33–35（像数字和运算符）——所以数学数据对代码任务有一点迁移，反之亦然。 SUPPORT = { "通用": list(range(0, 14)), "百科": list(range(0, 4)) + list(range(14, 22)), "对话": list(range(0, 4)) + list(range(22, 29)), "代码": list(range(0, 4)) + list(range(29, 36)), "数学": list(range(0, 4)) + list(range(33, 40)), "噪声": list(range(40)), } # Dirichlet 浓度：越小，每个上下文后面"可能的下一个词"越集中（越好预测）。 # 百科和数学取 0.5（内容更"丰富"，同样多的数据覆盖得更稀疏）；噪声取 1.0（几乎均匀，学不到什么）。 ALPHA = {"通用": 0.3, "百科": 0.5, "对话": 0.3, "代码": 0.3, "数学": 0.5, "噪声": 1.0} def _w(**kw): """把 {成分名: 权重} 变成按 COMPONENTS 顺序的向量。""" return np.array([kw.get(c, 0.0) for c in COMPONENTS]) # 原子数据来源："名字": (成分权重, 独特 token 数 N) # 网页（web）在本模块里由三个"子话题"组成；第 03 个脚本里会把它拆开（partition）。 SOURCES = { "web_news": (_w(通用=.50, 百科=.25, 对话=.05, 代码=.02, 数学=.03, 噪声=.15), 800_000), "web_chat": (_w(通用=.45, 对话=.35, 百科=.03, 代码=.02, 数学=.02, 噪声=.13), 800_000), "web_tech": (_w(通用=.35, 代码=.20, 数学=.15, 百科=.10, 对话=.05, 噪声=.15), 400_000), "wiki": (_w(百科=.65, 通用=.35), 5_000), "books": (_w(通用=.55, 对话=.45), 15_000), "code": (_w(代码=.85, 通用=.05, 噪声=.10), 30_000), "math": (_w(数学=.70, 百科=.20, 代码=.10), 2_500), "forum": (_w(对话=.35, 代码=.20, 数学=.15, 通用=.20, 噪声=.10), 10_000), # 下面两个只在 03 的"开发过程中领域会变"里出现 "stack_edu": (_w(代码=.90, 通用=.10), 8_000), # 新增：教育类代码（干净，没有噪声） "math_v2": (_w(数学=.75, 百科=.20, 代码=.05), 5_000), # 修订：重新处理过的数学数据（更多、更干净） } # "领域"（domain）= 我们给它分配比例的单位。一个领域可以由几个原子来源组成（按数据量比例抽）。 DOMAINS = { "web": ["web_news", "web_chat", "web_tech"], "wiki": ["wiki"], "books": ["books"], "code": ["code"], "math": ["math"], "forum": ["forum"], "web_news": ["web_news"], "web_chat": ["web_chat"], "web_tech": ["web_tech"], "stack_edu": ["stack_edu"], "math_v2": ["math_v2"], } BASE_DOMAINS = ["web", "wiki", "books", "code", "math", "forum"] # 01、02 用的 6 个领域 # 8 个下游任务（像 OLMix 的 52 个任务：数学、代码、问答……），每个也是成分的混合，不含噪声。 TASKS = { "知识问答": _w(百科=.6, 通用=.4), "阅读": _w(通用=.7, 对话=.3), "对话": _w(对话=.7, 通用=.3), "写代码": _w(代码=.9, 通用=.1), "讲代码": _w(代码=.5, 对话=.3, 通用=.2), "应用题": _w(数学=.6, 通用=.4), "证明": _w(数学=.6, 百科=.4), "科学": _w(百科=.5, 数学=.3, 通用=.2), } TARGET_R = 200_000 # "目标模型"的训练预算（token 数） TARGET_WIDTH = V # "目标模型"的大小：能区分全部 40 × 40 个上下文 R_STAR = 15.4 # 重复数据价值衰减的常数（Muennighoff et al. 2023 的拟合值 R*_D） class MixWorld: """ 多领域的合成世界。 【参数】seed：随机种子（决定各成分的"语法"和各来源的那一份数据）；sticky：成分的"黏性" （每个 token 以这个概率沿用上一个 token 的成分——像一篇文档里话题会持续一段）。 【属性】 names_tasks：任务名列表；QT：(任务数, V³) 每个任务的"真实"三元分布（展平）； floors：每个任务的熵下界（最好的模型也做不到比它低的损失）； src_counts：{来源: 那一份数据的三元计数 (V, V, V)}；src_N：{来源: 独特 token 数} 【思路】 1. 每个成分 c、每个上下文 (a, b)：下一个词的分布 T_c[a, b, :] ~ Dirichlet(ALPHA[c])，只放在 SUPPORT[c] 上； 2. 一个来源（或任务）= 一个"带成分的马尔可夫链"：状态是 (成分, 前两个词)。 用幂迭代求它的平稳分布，得到"(a, b, w) 三元组出现的概率" q[a, b, w]； 3. 每个来源的那一份数据 = 从 q 里抽 N 个三元组的计数（Poisson 近似多项分布）。 【边界】构造一次约 1 秒（14 条链的幂迭代）。所有随机数都来自 seed，重复构造得到完全相同的世界。 """ def __init__(self, seed=0, sticky=0.97): rng = np.random.default_rng(seed) self.T = np.zeros((len(COMPONENTS), V, V, V)) for k, c in enumerate(COMPONENTS): s = SUPPORT[c] self.T[k][:, :, s] = rng.dirichlet(np.full(len(s), ALPHA[c]), size=(V, V)) self.sticky = sticky self.src_q = {name: self._joint(w) for name, (w, _) in SOURCES.items()} self.src_N = {name: n for name, (_, n) in SOURCES.items()} rng_data = np.random.default_rng(seed + 1) self.src_counts = {name: rng_data.poisson(self.src_N[name] * self.src_q[name]) for name in SOURCES} self.names_tasks = list(TASKS) self.QT = np.stack([self._joint(w).ravel() for w in TASKS.values()]) self.floors = np.array([self._cond_entropy(q.reshape(V, V, V)) for q in self.QT]) self.first_perm = np.random.default_rng(seed + 7).permutation(V) # ------------------------------------------------------------------ 构造真实分布 def _joint(self, wc, iters=300): """ 求"带成分的马尔可夫链"的平稳三元分布 q[a, b, w]。 【参数】wc：成分权重（长度 6，会被归一化）；iters：幂迭代次数 【返回】(V, V, V) 数组，总和为 1 【思路】状态 = (成分 c, 前两个词 a, b)。一步转移：先决定下一个 token 的成分 （以 sticky 的概率不变，否则按 wc 重新抽），再按 T_c'[a, b, :] 抽下一个词 w，状态变成 (c', b, w)。 幂迭代：π ← π · 转移，迭代到收敛（黏性 0.97 意味着约 30 步换一次成分，300 次迭代足够）。 【numpy 小知识】np.einsum 用下标字符串写求和：'cab,cd->dab' 表示 ρ[d,a,b] = Σ_c π[c,a,b]·K[c,d]。 """ wc = np.asarray(wc, float) wc = wc / wc.sum() K = self.sticky * np.eye(len(wc)) + (1 - self.sticky) * wc[None, :] pi = np.einsum("c,ab->cab", wc, np.full((V, V), 1.0 / V ** 2)) for _ in range(iters): rho = np.einsum("cab,cd->dab", pi, K) pi = np.einsum("dab,dabw->dbw", rho, self.T) rho = np.einsum("cab,cd->dab", pi, K) return np.einsum("dab,dabw->abw", rho, self.T) @staticmethod def _cond_entropy(q): """H(w | a, b)（比特）：知道真实分布的模型能达到的最低损失。""" qab = q.sum(2, keepdims=True) cond = np.where(q > 0, q / np.maximum(qab, 1e-300), 1.0) return float(-(q * np.log2(cond)).sum()) # ------------------------------------------------------------------ 取数据 def domain_N(self, domain): """一个领域的独特 token 数（它所有原子来源的和）。""" return sum(self.src_N[s] for s in DOMAINS[domain]) def domain_joint(self, domain): """一个领域的真实三元分布（各原子来源按数据量加权）。03 的 DoReMi 用它算"领域上的损失"。""" parts = DOMAINS[domain] tot = self.domain_N(domain) return sum(self.src_N[s] / tot * self.src_q[s] for s in parts) @staticmethod def _effective_epochs(e, r_star=R_STAR): """ 看了 e 遍的数据"相当于"看了几遍新数据（Muennighoff et al. 2023 的有效数据公式）。 【参数】e：遍数（可以是小数）；r_star：衰减常数 【返回】e ≤ 1 时就是 e；e > 1 时 = 1 + R*·(1 − exp(−(e − 1)/R*)) 【思路】第一遍是全新的；之后每多一遍，价值按指数衰减。R* = 15.4 时： 4 遍 ≈ 3.7 遍新数据，16 遍 ≈ 10.6 遍，无穷多遍最多 ≈ 16.4 遍。 【边界】这是我们"写进世界"的假设，不是从模型里涌现出来的结论（讲义 §6 会说明）。 """ if e <= 1: return e return 1 + r_star * (1 - np.exp(-(e - 1) / r_star)) def counts(self, mix, R, seed=0, fresh=()): """ 按混合比例取训练数据，返回三元计数。 【参数】mix：{领域名: 比例}（比例之和应为 1）；R：总 token 数；seed：抽样的随机种子； fresh：一组领域名——假设这些领域"有无限多的新数据"，不用重复（只用于对照实验） 【返回】(V, V, V) 的浮点数组 【思路】领域 j 分到 p_j · R 个 token，相当于它的数据被看 f = p_j · R / N_j 遍： f < 1：每个三元组计数按二项分布抽一部分（随机取一部分文档的近似）； f ≥ 1：整份数据的计数 × 有效遍数（重复的部分打折）。 一个领域有几个原子来源时，每个来源都看 f 遍（按数据量比例抽）。 【边界】比例为 0 的领域不贡献数据；不同 seed 抽到的子集不同（模拟"换个随机种子重新训练"的噪声）。 """ rng = np.random.default_rng(seed) c = np.zeros((V, V, V)) for dom, p in mix.items(): if p <= 0: continue f = p * R / self.domain_N(dom) for s in DOMAINS[dom]: C = self.src_counts[s] if dom in fresh and f > 1: c += rng.poisson(f * self.src_N[s] * self.src_q[s]) elif f < 1: c += rng.binomial(C, f) else: c += self._effective_epochs(f) * C return c # ------------------------------------------------------------------ 模型 def logp(self, c, width=TARGET_WIDTH): """ 由计数得到模型的 log2 P(w | a, b)。 【参数】c：(V, V, V) 计数；width：模型大小，1–40。模型把"前两个词"里的第一个词 a 分成 width 类 （a 按一个固定的随机排列取余），同一类的上下文共用一组统计。width = 40 时能区分全部上下文； width = 1 时退化成"只看前一个词"的二元模型。 【返回】长度 V³ 的数组（展平的 log2 概率） 【思路】Witten–Bell 插值平滑（逐级回退）： P1(w) = (c1(w) + 1) / (n + V) （加一平滑的一元模型） P2(w | b) = (c2(b,w) + t2(b) · P1(w)) / (n2(b) + t2(b)) P3(w | h) = (c3(h,w) + t3(h) · P2(w | b)) / (n3(h) + t3(h))，h = (a 的类, b) 其中 n 是上下文出现的总次数，t 是它后面出现过多少种不同的词。 见过的多（n 大）→ 相信计数；后面的词种类多（t 大）→ 多回退。 【边界】从没见过的上下文直接用下一级的分布。重复的数据让 n 变大、t 不变 → 更相信计数 → 数据稀疏时"过于自信"。 """ cls = self.first_perm % width h = (cls[:, None] * V + np.arange(V)[None, :]).ravel() # 每个 (a, b) 属于哪个"桶" B = width * V c3 = c.reshape(V * V, V) idx = (h[:, None] * V + np.arange(V)).ravel() cb = np.bincount(idx, weights=c3.ravel(), minlength=B * V).reshape(B, V) # 同一个桶的计数加起来 tb, nb = (cb > 0).sum(1), cb.sum(1) c2 = c.sum(0) t2, n2 = (c2 > 0).sum(1), c2.sum(1) c1 = c2.sum(0) P1 = (c1 + 1) / (c1.sum() + V) den2 = n2 + t2 P2 = np.where(den2[:, None] > 0, (c2 + t2[:, None] * P1[None, :]) / np.maximum(den2, 1e-12)[:, None], P1[None, :]) P2r = np.tile(P2, (V, 1)) # 第 a·V + b 行 = P2[b] num = cb[h] + tb[h][:, None] * P2r den = (nb + tb)[h] P3 = np.where(den[:, None] > 0, num / np.maximum(den, 1e-12)[:, None], P2r) return np.log2(P3).ravel() def task_losses(self, logp): """8 个任务上的交叉熵（比特 / token）：−Σ q_task(a,b,w) · log2 P(w|a,b)，精确计算，不抽样。""" return -(self.QT @ logp) def domain_losses(self, logp, domains): """在若干领域自己的分布上的交叉熵（03 的 DoReMi 用）。""" return np.array([-(self.domain_joint(d).ravel() @ logp) for d in domains]) def train_eval(self, mix, R=TARGET_R, width=TARGET_WIDTH, seed=0, fresh=()): """ 一次"训练 + 评测"。 【参数】mix：{领域: 比例}；R：预算；width：模型大小；seed：数据抽样的种子；fresh：见 counts 【返回】长度 8 的数组：每个任务的损失（比特 / token，越低越好） 【边界】大约 3 毫秒。同样的参数和 seed 得到完全相同的结果。 """ return self.task_losses(self.logp(self.counts(mix, R, seed, fresh), width)) # ------------------------------------------------------------------ 启发式比例 def natural(self, domains): """自然分布：按每个领域的数据量分配（"有多少用多少"）。""" n = np.array([self.domain_N(d) for d in domains], float) return n / n.sum() def temperature(self, domains, tau): """温度采样：p_j ∝ N_j^(1/τ)。τ = 1 是自然分布，τ → ∞ 是均匀分布（mT5、XLM-R 在多语言数据上的做法）。""" n = np.array([self.domain_N(d) for d in domains], float) ** (1.0 / tau) return n / n.sum() def as_mix(domains, p): """把领域名列表和比例向量配成字典。""" return {d: float(x) for d, x in zip(domains, p)} if __name__ == "__main__": world = MixWorld() print("每个任务的熵下界（比特）：", " ".join(f"{t} {h:.3f}" for t, h in zip(world.names_tasks, world.floors))) print("各领域的独特 token 数：", {d: world.domain_N(d) for d in BASE_DOMAINS}) for name, p in [("自然分布", world.natural(BASE_DOMAINS)), ("均匀", np.ones(6) / 6), ("温度 τ=3", world.temperature(BASE_DOMAINS, 3))]: loss = world.train_eval(as_mix(BASE_DOMAINS, p)) print(f"{name}：比例 {np.round(p, 3)} 平均损失 {loss.mean():.4f}")