Ch 18 · OPD
封面/第四部分 · 后训练/Ch 18
第四部分 · 后训练 · 18

OPD — 多教师 → 单学生的合一

V4 用 on-policy reverse-KL 蒸馏替代最终 mixed-RL 阶段，并以全词表分布降低 token-level 估计的梯度方差。多教师的作用由任务数据与权重共同决定，不是 reverse KL 自动完成 teacher routing。

名词速通 · 一分钟看懂"OPD"

OPD = On-Policy Distillation：学生 $\pi_\theta$ 自己生成 trajectory，在这些状态上计算与领域 teacher 的加权 reverse KL，并使用完整词表分布训练统一 student

一句话：把 V3 的"混合 RL"整体替成"反向 KL 蒸馏"。 RL 只学 reward 期望（"做对了得多少分"），KL 蒸馏学整个分布（"在每个状态下 teacher 的概率分布是什么形状"）。 后者在每个 token 位置提供 128K 词表上的分布信息。报告的直接结论是：相比单 token KL 估计，完整分布能降低梯度方差并提高蒸馏稳定性；“无损合并”并不是论文保证。

OPD（On-Policy Distillation）
本章主角。从学生自己的策略采样 trajectory $y \sim \pi_\theta(\cdot \mid x)$，然后在每个 token 位置算 student 与 N 个 teacher 的 KL 损失。"on-policy" 是关键修饰词 —— 区别于 off-policy（用 teacher trajectory 训学生）。
KL Divergence（KL 散度）
$D_{\text{KL}}(P \| Q) = \sum_x P(x) \log \frac{P(x)}{Q(x)}$。不对称：$D_{\text{KL}}(P\|Q) \ne D_{\text{KL}}(Q\|P)$。这种不对称性不是 bug 而是设计变量 —— 选 forward 还是 reverse 决定了学生学到的"行为模式"。
Forward KL（mass-covering）
$D_{\text{KL}}(\pi_E\|\pi_\theta)$ 对 teacher 有质量、student 接近零的位置惩罚很大，因此在受限分布族中常表现出 mass-covering 倾向。若 student 表达能力充足，单 teacher 情况的最优解仍是完整复现 teacher。
Reverse KL（mode-seeking）
$D_{\text{KL}}(\pi_\theta\|\pi_E)$ 从 student 状态出发，对 student 落在 teacher 低概率区域施加较大惩罚；在容量受限、目标多峰等特定条件下常表现出 mode-seeking 倾向。多个 teacher 同时有非零权重时，最优分布更接近它们的加权几何组合，而不是自动选择其中一个。
Mode-Seeking vs Mass-Covering（直觉对比）
在“teacher 为双峰、student 被限制为单峰高斯”的示例中，forward KL 倾向覆盖两个峰，reverse KL 倾向贴近其中一个峰。这个结论依赖 student 分布族受限，不能直接当作任意神经网络和多 teacher 的解析结果。
On-Policy（在策略上采样）
损失里的采样分布是 $\pi_\theta$ 自己。每个 step 生成时用当前最新 policy。对应 reverse KL 的天然搭档 —— $D_{\text{KL}}(\pi_\theta \| \pi_E)$ 的期望就是在 $\pi_\theta$ 上取的，on-policy 让蒙特卡洛估计无偏。
Off-Policy 蒸馏（被否定的对象）
另一种做法是使用 teacher trajectories 进行监督蒸馏。它仍然有效，但训练状态分布与 student 实际生成时可能不同；on-policy 数据能直接覆盖 student 当前会到达的状态。
Full-Vocabulary OPD（全词表蒸馏）
每个 token 位置使用 teacher 的完整词表分布计算精确 KL。相比仅使用已采样 token 的 log-ratio，它利用了其它候选 token 的概率，并降低梯度估计方差；代价是更高的计算和调度复杂度。
Mixed RL（报告中的对照方法）
报告称，OPD 通过 logits-level alignment 把物理上分离的 expert 权重整合到统一参数空间，并在实践中规避传统 weight merging 或 mixed RL 常见的性能退化。报告没有进一步给出 reward 量纲、领域干扰或调参复杂度的分项实验。
一句话定位：OPD 在 student 自己生成的状态上，以领域数据和权重指定 teacher 监督，并计算全词表 reverse KL。它比 token-level 估计提供更低方差的梯度，但不能解释为 reverse KL 自动选择最强 teacher。

1. 为什么用 OPD 整合 specialist

V4 先通过 fine-tuning 和领域 RL 得到多个 specialist，再以多教师 OPD 作为主要整合方法。报告强调的直接优势有三点：

  • 按任务上下文选择相关教师：数学任务对齐数学 expert，编程任务对齐代码 expert；
  • 在 logits 层整合：不同 specialist 的能力通过输出分布监督进入同一个学生参数空间；
  • 全词表 KL：相比仅用采样 token 的 KL 估计，完整 logits 能降低梯度方差并提高训练稳定性。

报告称这种方法在实践中避开了传统权重合并或 mixed RL 常见的性能退化；它没有公布与这些方法的逐项消融，也没有把 specialist 定义成理论上的“最优 policy”。

数值演练 · reward 信号 vs 分布信号 单条 trajectory 长度 5000 token：
  • RL reward：通常为每条 trajectory 的一个或少量标量；
  • token-level KL estimate：每个 token 使用一个采样到的 log-ratio；
  • full-vocab KL：5000×128K 个 FP8 数约为 640MB，而不是 5MB；实际实现不会长期物化并保存全部 logits。
Full-vocabulary KL 确实利用了未采样 token 的概率信息，因此比单 token Monte Carlo 估计方差更低。报告没有给出“数据效率高几十到几百倍”的对比。

2. OPD 损失：多教师 reverse KL 加权和

给定 N 个 frozen 领域 expert $\{\pi_{E_1}, \dots, \pi_{E_N}\}$ 和领域权重 $\{w_i\}$，学生 $\pi_\theta$ 的 OPD 损失为：

$$ \mathcal{L}_{\text{OPD}}(\theta) \;=\; \sum_{i=1}^{N} w_i \cdot \mathbb{E}_{x \sim \mathcal{D}_i,\; y \sim \pi_\theta(\cdot \mid x)} \!\left[\, D_{\text{KL}}\!\big(\pi_\theta(\cdot \mid x, y_{<t}) \,\big\|\, \pi_{E_i}(\cdot \mid x, y_{<t})\big) \,\right] $$
📖 公式白话翻译

这条损失从内到外读：

  1. $D_{\text{KL}}(\pi_\theta \| \pi_{E_i})$：在当前 token 位置，学生分布与第 $i$ 个 teacher 分布的 reverse KL；
  2. $y \sim \pi_\theta$：trajectory 由学生自己采样 —— 这是 "on-policy" 的具体含义；
  3. $x \sim \mathcal{D}_i$：prompt 来自第 $i$ 个领域的数据集；
  4. $\sum_i w_i \cdots$：N 个领域的损失加权求和。

哪个 teacher 对当前样本起主要作用，取决于样本所属领域和分配给各 teacher 的权重 $w_i$。若多个 teacher 同时有正权重，reverse-KL 加权和不会自动忽略“较弱 teacher”；其无约束最优解与 teacher 概率的加权几何平均相关。

3. Reverse KL 与 Forward KL：适用条件与差异

这是整章最核心的洞察。先看两种 KL 的解析最优：

$$ \begin{aligned} \pi_\theta^{\star,\text{fwd}} &= \arg\min_{\pi_\theta} D_{\text{KL}}(\pi_E \| \pi_\theta) &&= \pi_E \;\;\text{（须覆盖全部 mass）} \\ \pi_\theta^{\star,\text{rev}} &= \arg\min_{\pi_\theta} D_{\text{KL}}(\pi_\theta \| \pi_E) &&= \pi_E \;\;\text{（也是它，但容量受限时塌到 mode）} \end{aligned} $$
📖 容量受限时两者发散

关键不在"无限容量下都收敛到 $\pi_E$"，而在容量受限（学生模型表达能力有限）时两者的近似行为：

  1. forward KL：在受限 student 分布族中常表现为 mass-covering；多个 teacher 的加权 forward KL 对应更接近加权算术混合的目标。
  2. reverse KL：在多峰目标与受限 student 分布族下常表现为 mode-seeking；多个 teacher 的加权 reverse KL 更接近加权几何组合，而不是离散选择某一 teacher。

V4 报告采用 reverse KL，并通过任务上下文、teacher assignment 与权重实现领域对应。这里的 mode-seeking 图示可以帮助理解单一多峰目标下的差异，但不能代替多 teacher OPD 的实际路由机制。

Demo · 1D 双峰 teacher 上 reverse / forward KL 的最优 student（容量限制为单峰高斯）
交互
teacher 双峰 $\pi_E$ forward KL student（mass-covering） reverse KL student（mode-seeking） 输出空间 $y$（如 token 位置上的 logit） 峰间距 = 3.0 · forward KL 输出形状："胖均值" · reverse KL 输出形状："单峰"

读图法：红色填充曲线是 teacher 的真实双峰分布（两个 mode 代表两个 teacher 的强项）。黄虚线是 forward KL 训出的最优单峰 student —— 胖且骑在两峰中间，两峰都覆盖但都不准（mass-covering）。绿实线是 reverse KL 的最优单峰 student —— 选其中一个峰精确复刻（mode-seeking）。
把“双峰间距”拉大，可以观察受限单峰 student 下的 mass-covering 与 mode-seeking 差异。这个 demo 不能证明 V4 会在多个 teacher 之间自动选择最强者。

4. 为什么 V4 采用 on-policy trajectories

光选 reverse KL 还不够。损失里的期望 $\mathbb{E}_{y \sim \pi_\theta}[\cdots]$ 必须从学生自己当前的策略采样 —— 这就是 "on-policy"。如果用 teacher 采样（off-policy 蒸馏），会出两个问题：

  1. 分布偏移（distribution shift）：teacher rollout 覆盖的是 teacher 的状态分布，而部署时学生会访问自己的状态分布；两者不一致时，训练信号可能集中在学生较少到达的位置；
  2. 学生状态上的监督不足：on-policy trajectory 让 teacher 分布在学生实际到达的位置上提供监督，包括学生已经偏离理想路径后的状态。纯 teacher rollout 对这些状态的覆盖通常较弱。
on-policy 的工程代价

On-policy 训练使用当前 policy 生成 trajectory，再计算 KL 损失，因此需要协调 rollout 与训练。Ch19 介绍的可抢占 rollout 和长序列并行机制用于提高这条流水的利用率与容错能力。

5. 为什么要做"全词表"，不退化到 token-level

一种工程偷懒做法是把 OPD 退化成 token-level RL：每个 token 位置只取学生与 teacher 的那一个 token 的 log 比值作为 advantage：

$$ A_t^{\text{token}} \;=\; \log \frac{\pi_E(y_t \mid x, y_{<t})}{\pi_\theta(y_t \mid x, y_{<t})} $$

把它当 RL advantage 塞进 GRPO loop，省掉算全词表 logit 的麻烦。论文给出三条否定理由：

  • 梯度方差较高：每个位置只使用采样 token 的 log-ratio，舍弃了其它词表候选的概率信息；报告称这种估计会造成高方差并可能导致训练不稳定；
  • 没真正利用 teacher 的分布信息：teacher 在 $y_t$ 上给的不仅是"$y_t$ 的概率"，更重要的是"除 $y_t$ 之外其它 token 的概率"—— 这才是判断 teacher "有多确定"的关键。token-level 完全丢掉这部分；
  • 等价 RL 的局限：token-level 蒸馏数学上就是 KL 的蒙特卡洛单样本估计，再塞进 RL 框架。这是"复用 RL loop"的便利，不是分布学习。

V4 选保留完整 logit 分布（$|V| > 100\text{K}$ 全词表），算精确 KL：

$$ D_{\text{KL}}(\pi_\theta \| \pi_E) \;=\; \sum_{v=1}^{|V|} \pi_\theta(v) \log \frac{\pi_\theta(v)}{\pi_E(v)} \quad \text{（每 token 位置都算一次）} $$
数值演练 · 全词表的工程账 单 step 训练里：
  • batch = 256 prompt × seq = 8K token = $2 \times 10^6$ 位置；
  • 每位置全词表 logit FP8 = $128\text{K} \times 1\text{ B} = 128\text{ KB}$；
  • 10 个 teacher 同框 = $2 \times 10^6 \times 128\text{ KB} \times 10 \approx \mathbf{2.6\text{ TB}}$ logit；
在这组示意假设下，单个 mini-batch 的纯 logit 张量约为 2.6 TB。即使大型集群的总显存可能超过这个数字，把 logits 同时物化并常驻仍不经济；Ch19 通过缓存 last hidden state、按 teacher 排序 batching，并在需要时重建 logits 来降低峰值占用。

6. 与 specialist 的接力

  • 每个 specialist 已经在 Ch16 走完 fine-tuning + GRPO，面向特定领域优化；
  • OPD 把 N 个 specialist 当 frozen teacher，不再更新它们的权重；
  • student 即"未来发布的 V4-Pro / V4-Flash"本体，从 Pro/Flash-Base 出发；
  • 训练目标：student 在所有领域都能 ≥ 对应 specialist 的能力 + 跨领域更强的泛化。
直觉对比 · mixed RL vs OPD

mixed RL 像让一个学生同时上 10 节不同老师的课，每个老师只给"对错"打分；
OPD 像让领域老师批改学生自己交上来的作业，并在每个 token 位置提供完整词表分布。具体由哪位老师监督，取决于样本领域与 teacher 权重。
前者主要使用 reward，后者直接对齐概率分布。不要把这种差异进一步简化成“reverse KL 自带 teacher router”。

7. 一句话总结

OPD 在 student 自己生成的状态上计算完整词表 reverse KL。相比 token-level KL estimate，它能降低梯度方差并更完整地利用 teacher 分布；多 teacher 的分工来自领域样本和权重设置，而不是 reverse KL 自动选师。相应的全词表计算、teacher 调度与 rollout 容错由 Ch19 的基础设施承担。