WDL-OPD: Weak-Driven On-Policy Distillation via Mixture-Constrained Co-Training

2026-08-10Machine Learning

Machine LearningArtificial Intelligence
AI summary

The authors study on-policy distillation (OPD), a way to make a smaller model (student) learn from a bigger one (teacher) using data generated by the student itself, which helps reduce mismatches. They propose a new method called WDL-OPD that uses two policies working together to stabilize learning by mixing their predictions and comparing them to the teacher. Experiments on large language models show that WDL-OPD improves accuracy on math problems and code generation compared to older methods. The authors caution that while their method helps stabilize learning, more testing is needed to confirm how broadly it works. They share detailed results and the exact training approach to support further study.

On-policy distillationStudent-teacher learningReinforcement learningPolicy co-trainingReverse KL divergenceModel stabilizationLarge language modelsQwen3 modelCode generationMATH500 benchmark
Authors
Zehao Chen, Gongxun Li, Tianxiang Ai, Yifei Li, Zixuan Huang, Wang Zhou, Tao Huang, Fuzhen Zhuang, Xianglong Liu, Jianxin Li, Deqing Wang, Yikun Ban
Abstract
On-policy distillation (OPD) aligns a student with a teacher on trajectories sampled from the student itself, reducing the train-test state mismatch of offline distillation. The same feedback loop can nevertheless be unstable: each update changes both the policy and the states on which the next update is computed. We introduce WDL-OPD, a mixture-constrained co-training method with two trainable policies. An anchor policy generates every rollout, an auxiliary policy evaluates the same visited states, and a geometric mixture of their token distributions is matched to a frozen teacher by reverse KL. Both policies receive gradient. We show that freezing the auxiliary recovers an anchor-plus-contrast proxy target closely related to OPD$^2$ and W2S-OPD, whereas joint training creates branch-level degrees of freedom that a static delta cannot express. In recorded Qwen3 experiments at 1.7B and 4B scale, WDL-OPD produces the strongest student checkpoint in each of four scale-domain settings. It raises MATH500 accuracy from 0.630 to 0.685 at 4B and from 0.521 to 0.585 at 1.7B. In code generation, seven single-policy OPD configurations exhibit entropy growth or trajectory degradation, while co-training reaches independently re-evaluated development scores of 0.637 and 0.375. Because several comparisons differ in curriculum or initialization, these results support a stabilization hypothesis rather than a universal causal claim. We provide the exact training algorithm, failure evidence, and the controlled comparison matrix needed to test that hypothesis.