Filtrer, puis repondérer : Repenser la granularité de l'optimisation dans la distillation on-policy

Résumé

La distillation on-policy (OPD) dans les grands modèles de langage évolue d'une supervision KL complète sur les trajectoires vers des paradigmes d'apprentissage plus sélectifs. Les méthodes récentes d'OPD se concentrent de plus en plus sur la sélection des trajectoires à apprendre, des tokens les plus informatifs et des signaux de supervision les plus fiables. Motivés par cette tendance, nous repensons la granularité d'optimisation de l'OPD et proposons \fireicon\ FiRe-OPD (Filter, then Reweight), qui ajuste conjointement les signaux de supervision aux niveaux des trajectoires et des tokens. En détail, FiRe-OPD filtre d'abord les trajectoires pour éliminer les échantillons de rollout de faible qualité, puis applique un rééquilibrage doux au sein des trajectoires retenues afin de mettre en avant les tokens informatifs. Comparé à une sélection dure de tokens, FiRe-OPD exploite un mécanisme de pondération douce pour atténuer efficacement la perte d'information et améliorer la stabilité de l'optimisation, permettant ainsi une optimisation OPD plus fine. Nous validons l'efficacité de FiRe-OPD dans des contextes de fort-à-faible, de professeur unique et de professeurs multiples, et démontrons sa supériorité par rapport aux méthodes OPD récentes au niveau des tokens (par exemple, +6,25 sur AIME 2024 dans le cadre fort-à-faible, +18,81 sur Miner dans le cadre multi-professeurs). Notre code est disponible à l'adresse 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.