Divergence-Routed On-Policy Distillation for Language Models
On-policy distillation (OPD) trains a student on its own trajectories with reverse-KL, but reverse-KL gives weak gradients for teacher-supported tokens where the student assigns very low probability, causing teacher-mode under-coverage. This paper proposes Divergence-Routed OPD (DrOPD), which uses cross-sampling to show that teacher-student Jensen-Shannon divergence pinpoints behaviorally critical positions better than single-model entropy, then adds forward-KL selectively at the highest-JS positions while keeping reverse-KL everywhere. Across six math benchmarks and two student-teacher pairs, DrOPD reduces severe under-coverage by 37.9% and improves mean Avg@8 by up to 1.84 points, offering a principled token-level strategy for allocating distillation supervision.
DrOPD is a cleanly motivated and methodically executed paper that correctly diagnoses a gradient failure in standard RKL-based OPD and proposes a simple, theoretically grounded fix. Its strongest contributions are the cross-sampling diagnostic and the fixed-context under-coverage analysis, which together make a compelling case for teacher-student disagreement as a routing signal. That said, the downstream benchmark improvements are modest, the experiments are confined to math reasoning with closely sized Qwen3 models, and several relevant contemporaneous methods are neither benchmarked nor released. Overall, the paper offers solid mechanistic insights but limited practical headroom.
The cross-sampling diagnostic is rigorous. Divergence-based routing reaches near-teacher accuracy with fewer than 10% of tokens intervened, while entropy-based routing lags far behind, directly supporting the claim that single-model uncertainty is a poor proxy for teacher-student disagreement. The gradient analysis in Section 2.3 is equally sound: deriving that the RKL logit gradient vanishes as $p(v) \to 0$ while FKL retains a $-q(v)$ signal gives a clean theoretical justification for adding forward-KL supervision at high-disagreement positions.
The empirical gains are small and inconsistent across scales. For the 1.7B-4B pair, mean Avg@8 rises by only +1.84 percentage points, and for the 4B-8B pair the gain shrinks to +0.59, with Minerva actually declining. The cross-sampling evidence that underpins the entire routing argument comes from just 40 AMC23 problems with eight responses each—a sample of 320 trajectories that is far too small to generalize. Moreover, DrOPD is not free: training costs roughly 8% more per step and 7.7% more cumulative time than standard OPD, a trade-off the authors do not discuss.
The scope is narrow. Experiments cover only mathematical reasoning using Qwen3 models with modest size gaps, and several recent methods such as TIDE and TRACE are mentioned only in related work rather than evaluated. The reliance on a top-16 union candidate set may truncate low-probability but important teacher tokens, and baselines like EOPD are reproduced in-house, raising the risk of tuning bias.
The fixed-context under-coverage analysis is the strongest empirical link between the proposed mechanism and outcome: on held prefixes the severe-gap fraction falls from 4.48% under OPD to 2.79% under DrOPD. However, the paper never establishes that these recovered tokens causally drive the benchmark improvements rather than merely correlating with them. Comparisons to related work are descriptive rather than empirical; TIDE, TRACE, and AOPD appear only in the related-work section without shared-benchmark numbers, and EOPD is reproduced in-house rather than run from official code. A fetched excerpt from the original EOPD paper confirms the high-entropy forward-KL strategy but does not specify the 0.8 threshold used here, leaving ambiguity about reproduction fidelity.
The authors provide detailed hyperparameters, data, framework, and model specifications in Appendix A, which is sufficient for reproduction in principle. However, no code or checkpoints are released, and several baselines—most notably EOPD—are reproduced by the authors rather than taken from official implementations, introducing possible tuning bias. The teacher models are GRPO-trained variants that must be reproduced independently, and the training was run on 8*A800 GPUs, which limits accessibility. The method itself is simple enough to reimplement, but an exact independent reproduction would depend on precise access to the same GRPO-trained teacher weights and the verl framework version used.
On-policy distillation (OPD) transfers knowledge from a teacher model to a student along student-generated trajectories, typically by minimizing the reverse Kullback-Leibler (RKL) divergence. However, RKL can provide weak corrective signals for teacher-supported tokens to which the student assigns very low probability. This raises two questions: which positions are most consequential for the behavioral gap between the student and teacher, and what additional supervision is needed at these positions? Through cross-sampling, we find that divergence-based routing more effectively pinpoints positions contributing to the teacher-student performance gap than entropy-based routing. Further analysis reveals teacher-mode under-coverage at high Jensen-Shannon (JS) divergence positions, where standard OPD increases the large teacher-student probability gaps. Motivated by these findings, we propose Divergence-Routed On-Policy Distillation (DrOPD). DrOPD retains RKL supervision at all positions and selectively adds forward Kullback-Leibler (FKL) divergence supervision at positions with high JS divergence. Experiments across six math benchmarks show that DrOPD improves over standard OPD and entropy-based routing. DrOPD also reduces severe under-coverage cases by 37.9 percent relative to OPD, supporting targeted FKL as an effective complement to OPD.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.