필터링 후 재가중: 온-정책 증류에서 최적화 세분성 재고찰

초록

대규모 언어 모델에서의 온정책 증류(On-Policy Distillation, OPD)는 전체 궤적에 대한 KL 감독(full-trace KL supervision)에서 보다 선택적인 훈련 패러다임으로 전환되고 있다. 최근 OPD 방법들은 어떤 궤적(trajectory)을 학습할지, 어떤 토큰이 가장 유용한 정보를 제공하는지, 그리고 어떤 감독 신호가 가장 신뢰할 수 있는지를 선택하는 데 점점 더 초점을 맞추고 있다. 이러한 추세에 착안하여, 본 연구는 OPD의 최적화 세분성(granularity)을 재고하고 궤적 및 토큰 수준에서 감독 신호를 동시에 조정하는 \fireicon\ FiRe-OPD(Filter, then Reweight)를 제안한다. 구체적으로, FiRe-OPD는 먼저 궤적을 필터링하여 품질이 낮은 롤아웃 샘플을 제거한 후, 유지된 궤적 내에서 소프트 재가중(soft reweighting)을 적용하여 정보가 풍부한 토큰을 강조한다. 하드 토큰 선택과 비교하여, FiRe-OPD는 소프트 가중 메커니즘을 활용하여 정보 손실을 효과적으로 완화하고 최적화 안정성을 향상시킴으로써 더 세분화된 OPD 최적화를 달성한다. 우리는 강자에서 약자(strong-to-weak), 단일 교사(single-teacher), 다중 교사(multi-teacher) 설정에서 FiRe-OPD의 효용성을 검증하고, 최근의 토큰 수준 OPD 방법보다 우수함을 입증한다(예: 강자에서 약자 설정에서 AIME 2024에서 +6.25, 다중 교사 설정에서 Miner에서 +18.81). 코드는 https://github.com/YuYingLi0/FiRe-OPD에서 확인할 수 있다.

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.