Source-linked AI summary
Filter, Then Reweight: Rethinking Optimization Granularity in On-Policy Distillation
Yuying Li, Leqi Zheng, Yongzi Yu, Wenrui Zhou, Xuchang Zhong, Xing Hu, Jing Jin, Hangjie Yuan, Tao Feng
TL;DR
On-policy distillation must handle unequal learning value and reliability across trajectories and tokens. FiRe-OPD filters low-confidence trajectories and softly reweights tokens, consistently improving over standard OPD and recent baselines across three distillation scenarios.
Problem
Standard OPD uniformly supervises full trajectories despite unequal learning value and signal reliability across trajectories and tokens.
Method
FiRe-OPD filters low-confidence trajectories and continuously weights tokens using teacher confidence and student confusion.
Results
FiRe-OPD consistently improves over standard OPD and recent baselines across three distillation scenarios, achieving +6.25 on AIME 2024 in strong-to-weak distillation.
Takeaways & Limitations
Hard filtering is more effective for trajectories, whereas soft weighting surpasses hard token selection at the token level.
Takeaways & Limitations
FiRe-OPD treats tokens independently and does not model how erroneous prefixes may degrade subsequent teacher signals.
Abstract
from arXiv · showhide
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.
1 Introduction
FiRe-OPD addresses OPD’s granularity and reliability limitations by filtering trajectories and softly reweighting tokens. The framework argues that hard trajectory filtering and soft token weighting are preferable at their respective optimization levels.
- Motivation: OPD learns from student-generated trajectories with dense token-level supervision, avoiding supervised fine-tuning’s train-inference mismatch and reinforcement learning’s sparse outcome rewards.
- Motivation: Standard OPD applies uniform full-trajectory KL supervision even though trajectories and tokens differ in learning value and supervision reliability.
- Limitations: Existing methods isolate trajectory- or token-level optimization, while hard token selection can cause nonsmooth optimization and discard useful supervision signals.
- FiRe-OPD: FiRe-OPD jointly filters low-likelihood trajectories and applies continuous token importance weights using teacher confidence and student confusion.
- Contributions: FiRe-OPD establishes that hard filtering is more effective for trajectories, whereas soft weighting surpasses hard selection for tokens across multiple distillation settings.
2 Related Work
Related work spans classical and sequence-level off-policy distillation, broader capability transfer in LLMs, and recent efforts to reduce noisy token supervision. FiRe-OPD is characterized by combining trajectory-level filtering with token-level importance weighting, addressing limitations of hard token selection.
- Off-Policy Distillation: Off-policy distillation transfers teacher knowledge through output distributions or complete teacher-generated responses, while LLM distillation increasingly targets reasoning and alignment.Classical knowledge distillation matches teacher output distributions, whereas sequence-level distillation uses complete teacher-generated responses.
- FiRe-OPD: FiRe-OPD combines trajectory-level filtering with token-level importance weighting.
- Token Selection: Recent token-selection methods discard low-value tokens to reduce noisy supervision, but hard selection can lose useful information and yield brittle optimization signals.
3 Methodology
FiRe-OPD refines standard on-policy distillation by filtering low-quality trajectories before applying soft token-level reweighting. It prioritizes paths with aligned teacher supervision and token positions where the teacher is confident while the student remains confused, preserving gradient contributions across positions.
- Standard OPD: Standard OPD uses PPO-style clipped objectives with token-level teacher–student log-likelihood-ratio advantages and treats all trajectories and token positions uniformly.The clipping parameter is set to ϵ = 0.2.
- FiRe-OPD overview: FiRe-OPD combines trajectory-level filtering with token-level soft reweighting because distillation-signal quality varies across both levels.This replaces uniform supervision with complementary granularity-aware mechanisms.
- Trajectory-level filtering: The trajectory importance score uses the teacher’s normalized log-probability on each student-generated rollout to measure teacher–student distributional alignment.Low teacher log-probability indicates a large distribution gap and unreliable token-level guidance along that path.
- Trajectory-level filtering: FiRe-OPD ranks rollouts by s(y), discards the bottom p% with p = 20 by default, and optimizes only the surviving trajectories.The retained paths are intended to lie within the teacher’s competence region, where supervision is more coherent and meaningful.
- Token-level soft reweighting: Token importance combines teacher confidence and student confusion multiplicatively, emphasizing positions where guidance is reliable and genuinely needed.Unlike hard truncation, soft weighting retains gradient contributions from all positions in proportion to relative informativeness.
4 Experiment
FiRe-OPD is evaluated across strong-to-weak, single-teacher, and multi-teacher distillation settings, using mathematical and code-generation benchmarks against standard OPD and recent improvements. It achieves the strongest reported results across scenarios, while ablations and strategy comparisons support hard trajectory filtering with soft token weighting.
- Evaluation Setup: The evaluation spans strong-to-weak, single-teacher, and multi-teacher scenarios, with eight mathematical benchmarks reporting Avg@8 and three code benchmarks reporting pass@1.Baselines include standard OPD, ExOPD, TIP, REOPOLD, EOPD, and Uni-OPD, trained under the same data, model, and compute budget.
- Strong-to-Weak Distillation: FiRe-OPD achieves 60.83% average accuracy in strong-to-weak distillation, surpassing ExOPD by 0.67 points and standard OPD by 2.13 points.It gains +6.25 on AIME 2024, +4.17 on AIME 2025, +3.75 on HMMT Feb, and +2.48 on MATH-500.
- Single-Teacher Distillation: FiRe-OPD reaches 61.74% average accuracy in single-teacher distillation, improving over standard OPD by 0.53 points and ExOPD by 0.90 points.Notable gains include +3.33 on AIME 2024 and +2.08 on HMMT Nov.
- Multi-Teacher Distillation: FiRe-OPD achieves 51.88% on multi-teacher math reasoning and 64.16% on code generation, improving over OPD by +4.84 and +4.37, respectively.It gains +18.81 on MinervaMAT, +4.59 on AIME 2024, +9.77 on HumanEval+, and +2.11 on MBPP+.
- Component Ablation: The full FiRe-OPD score is 60.83, while removing student confusion drops performance by -2.24, trajectory filtering by -1.84, and teacher confidence by -0.96.Student confusion is the dominant token-level signal, while teacher confidence is complementary.
- Soft Weighting vs. Hard Truncation: FiRe-OPD’s Hard trajectory filtering + Soft token weighting achieves 60.83, outperforming Hard+Hard at 58.23, Soft+Soft at 58.68, and Soft+Hard at 58.55.Hard filtering removes unreliable trajectories, whereas soft weighting preserves gradient contributions according to token informativeness.
5 Conclusion
FiRe-OPD is a dual-granularity on-policy distillation framework that filters low-confidence trajectories and continuously reweights tokens using teacher confidence and student confusion. Across three distillation scenarios, it consistently improves over standard OPD and recent baselines, while ablations show asymmetric contributions from teacher and student signals across granularities.
- Framework: FiRe-OPD combines trajectory filtering with continuous token-level weighting for on-policy distillation.It filters low-confidence trajectories and assigns token weights based on teacher confidence and student confusion.
- Evaluation: The framework was evaluated across three distillation scenarios on math reasoning and code generation benchmarks.
- Results: FiRe-OPD consistently improves over standard OPD and recent baselines.
- Ablations: Ablation studies show that teacher and student signals contribute asymmetrically across optimization granularities.
6 Limitations
FiRe-OPD’s adaptive distillation granularity remains an open design space, particularly for modeling erroneous prefixes and exploring higher-level weighting units aligned with chain-of-thought structure.
- 6 Limitations: The design space for adaptive distillation granularity remains largely unexplored.
- 6 Limitations: FiRe-OPD treats tokens independently and does not model how erroneous prefixes may degrade subsequent teacher signals.
- 6 Limitations: Step-level and segment-level weighting are promising future directions because they align more naturally with chain-of-thought structure.
A Full Ablation Results
The ablations identify robust entropy-aware weighting and trajectory filtering choices in strong-to-weak distillation. Performance peaks with moderate teacher-confidence scaling and a 20% trajectory-filtering percentile, while extremes can hurt training.
- Sensitivity to α and β: 60.83% average accuracy is achieved at α = 1.0 with β fixed at 1.0, while α = 5.0 remains competitive at 60.64%.The results indicate that moderately amplifying teacher-confidence signals is beneficial and that performance is not overly sensitive to α.
- Sensitivity to Trajectory Filtering Percentile: p = 20% achieves 60.83% average accuracy for trajectory-level filtering and is the default setting.Lower filtering retains noisy off-distribution trajectories, whereas aggressive filtering discards potentially useful training signals.