BranchOPD：用分支采样让在线策略蒸馏更高效

从单轨迹到分支采样：OPD 的瓶颈在哪

在线策略蒸馏（On-Policy Distillation，OPD）是让小模型从大模型学习的主流路径之一。它的核心做法是：让学生模型自己生成轨迹，再由教师模型在这些轨迹上提供监督信号。相比传统的离线蒸馏，这种方式让学生在自己真实会犯错的分布上接受指导，监督更贴合实际推理路径。

但这套机制存在一个结构性限制。一条轨迹在每个决策点只能实现一种延续（continuation），这意味着在有限的采样预算下，监督的覆盖面被严重压缩。模型在那些关键的、不确定的岔路口到底该怎么走，往往得不到足够密集的教师指导。换句话说，单轨迹采样把大量监督机会浪费在了那些模型本就能走对的部分上。

在线策略蒸馏与离线蒸馏的根本区别在于数据分布的对齐方式。离线蒸馏使用教师模型预先生成的静态数据集训练学生，学生接触到的是教师的行为分布，而非自身真实推理时产生的分布。这导致"协变量偏移"（covariate shift）问题：学生在训练时从未见过自己容易犯错的状态，推理时一旦偏离教师轨迹，就缺乏纠错能力。在线策略蒸馏通过让学生主动生成轨迹来规避这一问题，但随之而来的代价是：每条轨迹是顺序生成的单一路径，在计算预算固定时，整个状态空间的覆盖密度极为稀疏，而模型最容易出错的"岔路口"往往又是低概率路径，被采样到的频率天然偏低。

BranchOPD 的核心思路

这篇研究提出的 BranchOPD（Branching On-Policy Distillation）给出的方案是：不要在整条轨迹上平均用力，而是找出最值得补充监督的决策点，在那里展开分支。

具体来说，方法包含三个关键步骤：

• 信息量决策点筛选：利用学生模型自身的不确定性（uncertainty）以及教师与学生之间的分歧（teacher-student disagreement）来定位哪些决策点最值得分支。这两个信号的组合，本质上是在找「学生没把握、而且和老师想法不一致」的地方。
• 共享前缀重采样：从同一个前缀出发，重新采样学生策略的多条延续，而不是从头生成多条独立轨迹。
• token 级教师监督：对主轨迹和分支都施加 token 级别的教师监督。

在不同任务中，分支操作的粒度也不同——数学任务里是在 token 级决策上分支，智能体（agent）任务里则是在回合内（within-turn）决策上分支。这种任务自适应的粒度设计，让方法能同时覆盖推理和交互两类场景。

这里的"不确定性"与"师生分歧"是两个独立但互补的信号。不确定性通常通过学生模型在该决策点的输出概率分布熵（entropy）来估计：熵越高，说明学生自身对下一步该怎么走越没把握。师生分歧则度量的是教师与学生在同一状态下给出的动作或 token 分布之间的差距，常用 KL 散度衡量。单独使用不确定性可能把预算浪费在学生"乱猜"但实际无关紧要的地方；单独使用师生分歧则可能优先处理学生确信但与教师相悖的点，这类点通常是学生存在系统性偏差的位置。两者结合，才能精准定位"学生不确定、且与教师分歧大"的高价值决策点，使分支采样的信息增益最大化。

共享前缀：省下近三成生成成本

BranchOPD 一个很务实的优点是成本控制。传统的 Multi-Rollout 做法需要生成多条完全独立的轨迹，重复计算了大量相同的前缀部分。而 BranchOPD 通过共享前缀复用（shared-prefix reuse），只在分支点之后才产生新的生成开销。

论文给出的数据相当直接：相比七条独立轨迹的方案，共享前缀复用在 ALFWorld 上减少约 19% 的学生生成 token 成本，在 WebShop 上减少约 31%。对于动辄需要大规模采样的蒸馏训练来说，这种节省意味着在同等预算下能获得更广的覆盖，或在同等覆盖下显著降低算力消耗。

协议边界修复 PBR

研究还引入了一项辅助机制——协议边界修复（Protocol-Boundary Repair，PBR）。它在协议约束下施加有针对性的监督，用来缓解重复生成和无效生成的问题。在智能体任务中，模型容易陷入重复动作或产生被浪费的 token，PBR 相当于给这些边界情况打了一个补丁，让生成更干净、更符合任务协议。

在智能体任务（如 ALFWorld、WebShop）中，模型的输出需要遵循特定的交互协议格式——例如动作必须属于合法动作集，或响应需要符合特定的结构化模板。分支采样在从中间状态重新生成延续时，新生成的序列可能违反这些协议约束，产生格式错误或循环重复的无效动作，这类 token 不仅不能提供有效的训练信号，还会污染损失计算。PBR 的作用是识别这类边界违规，并在其上施加针对性的监督修正，相当于在生成层面设置了一道过滤与修复机制，确保分支轨迹的可用性和训练信号的质量。

实验结果：全面超越两种基线

BranchOPD 的对比对象是 Vanilla OPD 和 Multi-Rollout OPD（n=7），结果显示它在所有评测基准上都胜出：

• 数学任务：在 avg@16（对 16 个采样答案求平均准确率）指标上，比 Vanilla OPD 提升 18.55 个百分点。
• ALFWorld：在 ID/OOD（分布内/分布外）平均成功率上，比 Vanilla OPD 提升 13.8 个百分点。

值得关注的是，这些收益不仅相对最朴素的基线成立，也超过了消耗更大的 Multi-Rollout 方案。也就是说，BranchOPD 既更省，又更强——这在训练方法的改进里并不常见。

为什么在不确定处分支有效

论文的机制分析（mechanistic analyses）和筛选消融实验（selection ablations）回答了一个关键问题：分支点的选择是否真的重要？

结论是肯定的。在不确定的决策点上分支，能够扩展后续行为的覆盖面（expand subsequent behavior coverage），其表现优于两种对照策略：随机选择分支点，以及固定在每个回合开头分支（turn-initial branching）。这说明方法的收益并非来自单纯增加采样量，而是来自「把采样预算花在刀刃上」这一精准定位能力。

从方法论角度看，这给在线策略蒸馏指出了一个值得深挖的方向：监督的价值分布是高度不均匀的，与其均匀铺开，不如用模型自身的信号来动态分配监督资源。这种「以不确定性为导向」的采样思路，对强化学习和蒸馏领域都有迁移价值。

小结

BranchOPD 的贡献可以概括为：用学生不确定性和师生分歧定位关键决策点，通过共享前缀分支采样扩大监督覆盖，同时显著降低生成成本，并辅以 PBR 处理协议边界问题。它在数学和智能体任务上同时取得了明显增益，为有限采样预算下的高效蒸馏提供了一条清晰的技术路线。