面向大语言模型的混合策略蒸馏方法 Hybrid Policy Distillation for LLMs
将前向KL与反向KL按token级信号混合，结合off-policy数据与轻量on-policy采样的统一蒸馏框架
前置知识
知识蒸馏（KD）
知识蒸馏是把一个强大的「教师」模型的输出知识迁移到更小、更便宜的「学生」模型中的一类训练范式。根据能否访问教师内部概率分布，可分为 black-box KD（只能拿到教师生成的文本，等价于监督微调 SFT）和 white-box KD（可拿到教师对每个 token 的完整概率分布，用 KL 散度之类的分布级目标来对齐）。
本文的核心是把 SFT、forward KL、reverse KL 等方法统一成同一个 reweighted log-likelihood 框架，因此必须先理解 KD 与 SFT 的区别。
前向 KL 与反向 KL（FKLD / RKLD）
FKLD 定义为 $D_{KL}(p \| q_\theta) = \mathbb{E}_{a\sim p} \log \frac{p(a)}{q_\theta(a)}$，鼓励学生模型覆盖教师的所有模式（mode-covering），但容易让概率分布过度平滑；RKLD 则是 $D_{KL}(q_\theta \| p) = \mathbb{E}_{a\sim q_\theta} \log \frac{q_\theta(a)}{p(a)}$，鼓励学生聚焦教师的高概率模式（mode-seeking），但可能在学生与教师差距大时梯度高方差、训练不稳定。
HPD 整个设计就是在 token 级别把 forward 和 reverse KL 的优势混合起来，不懂这两个目标就理解不了论文的核心机理。
Off-policy 与 On-policy 数据（OPD）
Off-policy 蒸馏指学生只用一个固定的、由教师（或人工）生成的数据集 $\mathcal{D}$ 来模仿；On-policy 蒸馏（OPD）则是让学生在训练时自己采样 rollout，再用这些 rollout 做对齐。On-policy 数据更对齐训练-推理行为，但计算开销大；off-policy 高效但有 train-inference mismatch。
HPD 想要兼得两者，所以必须理解这两种数据源的优缺点。
K1 估计器（K1 estimator）与重要性采样
KL 散度的精确计算需要遍历整个词表，最简单的 Monte Carlo 估计器即 K1：$K_1 = \frac{1}{N} \sum_i \log \frac{q_\theta(a^{(i)})}{p(a^{(i)})}$，$a^{(i)} \sim q_\theta$。它是 RKLD 的无偏估计，但方差大且 log-ratio 大量为负。Schulman 2020 提出的 negative K1 常被用作 token 级奖励信号。
HPD 直接用 negative K1 同时作用在 offline 专家 token 和学生采样 token 上，并据此设计权重，没有它就读不懂公式 11-14。
研究动机
把大语言模型蒸馏到小模型时，现有三类因素——「KL 的方向（前向 vs. 反向）」「优化策略（loss vs. reward）」「数据来源（off-policy vs. on-policy）」——彼此纠缠，单独调任何一项都得不到最佳效果。具体痛点有三：FKLD 这种「模式覆盖」型分布在学生容量不足时会让概率质量摊得太薄（over-smoothed），损害生成质量；RKLD 这种「模式聚焦」型分布在师生差距大时梯度方差高，训练不稳定；而 on-policy 蒸馏需要学生自己生成完整 rollout，计算成本极高、且教师侧仍然要做一次完整 forward 带来 distribution shift；off-policy 蒸馏又会因为学生没见过的状态导致 train-inference 行为不一致（如图 1d 中 SFT 的测试时熵与训练时严重偏离）。在数学推理这种长生成任务上，这些问题尤为明显：作者在表 2 中指出，用 LLaMA-3 8B 去蒸馏 1B 学生时，FKLD、RKLD、JSD 等方法平均仅有 15 左右分，远不能与 SFT 后接 OPD 的两阶段管线相比，而后者又贵很多。
本文的目标是本文的核心目标是设计一种 single-stage、统一 token 级的蒸馏目标 HPD（Hybrid Policy Distillation），它要在三个维度上同时取得好处：使用 FKLD 的「避免漏掉模式」特性，使用 RKLD 的「聚焦主模式 + 抑制不合理采样」特性，并且只做极轻量的 on-policy 近似（仅采样一个 next token 而非完整 rollout），从而比 SFT + OPD 的两阶段管线更便宜，又能比单独的 SFT / FKLD / RKLD / JSD 更准、训推行为更对齐。
与已有工作不同的是，现有工作大多只动了「三大设计轴」中的一个：比如 SeqKD 只调数据（用教师生成当 SFT 数据），MiniLLM/GKD 只强调 reverse KL 用于 on-policy，ABKD 用 $\alpha\text{-}\beta$ divergence 整体调整概率质量分配，DistiLLM-2 用 contrastive 同时增专家似然 / 抑学生似然，但没有一篇把这些因素放进同一个 token 级 reweighted log-likelihood 框架来分析、并提供一种 principle + efficient 的混合机制。HPD 的独特切入角度正是：用同一套 reweighted likelihood 语言同时描述 SFT / FKLD / RKLD / JSD，并在此基础上自然地「内嵌」一种 hybrid 前反向 KL 的 token 级权重，让 forward KL 提供 dense signal、reverse KL 提供 mode-seeking + 抑制异常 sample 的修正，仅以一个额外的 student-sampled token 计算 negative K1 即可，几乎无额外推理成本。
核心方法
直觉上，HPD 想做的事是「在每一个 token 位置，先看学生是否低估了 ground-truth（专家）token；若低估就加大学习权重（前向 KL 风格的补救），再看学生自己采样的某个非专家 token 是否被高估（reverse-KL 风格的修正）；若也被高估，就把概率 mass 重定向到专家 token，并同时把该专家 token 的前向 KL 权重再翻倍」。技术路线上，论文先把 SFT / FKLD / RKLD 全部改写成同一个 reweighted log-likelihood 目标 $\mathcal{L}(\theta) = -\mathbb{E}_{(s_t,a_t)\sim D_\pi} w(a_t|s_t) \log q_\theta(a_t|s_t)$，再设计 $w^*(t)$ 和 $w(t)$ 两个 token 级权重：在 offline 专家 token $a^*_t$ 上用 negative K1 估计器 $k_1 = q_\theta(a^*_t|s_t)[\log p(a^*_t|s_t) - \log q_\theta(a^*_t|s_t)]$ 做门控，正负号决定触发前向 KL 或反向 KL 抑制；在学生采样的非专家 token $a_t \sim q_\theta$ 上仅当 $k'_1<0$ 才下发负权重，并把节省下来的概率质量通过双倍权重重定向到专家 token。这种 hybrid 是「在同一 token、同一 batch 内自动决定」要 forward 还是 reverse，因此不需要调任何额外超参。
和已有方法的本质区别在于：(1) HPD 不是像 ABKD / Wu et al. 2025 那样用固定权重把 $\alpha\text{-}\beta$ divergence 或 weighted FKLD+RKLD 组合起来，而是按 token 状态「动态切换 + 协同」，关键区别是 HPD 引入了「专家 token 重定向机制」——当学生在采样非专家 token 上 high-confidence 错估（$k'_1<0$）时，把专家 token 的前向 KL 权重翻倍（式 14），从而把被错误高估的概率拉回正确位置；(2) HPD 不需要完整的 on-policy rollout，只在 ground-truth prefix $s_t$ 上采样一个 next token，计算代价与一次 forward 等价，但达到了 GKD / MiniLLM 的 on-policy 效果；(3) 整套机制完全基于公式 10 给出的「正权重推升、负权重重分配概率 mass」的梯度解释，提供了 FKLD 提供 dense signal、RKLD 抑制过平滑行为的统一视角。
方法步骤详情
HPD 的训练步骤在 Algorithm 1 中给出，对每个 $(s_t, a^*_t)$ 输入 student $q_\theta$、teacher $p$、dataset $\mathcal{D}$：(1) 计算专家 token 的对数概率 $(\log q^*, \log p^*) = (\log q_\theta, \log p)(a^*_t|s_t)$；(2) 用式 (11) 计算专家 token 上的 negative reverse-K1 gap：$k_1 = q_\theta(a^*_t|s_t)(\log p^* - \log q^*)$，正值代表学生低估了专家 token；(3) 在同一状态 $s_t$ 下让学生采样一个不同的 token $a_t \sim q_\theta(\cdot|s_t)$，同样算其 K1 gap：$k'_1 = q_\theta(a_t|s_t)[\log p(a_t|s_t) - \log q_\theta(a_t|s_t)]$；(4) 按式 (14) 决定专家权重 $w^*_t$：当 $k_1>0$ 且 $k'_1<0$ 时取 $2p(a^*_t|s_t)+k_1$（双倍重定向），当 $k_1<0$ 时取 $k_1$（反向抑制），否则取 $p(a^*_t|s_t)+k_1$；(5) 按式 (13) 决定采样 token 权重 $w_t = \mathbb{I}[a_t\neq a^*_t]\cdot\mathbb{I}[k'_1<0]\cdot k'_1$，仅当其为负时才下发修正；(6) 用式 (15) 的 reweighted log-likelihood $\mathcal{L}_{HPD} = -\mathbb{E}[w^*_t \log q_\theta(a^*_t|s_t) + w_t \log q_\theta(a_t|s_t)]$ 反向更新 $\theta$。整套流程只比 SFT 多了：教师的一次额外 forward + 学生采样一次额外 token + 两组 K1 计算，没有 teacher rollout、没有 student rollout。
技术新颖性
技术新颖性体现在三方面：第一，给出了一种统一的 reweighted log-likelihood 视角（式 9 + 表 1），把 SFT / FKLD / RKLD / JSD 用同一个 $w(a_t|s_t)$ 表征并解释「正权推升、负权重分配」的梯度传播机制（式 10），这是此前工作都没系统做过的；第二，提出「专家 token 在 $k_1>0$ 且 $k'_1<0$ 时权重翻倍」的 Reinforce 设计，巧妙地利用了同一 batch 内采样 token 的负 K1 信号来指导 expert token 权重的强化，避免了显式的 coefficient 调参；第三，把 on-policy 蒸馏从「完整 rollout」简化到「单一 next-token 采样」，在和 GKD / MiniLLM 同样的效果区间里获得了 OPD 的 train-inference alignment 优势，但开销降了一个量级。
实验结果
在 Qwen2.5（7B→1.5B/3B）和 LLaMA-3（8B→1B/3B）两个模型族六组学生-教师组合上，HPD 在 AIME24/25、AMC、MATH、OlympiadBench、GPQA 六个数学推理基准上一致优于 SFT、SeqKD、RKLD、JSD 等所有基线（表 2），且绝大多数差异统计显著（$p<0.01$）。最具代表性的提升是：LLaMA-3 1B 在 off-policy 下 HPD 平均分 18.01，相比 SFT 的 13.88 提高约 30%，相比 RKLD 的 15.07 提高约 19.5%；3B 规模时 HPD 平均分 34.56 vs SFT 29.45（+17.3%），相比 RKLD 30.41 提升约 13.6%。同时 HPD 把 Qwen 2.5 3B 平均分从 MS（monkey see）的 28.25 推到 39.83，相对提升 41.0%，把 LLaMA 3 3B 从 19.43 推到 34.56，提升 77.9%——这种「跨过 model size」的效果在论文中用「学生获得接近更大模型能力」来描述。代码任务上（表 4），DS-Coder 6.7B→1.3B 平均 pass@1 从 KD 64.60 / JSD 64.10 提到 HPD 66.35，Qwen-Coder 7B→1.5B 从 KD 72.45 / JSD 76.00 提到 77.35。Personalization 任务（表 3，Qwen 2.5 7B→1.5B，对话）上，HPD 取得 AlpacaEval2 LC 13.75%、AE-WR 14.25%、Arena-WR 21.80% —— 三项均最高，MT-1T 7.23、MT-2T 5.84 同样领先。On-policy 推理（表 5）上，HPD 单阶段就比 SFT+OPD 两阶段的 29.56 (Qwen) / 17.06 (LLaMA) 更优（30.24 / 18.01），HPD+OPD 进一步推到 33.41 / 20.40。Figure 1 显示 HPD 训练时熵不坍塌、KLD 持续下降、性能单调提升；Figure 1d 显示测试时熵与训练时高度一致，说明训推行为对齐。Figure 2 在 OPD 框架下进一步证明 HPD 初始化带来更小的 KL、更平稳的 mean advantage 估计和更一致的测试熵。图 3 消融显示，去掉学生采样会导致训练 quick plateau，去掉 Reinforce 操作会让 KL 收敛变慢。表 6/7 的 HPD+DPO 实验证明 HPD 的初始化让偏好对齐阶段进一步获得 +11.92 个 Arena 综合点，且可做 self-distillation evolving 持续提升（最多 +7.94）。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 数学推理 (Qwen 2.5 7B→1.5B) | 6个基准平均分 (AIME24/25 avg@32, AMC/Math/Oly/GPQA avg@8) | 30.24 | SFT 22.80 / SeqKD 26.27 / RKLD 26.15 / JSD 26.43 | 相对最佳基线 JSD 提升 +14.5%；相对 SFT 提升 +32.6% |
| 数学推理 (Qwen 2.5 7B→3B) | 6基准平均分 (avg) | 39.83 | SFT 32.67 / SeqKD 35.55 / RKLD 32.39 / JSD 36.61 | +8.8% vs JSD，+21.9% vs SFT |
| 数学推理 (LLaMA-3 8B→1B) | 6基准平均分 (avg) | 18.01 | SFT 13.88 / RKLD 15.07 / JSD 15.62 | +15.3% vs JSD，+29.8% vs SFT |
| 数学推理 (LLaMA-3 8B→3B) | 6基准平均分 (avg) | 34.56 | SFT 29.45 / RKLD 30.41 / JSD 29.95 | +13.6% vs RKLD，+17.4% vs SFT |
| 对话个性化 (Qwen 2.5 7B→1.5B, Ultrafeedback) | AE-LC / AE-WR / Arena-WR | 13.75 / 14.25 / 21.80 | SFT 12.74/13.72/18.10；RKLD 11.26/12.00/17.80；JSD 13.48/13.89/20.20 | 三项指标均最高 |
| 代码生成 pass@1 (DS-Coder 6.7B→1.3B) | HumanEval+ & MBPP+ pass@1 平均 | 66.35 | KD 64.60 / RKLD 61.60 / JSD 64.10 | +2.7% vs JSD |
| 代码生成 pass@1 (Qwen-Coder 7B→1.5B) | HumanEval+ & MBPP+ pass@1 平均 | 77.35 | KD 72.45 / RKLD 75.85 / JSD 76.00 | +1.8% vs JSD |
| On-policy 推理 (Qwen 2.5 7B→1.5B) | 6基准平均分 | HPD 单独 30.24；HPD+OPD 33.41 | SFT+OPD 29.56 | HPD 单阶段已胜 SFT+OPD；两阶段再 +13.0% |
| On-policy 推理 (LLaMA-3 8B→1B) | 6基准平均分 | HPD 18.01；HPD+OPD 20.40 | SFT+OPD 17.06 | 单阶段 +5.6%，两阶段 +19.6% |
| DPO 后偏好对齐 (Qwen 2.5 3B-Base) | AE-LC + AE-WR + Arena-WR 综合 Δ | HPD+DPO 综合 +11.92 | SFT+DPO +1.81；RKLD+DPO +9.74 | 比次优 RKLD 初始化多 +2.18 综合点 |
局限与改进
作者承认的局限性主要集中在三方面：(i) HPD 仍依赖 teacher 能产出与训练数据分布高度一致的输出，作者在附录和 Section 5.1 中假设 teacher distribution ≈ empirical dataset distribution 才能避免反向 forward 教师 forward 引入的误差；(ii) negative K1 估计器是 RKLD 的无偏估计但方差高，作者提到在高师生差距下仍有少数训练样本出现局部不稳定，这点和 reverse KL 本身的局限一致；(iii) 论文只在 1B–7B 量级、单一任务族（reasoning、dialogue、code）上验证，更大或更小模型、长上下文多模态尚未涉及。本人观察到的额外限制包括：(a) HPD 的「双倍 weight」仅在 $k_1>0$ 且 $k'_1<0$ 时触发，这个二元决策可能对初始化很敏感，作者虽称「没有引入超参」，但 $p(a^*_t|s_t)$ 这一项在实际数据集里有偏时会放大偏差；(b) 实验统一使用 batch size 256、训练 2k steps，对超长时间训练的稳定性还需要验证；(c) HPD 没有和同样号称轻量的 DistiLLM-2 contrastive loss 直接对比，仅靠论文自报数据无法判断 HPD 在 contrastive 任务上的相对位置。
独立分析的弱点
独立分析可以看到四点可改进空间：(1) HPD 的「专家 token 权重翻倍」是 hard-threshold 的二元判定（$k_1>0\&k'_1<0$），在 $k_1$ 接近 0 的边界 case 会闪烁切换方向，建议改成 soft-gate（例如 $w^*_t = (p+k_1) + \sigma(k_1)\cdot\sigma(-k'_1)\cdot p$ 类的连续形式）以平滑训练；(2) 负 K1 的高方差没做方差缩减（variance reduction），可以像 Schulman k3、leave-one-out 或 RLOO 那样做 baseline 减除，提升稳定性；(3) HPD 当前只 sample 一个 next token，对长 CoT 推理中下游 token 影响没建模，远端 credit assignment 缺乏，可考虑扩展到 multi-step student sampling + 截断重要性采样；(4) 论文没有面对极端小的模型（<1B）和极端大的师生 gap（>10×参数量）做鲁棒性测试，这种场景下「双倍 weight」机制可能反而放大教师噪声，建议增加 group-wise 或 capacity-aware 的权重衰减。
未来方向
作者在结论中明确提到：「Future work can explore applying HPD earlier in training, such as during mid-training or even pre-training, to assess its potential benefits for LLMs」，这是最直接的延伸方向。基于此工作可延伸的方向至少有四：(a) 把 HPD 从后训练阶段迁移到 mid-training / pre-training，研究其对 scaling law 的影响（论文的核心机理是 token 级 reweighting，理论上应该和数据规模同步受益）；(b) 与 RLHF / DPO 进一步耦合：Section 7 已展示 HPD 作为 DPO 初始化的优越性，下一步可以做 HPD-DPO joint training，让对齐阶段持续用 hybrid KL signal 监控 reference drift；(c) 引入 multi-token 或 block-level 的 K1 estimator，将 HPD 从 next-token 粒度扩展到 n-gram 粒度，更适合代码和长链推理的局部一致性；(d) 把 HPD 应用到 reasoning 模型（如 R1 系列）与 base 模型的 distillation，论文已经展示了 +77.9% 的 LLaMA-3 3B 增益，但只在 7B/8B 教师上，扩展到 70B 量级验证 scaling 行为是直接可做的。
复现评估
复现评估整体良好：(1) 代码已开源在 https://github.com/zwhong714/Hybrid-Policy-Distillation；(2) 使用公开数据集（OpenR1-Math-8192、Ultrafeedback、HumanEval、MBPP、OlympiadBench、MATH-500、GPQA、AIME24/25、AMC）均为标准 benchmark；(3) 模型覆盖 Qwen2.5 (1.5B/3B/7B) 和 LLaMA-3 (1B/3B/8B) 的开源权重，对计算资源要求是单卡 A100/H100 即可跑 1.5B/3B 蒸馏（batch size 256，总训练约 2k 步），7B 教师则需要更多显存，论文用了 teacher 已经在 OpenR1 数据上 GRPO 微调过的版本作 teacher，这一部分未必所有团队能直接复现；(4) 评估使用 EvalPlus、HuggingFace Open-R1 标准流程，统计显著性用 t-test 标注。复现难度估计中等：核心算法仅几十行改动，但 teacher 的 GRPO 微调、6+ 个 benchmark 的统一 eval、t-test 统计检验会增加工作量。
论文图表
