arXiv:2605.09725cs.CV2026-05被引 3

通过优选最优教师轨迹,提升强化学习推理的训练效率与准确性。

On-Policy Distillation with Best-of-N Teacher Rollout Selection

论文配图:On-Policy Distillation with Best-of-N Teacher Rollout Selection
图 1 · 摘自论文原文
  • 从多个教师轨迹中选取最正确且与学生行为匹配的进行监督
  • 在复杂题目上相比标准方法提升显著,最高提升达12.3%
  • 适合需要高效推理训练的AI模型研究者使用

在线策略蒸馏(OPD)作为一种数据高效的后训练方法,可在避免强化学习奖励依赖和标准监督微调中的灾难性遗忘问题的同时提升推理能力。然而,标准OPD通常在噪声干扰的学生采样上下文中计算教师监督,且仅依赖单个随机教师轨迹,导致监督信号方差大:采样的教师轨迹可能错误、无信息或与学生当前推理行为不匹配。为此,我们提出BRTS框架,通过从采样教师轨迹池中选择最优轨迹来增强标准学生上下文蒸馏。BRTS引入教师上下文监督分支,采用优先规则筛选:优先选择正确的轨迹,其次选择与学生行为最契合的。当多个正确轨迹存在时,选对齐度最高的;若无条件教师样本在难题上失败,则触发基于真值条件的恢复步骤以生成自然推导过程。所选轨迹用于提供可靠教师上下文监督,并辅以额外损失项。在AIME 2024、AIME 2025和AMC 2023上的实验表明,BRTS在困难推理基准上优于标准OPD,尤其在高难度数据集上提升最大。代码已开源。

原文摘要 · Abstract (English)

On-policy distillation (OPD), which supervises a student on its own sampled trajectories, has emerged as a data-efficient post-training method for improving reasoning while avoiding the reward dependence of reinforcement learning and the catastrophic forgetting often observed in standard supervised fine-tuning. However, standard OPD typically computes teacher supervision under noisy student-generated contexts and often relies on a single stochastic teacher rollout per prompt. As a result, the supervision signal can be high-variance: the sampled teacher trajectory can be incorrect, uninformative, or poorly matched to the student's current reasoning behavior. To address this limitation, we propose BRTS, a Best-of-N Rollout Teacher Selection framework for on-policy distillation. BRTS augments standard student-context OPD with a teacher-context supervision branch constructed from the curated teacher trajectory. Rather than distilling from the first sampled teacher rollout, BRTS samples a small pool of teacher trajectories and selects the auxiliary trajectory using a simple priority rule: correctness first, student alignment second. When multiple correct teacher trajectories are available, BRTS chooses the one most aligned with the student's current behavior; when unconditioned teacher samples fail on harder prompts, it invokes a ground-truth-conditioned recovery step to elicit a natural derivation. The selected trajectory is then used to provide reliable teacher-context supervision inside the OPD loop, augmented with an auxiliary loss on the teacher trajectory. Experiments on AIME 2024, AIME 2025, and AMC 2023 show that BRTS improves over standard OPD on challenging reasoning benchmarks, with the largest gains on harder datasets. Our code is available at https://github.com/BWGZK-keke/BRTS.

推理增强知识蒸馏强化学习AI训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。