Length inflation in model distillation caused by token disagreement
When EOS Tokens Disagree: Understanding Length Inflation in On-Policy Distillation
Machine Learning
Summary
Sometimes when one AI model learns from another, its answers get too long and fill up the allowed space. The researchers found that this happens because the two models disagree on which special 'stop' word means to end a response. Even if both models say they use the same stop words, they place different importance on those words, causing confusion. The team shows that treating different stop words that mean the same thing as one shared stop signal helps reduce these overly long answers. However, there are still other reasons that cause length increase, especially late in training.
What this means in practice
- •For machine learning engineers: Reduce excessive response lengths in AI models by implementing shared stop token handling during knowledge distillation.
- •For software developers building chatbots: Improve chatbot response generation efficiency by correcting termination token mismatches between models after fine-tuning.
Authors
Yuxiao Yang, Tianrun Yu, Shangzhe Li, Kaixiang Zhao, Xuchao Zhang, Chetan Bansal, Huaxiu Yao, Taylor W. Killian, Weitong Zhang
Abstract
We study length inflation in on-policy distillation (OPD), where student responses can become excessively long and even exhaust the generation budget. We identify \emph{termination-token mismatch} between base students and post-trained teachers as an important source of this behavior. Across Qwen3, Llama, and Gemma, the two models can place their stopping probability on different EOS tokens, even when their declared stopping sets are identical. This mismatch can suppress the student's preferred termination action without reliably transferring the teacher-preferred alternative. We show that aligning the decoding stopping set alone is insufficient, while treating functionally equivalent EOS tokens as a shared semantic stopping action substantially mitigates mismatch-induced length inflation across all three model families. To further understand how termination behavior evolves over training, we study OPD across different K2-Horizon training stages. This stage-wise analysis shows that termination preferences can shift substantially during training, while also revealing a distinct length inflation late in the OPD run that persists beyond termination alignment. Together, these results identify termination mismatch as an important, but not exhaustive, source of OPD length dynamics. We release an implementation incorporating the proposed termination-handling corrections.