Zhang Honghai

Author directory

2026

On-policy distillation reduces train-inference mismatch by supervising student-generated trajectories, but full-vocabulary token-level distribution matching remains expensive for large language models. A common approximation truncates reverse-KL distillation to the student’s top-k tokens, which preserves efficiency but drops tail-token supervision and ignores support-level mass mismatch between the student and teacher. We analyze these approximation gaps and propose Stochastic Support Sampling with Teacher Mass Reweighting, which samples fixed-size supports from the student distribution and uses teacher support mass as a reliability weight for the conditional distillation loss. Under the same support budget, our method outperforms the strongest baseline by +1.4 hard-math average points on Qwen3-8B→Qwen3-1.7B-Base, and by +1.1 math average points and +2.8 code average points on DeepSeek-R1-Distill-Llama-8B→Meta-Llama-3.1-8B-Base.