Divergence methods improve knowledge transfer for language models
Divergence-Routed On-Policy Distillation for Language Models
Artificial Intelligence
Summary
Training smaller language models to learn from larger ones often uses a technique called on-policy distillation, which sometimes struggles to correct the small model’s mistakes strongly enough. The paper’s authors discovered that checking differences in how models assign probabilities (divergences) helps identify where errors matter most. They propose a new method that adds a different type of correction where the models disagree greatly, improving the smaller model’s learning and reducing important mistakes. Tests on math problems show this new approach works better than previous ones.
What this means in practice
- •For natural language processing engineers: Improve training methods for compact language models to better match larger models’ behavior on complex text generation tasks.
- •For automated math tutoring developers: Enhance model training to produce more accurate and reliable automated solutions for math problem-solving scenarios.
Authors
Tiantian Peng, Guannan Lv, Hongyang Li, Qiuhe Hong, Hang Cheng, Hongjian Dou, Xiangyu Wu, Yuyang Liu
Abstract
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.