KL 的 KL：使用控制变量基线的在线策略蒸馏 KL for a KL: On-Policy Distillation with Control Variate Baseline

方法的具体执行步骤分为四个阶段。第一阶段：定义 OPD 的 per-token reward 和梯度。给定 prompt $x \sim \mathcal{D}$，student $\pi_\theta$ 自回归采样 $y = (y_1,\dots,y_{|y|})$，per-token reward $r_t(c_t,y_t) = \log \pi_T(y_t|c_t) - \log \pi_\theta(y_t|c_t)$ 作为不传梯度的固定值，base OPD 梯度是 $\nabla_\theta J_{\mathrm{OPD}} = \mathbb{E}[\sum_t r_t \nabla_\theta \log \pi_\theta(y_t|c_t)]$。第二阶段：闭式计算 value baseline。在同一个 forward pass 中，已经算出 student 分布 $\pi_\theta(\cdot|c_t)$ 和 teacher 分布 $\pi_T(\cdot|c_t)$，直接计算 $b(c_t) = -D_{\mathrm{KL}}(\pi_\theta(\cdot|c_t) \Vert \pi_T(\cdot|c_t)) = -\sum_v \pi_\theta(v|c_t)[\log \pi_\theta(v|c_t) - \log \pi_T(v|c_t)]$；这是一个标量，对该 token 不传梯度（用 detach 截断）。第三阶段：用 baseline 构造 advantage 与最终梯度。把 advantage 定义为 $a_t(c_t,y_t) = r_t(c_t,y_t) - b(c_t) = r_t + D_{\mathrm{KL}}(\pi_\theta \Vert \pi_T)$，vOPDfull-V 梯度就是 $\nabla_\theta J_{\mathrm{vOPD}^{\mathrm{full\text{-}V}}} = \mathbb{E}[\sum_t a_t \nabla_\theta \log \pi_\theta(y_t|c_t)]$。注意反向时 $\nabla_\theta \log \pi_\theta$ 只在被采样的单 token $y_t$ 上计算，与 base OPD 完全一致，没有引入任何额外反向 token。第四阶段（可选）：top-k 近似加速。把 baseline 中的 KL 限制在 student 的 top-k 支撑 $S_t$ 上、重整化分布后得到 $\hat{b}(c_t) = -D_{\mathrm{KL}}(\bar{\pi}_\theta \Vert \bar{\pi}_T)$，得到 vOPDtop-k 的 advantage 与梯度；因为 $\hat{b}(c_t)$ 不依赖于被采样的 $y_t$、只依赖于 $c_t$ 与 $\pi_\theta,\pi_T$，所以梯度仍然无偏。实际训练流程是：在已有 rollout 上，对每个 token 计算 KL 标量，加到 per-token reward 上得到 advantage，乘以 $\nabla_\theta \log \pi_\theta(y_t|c_t)$ 反向，权重大小为 advantage 的代数值。

技术新颖性体现在四个层面。第一，理论层面：首次为 OPD 显式推导出 value function 的闭式表达，把 OPD 严格映射到 policy-gradient RL 框架，并证明 baseline 满足无偏性与方差缩减条件（含详细的 tr(Var) 推导 Eq.(11)，指出方差缩减与 $D_{\mathrm{KL}}^2$ 成正比）。第二，算法层面：vOPD 是第一个'零额外 critic、零额外 rollout、零额外 backward token'的 OPD 稳定化方案，把 RL 中 base OPD 一直没用上的工具捡了回来。第三，理论澄清：作者严格论证了为什么 OPDfull-V 用不上 baseline——因为它的梯度本身就是对 $y$ 的精确期望，减去 $\sum_v \pi_\theta \nabla \log \pi_\theta = \nabla 1 = 0$，baseline 项自动消失，从原理上解释了两种方法在'轻量 vs 精确'轴上的对偶关系。第四，'baseline 位置差异'的洞察：top-k 放在 loss（OPDtop-k）会引入偏差，但放在 baseline（vOPDtop-k）因为不依赖于采样 token $y_t$ 而保持无偏，这个看似微小的位置差异导致实际性能天差地别——后者带来 +3%，前者几乎无收益（+0.4%）。这一观察对后续所有'在 KL 蒸馏中引入近似'的工作都有方法论意义。

独立分析下，vOPD 在以下几方面仍有可改进空间。第一，baseline 的极端行为：当某些上下文 $c_t$ 上 student 与 teacher 几乎完全不一致时（如 KL > 5 nats），baseline 会把所有负 reward token 的 advantage 几乎归零，从 Figure 1 看这些 token 出现在 y=x 上方的密集带中，作者解释为'压住噪声'，但缺乏直接证据证明这些 advantage≈0 的 token 真的全部是噪声而非被遗漏的有用信号——可能存在 teacher 在某条生成路径上的罕见但正确的偏离被 baseline 误杀。第二，wall-clock 收益依赖具体实现，Figure 3 显示在 4B 模型上 OPDtop-k 比 vOPDfull-V 还慢（32.88 vs 33.46 秒），这与理论上 OPDtop-k 用 k=20 应该明显更快的预期不符，提示 backward pass 之外的工程开销（如 tokenizer、logits gather）占据可观比例，未来需要更细致 profiling 与 kernel fusion。第三，论文只在 Qwen3 与 Olmo-3 两个家族上做了实验，且都是同源（teacher 与 student 同架构）蒸馏，对跨架构 OPD（如 Llama→Qwen）这种 token 化方式不一致的情形没有验证。第四，论文没有分析 vOPD 与近期其他 RL 改进（如 DPO-KL、SimPO-style 偏好优化）的兼容性，也没有与 RLOO/GRPO 类方法在'用 teacher 信号替代稀疏 reward'的角度做更系统的对比。改进方向包括：(1) 引入 clipped baseline 或 KL 自适应上界来避免极端零化问题；(2) 与 sequence-level KL 或带 reference model 的策略对齐目标结合，验证扩展性；(3) 探索把 baseline 做成 learned value head 的混合方案以在 top-k=5 这类极粗近似下进一步提升稳定性。

作者明确提出的未来方向包括四个：(i) 在 70B+ 规模上验证 vOPD 是否保持线性 scaling；(ii) 进一步工程优化使 vOPDtop-k 的 wall-clock 优于 vOPDfull-V（当前两者在 4B 模型上几乎并驾齐驱）；(iii) 探索 vOPD 在 self-distillation 场景的适用性，即去掉外部 teacher、用 student 自身不同 checkpoint 或 ensemble 作为伪 teacher；(iv) 把 token-level KL 扩展到 sequence-level KL，研究对应的闭式 value baseline。基于本文成果，自然可延伸的方向还有：(a) 把 vOPD 的 baseline 思想推广到带 sparse reward 的 RLVR（如把 verifiable reward 的 per-token advantage 通过闭式 KL baseline 调平），与 GRPO 形成有趣对比；(b) 探索 vOPD 的 advantage shaping 与 clipping 组合，例如借鉴 PPO 的 ratio clipping 思路防止 teacher-student 分布快速漂移；(c) 把控制变量 baseline 推广到 dual-direction KL（如同时用 forward KL $D_{\mathrm{KL}}(\pi_T \Vert \pi_\theta)$ 做对称 baseline）以研究对 mode-covering vs mode-seeking 行为的影响；(d) 既然 vOPDtop-k 对 k 不敏感，可以进一步把 baseline 计算 offload 到 CPU 或独立线程，与主反向流并行，潜在进一步压低 wall-clock。

作者提供了开源代码仓库 https://github.com/holi-lab/vOPD，方便后续研究者复现。论文中详尽列出了所有实验设置：模型使用 Qwen3-1.7B / 4B-Base、Olmo-3-7B-Think / 7B-Base、Qwen3-0.6B-Base 等，训练数据使用 DAPO-Math-17K 的英文子集（14K 样本、单 epoch）与 SciKnowEval 化学子集（75/5/20 划分），评测使用 MATH500、Minerva Math、AMC23、AIME24/25、SciKnowEval、GPQA-Diamond 六个基准，使用 avg@n 与 pass@n（n=8 或 32）。所有 vOPD 变体的默认 k=20，遵循 Li et al. (2025) 的设定。实验硬件明确为单卡 NVIDIA H200。整体复现门槛适中：模型规模最高 7B、单卡 H200 即可运行、消融实验明确给出 k 的鲁棒性范围。但仍有一些细节可能影响严格复现：(1) OPD 各变体（特别是 OPDfull-V）的具体实现细节（kernel 选择、log-sum-exp 计算方式）会影响 wall-clock，论文未完全公开；(2) 训练超参（学习率、batch size、序列长度、采样 temperature、clip ratio 等）在主文中未完整列出，需要查附录；(3) 不同随机种子下的方差没有报告，Figure 3 给出 wall-clock 的误差棒但其他指标没有平均多次运行；(4) SciKnowEval 化学子集的 train/eval/test 划分方式需要依赖引用文献 [11,15,35,48]。整体而言，理论上无偏性、闭式 KL baseline 推导与 base OPD 单 token 反向的结构都是可以直接验证的；想完整复现具体精度数字则需要较细致的工程对齐。