New method improves AI learning by better matching teacher and student choices
Distillation as Probability Transport: Routed On-Policy Distillation
Machine LearningComputation and Language
Summary
Sometimes, AI teaches another AI by showing examples of what to do, but the typical way only looks at simple yes/no feedback for each step without explaining how to fix mistakes overall. The authors suggest a way to think of teaching as moving chunks of probability from one choice to another, which means the student AI gets clearer guidance on where to improve. They create a method called RouteOPD that matches parts of the student's decisions needing change with specific parts the teacher prefers, making the learning more precise. Tests show RouteOPD helps the student AI learn better and follow the teacher’s guidance more closely than older methods.
on-policy distillationprobability transportteacher-student learninglog-odds optimizationreinforcement learningdistribution matchingreverse-KL divergencemathematical reasoning benchmarks
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.