基于 Best-of-N 教师轨迹选择的同策略蒸馏
计算机视觉与模式识别
2026-05-14 v2
摘要
同策略蒸馏(OPD)在学生模型自身的采样轨迹上对其进行监督,已成为一种数据高效的训练后方法,能够提升推理能力,同时避免强化学习的奖励依赖性以及标准监督微调中常见的灾难性遗忘。然而,标准 OPD 通常在含噪的学生生成上下文下计算教师监督,且往往依赖每个提示仅进行一次随机教师轨迹采样。因此,监督信号可能具有高方差:采样的教师轨迹可能是不正确的、无信息量的,或与学生当前的推理行为匹配度差。为解决这一局限性,我们提出了 BRTS,一种用于同策略蒸馏的 Best-of-N 轨迹教师选择框架。BRTS 利用从精选教师轨迹构建的教师上下文监督分支,增强了标准学生上下文 OPD。BRTS 并非从首次采样的教师轨迹进行蒸馏,而是采样一个小的教师轨迹池,并使用简单的优先级规则选择辅助轨迹:首先考虑正确性,其次考虑学生对齐度。当存在多个正确的教师轨迹时,BRTS 选择与学生当前行为最对齐的轨迹;当无条件的教师样本在较难的提示上失败时,它调用真值条件恢复步骤以引出自然的推导。所选轨迹随后用于在 OPD 循环内提供可靠的教师上下文监督,并在教师轨迹上增加辅助损失。在 AIME 2024、AIME 2025 和 AMC 2023 上的实验表明,BRTS 在具有挑战性的推理基准上优于标准 OPD,在较难的数据集上提升最大。我们的代码可在 https://github.com/BWGZK-keke/BRTS 获取。
关键词
引用
@article{arxiv.2605.09725,
title = {On-Policy Distillation with Best-of-N Teacher Rollout Selection},
author = {Ke Zhang and Yunjie Tian and Dongdi Zhao and Yijiang Li and Yuanye Liu and Vishal M Patel and Di Fu},
journal= {arXiv preprint arXiv:2605.09725},
year = {2026}
}
备注
10 pages, 5 figures