通过筛选轨迹再加权,实现更精细的强化学习蒸馏优化
Filter, Then Reweight: Rethinking Optimization Granularity in On-Policy Distillation

- 先过滤低质量轨迹,再对保留轨迹中的关键token进行软加权
- 在强-弱教师设置下比现有方法提升6.25分,多教师设置提升18.81分
- 适合追求高精度蒸馏的LLM研究者和工业部署场景
大型语言模型中的在线策略蒸馏(OPD)正从全轨迹KL监督转向更选择性的训练范式。近期方法聚焦于选择哪些轨迹、哪些标记更具信息量、哪些监督信号更可靠。受此趋势启发,我们重新思考OPD的优化粒度,提出FiRe-OPD(Filter, then Reweight),在轨迹和标记两个层面联合调整监督信号。具体而言,FiRe-OPD首先过滤掉低质量的采样轨迹,然后在保留轨迹内对信息量高的标记进行软加权。相比硬性标记选择,该方法通过软加权机制有效缓解信息丢失,提升优化稳定性,从而实现更细粒度的OPD优化。我们在强-弱教师、单教师及多教师设置下验证了其有效性,结果表明其优于近期的标记级OPD方法:在强-弱教师设置中,AIME 2024得分提升6.25;在多教师设置中,Miner得分提升18.81。代码已公开于https://github.com/YuYingLi0/FiRe-OPD。
原文摘要 · Abstract (English)
On-Policy distillation (OPD) in large language models is shifting from full-trace KL supervision toward more selective training paradigms. Recent OPD methods increasingly focus on selecting which trajectories to learn from, which tokens are most informative, and which supervision signals are most reliable. Motivated by this trend, we rethink optimization granularity of OPD and propose \fireicon\ FiRe-OPD (Filter, then Reweight), which jointly adjusts supervision signals at both trajectory and token levels. In details, FiRe-OPD first filters trajectories to remove low-quality rollout samples, and then applies soft reweighting within the retained trajectories to emphasize informative tokens. Compared with hard token selection, FiRe-OPD leverages a soft-weighting mechanism to effectively mitigate information loss and enhance optimization stability, thereby achieving finer-grained OPD optimization. We validate the effectiveness of FiRe-OPD across strong-to-weak, single-teacher, and multi-teacher settings, and demonstrate its superiority over recent token-level OPD methods ( (e.g., +6.25 on AIME 2024 in strong-to-weak, +18.81 on Miner in multi-teacher). Our code is available at https://github.com/YuYingLi0/FiRe-OPD.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。