Probability transport improves on-policy knowledge transfer methods
Distillation as Probability Transport: Routed On-Policy Distillation
Machine LearningComputation and Language
Summary
It can be hard for one AI model (the student) to learn from another (the teacher) because the usual methods only tell the student whether to increase or decrease the chance of choosing certain words, without explaining where that probability should move. The authors propose a new way called RouteOPD that treats this learning process like moving probability from places where the student guesses too much to places where the teacher thinks more focus is needed. This approach guides the student more precisely on how to adjust its probabilities, leading to better learning and performance on difficult reasoning tasks.
What this means in practice
- •For machine learning engineers: Improve the accuracy of AI models learning from teacher models by explicitly directing probability updates between tokens.
- •For natural language processing teams: Enhance language model training on reasoning tasks by using probability transport to better align student models with teacher guidance.
Authors
Tianle Xia, Lingxiang Hu, Yiding Sun, Linfang Shang, Ming Xu, Lan Xu, Ning Zheng, Wei Xu, Jie Jiang
Abstract
On-policy distillation (OPD) transfers teacher knowledge on student-generated trajectories, but efficient sampled objectives reduce the teacher distribution to scalar credit on individual tokens. Such credit indicates whether a token should gain or lose probability, yet leaves the corresponding redistribution unspecified. We recast OPD as teacher-guided probability transport and propose RouteOPD (Routed On-Policy Distillation), which decomposes local teacher--student disagreement into student-excess sources and teacher-deficit destinations and couples them into explicit transport pairs. RouteOPD optimizes pairwise log-odds toward jointly realizable targets obtained from a bounded teacher potential, while adapting the transport budget to the concentration of teacher demand. This formulation directs updates toward teacher-preferred destinations and controls their magnitude within a single transport operator. Experiments across four teacher--student settings and four mathematical-reasoning benchmarks demonstrate that RouteOPD consistently outperforms sampled reverse-KL OPD, with improvements accompanied by higher routing fidelity and lower background leakage. These results demonstrate the effectiveness of explicitly modeling probability transport in on-policy distillation.