One-sentence Summary

Researchers at the University of North Carolina at Chapel Hill, Brigham Young University, and Microsoft show that termination-token mismatch between base students and post-trained teachers drives length inflation in on-policy distillation across Qwen3, Llama, and Gemma, and they propose treating functionally equivalent EOS tokens as a shared semantic stopping action, which substantially mitigates this mismatch-induced inflation, while stage-wise K2-Horizon analysis reveals additional late-run length dynamics.

Key Contributions

• The paper identifies termination-token mismatch between base students and post-trained teachers as a key source of length inflation in on-policy distillation, showing across Qwen3, Llama, and Gemma that models can prefer different EOS tokens even when their declared stopping sets are identical, which suppresses the student's preferred termination action without reliably transferring the teacher-preferred alternative.
• The paper demonstrates that aligning the decoding stopping set alone is insufficient, whereas treating functionally equivalent EOS tokens as a shared semantic stopping action substantially mitigates mismatch-induced length inflation across all three model families.
• The paper presents a stage-wise analysis of on-policy distillation across K2-Horizon training stages, showing that termination preferences can shift substantially during training and that a distinct late-training length inflation persists beyond termination alignment, and it releases an implementation incorporating the proposed termination-handling corrections.

Introduction

On-policy distillation (OPD) trains a student on its own sampled rollouts with dense token-level supervision from a post-trained teacher, making it a practical way to transfer strong behavior to smaller or less capable models. However, OPD often produces progressively longer responses with truncation or repetition, and prior work has attributed this length inflation to objective-level effects, rollout degradation, or training instability. The authors identify termination-token mismatch as a direct and diagnosable cause: base students and post-trained teachers may encode the same semantic stop decision with different EOS tokens, even when their declared stopping sets match. Because the teacher evaluates only student-generated tokens, the student's own termination action can be suppressed. The authors 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 Qwen3, Llama, and Gemma.

Method

The authors identify that modifying only the decoding interface is insufficient to address termination mismatch in distillation. The mismatch must be corrected within the distillation signal itself by reconciling the termination probabilities of the teacher and the student. Let $\mathcal{E}_{\mathrm{EOS}}$EEOS​ denote the termination-equivalent tokens under the rollout protocol, and let $e_{\star} \in \mathcal{E}_{\mathrm{EOS}}$e⋆​∈EEOS​ be a canonical EOS token supported by the base student. To address this, the authors propose three probability-level alignment corrections.

The first correction, teacher-side EOS mapping, maps the probability mass the teacher assigns to all termination-equivalent tokens to the canonical student EOS token. This is formulated as:

In practice, negligible probability is retained on other EOS tokens for numerical stability, while the canonical probability is adjusted to preserve the total mass. The student distribution remains unchanged, and rollout terminates exclusively on $e_{\star}$e⋆​.

The second correction, semantic EOS class, treats all tokens in $\mathcal{E}_{\mathrm{EOS}}$EEOS​ as realizations of a single semantic action, STOP, rather than selecting a canonical surface token. For any policy $\pi \in \{\pi_{\theta}, \pi^E\}$π∈{πθ​,πE}, the aggregated probability is defined as:

while $\bar{\pi}(a \mid s_t) = \pi(a \mid s_t)$πˉ(a∣st​)=π(a∣st​) for non-EOS tokens. If a sampled token $y_t$yt​ is an EOS token, it is replaced by $\bar{y}_t = \mathrm{STOP}$yˉ​t​=STOP. The standard sampled-token OPD update is then applied directly:

This ensures EOS samples are supervised through the total termination probability without altering non-EOS tokens.

The third correction, canonical single-EOS action space, maps the teacher EOS probability mass to $e_{\star}$e⋆​ and removes all other tokens in $\mathcal{E}_{\mathrm{EOS}}$EEOS​ from the student sampling distribution, renormalizing the remaining probabilities. This creates a consistent action space with a single canonical termination action.

The authors note that these three corrections yield similar empirical results. Teacher-side EOS mapping is utilized for explicit, known mismatches because it simply modifies the teacher distribution. However, semantic EOS aggregation is adopted as the default for cross-family experiments because it avoids the requirement of designating a canonical surface form, accommodating models that expose the same declared EOS set but prefer different tokens.

The shared termination failure is distinct from its effect on measured task performance. Cross-template evaluation can reduce observed response length, and performance gains from OPD can persist despite severe length inflation, particularly under specific evaluation templates. The apparent effect on accuracy also depends on the grader, as certain evaluation styles are more sensitive to long redundant continuations. Consequently, length inflation does not necessarily imply a comparable loss of reasoning capability.

As shown in the figure below:

Experiment

The experiments evaluate sampled-token on-policy distillation between base student models and post-trained or instruct teachers from Qwen3, Llama 3.2, Gemma 3, and K2-Horizon on single-turn math reasoning tasks. They show that length inflation and clipping arise when teacher and student distribute termination probability across different equivalent EOS tokens, even if the decoding stopping sets are aligned, and that probability-level corrections such as semantic EOS aggregation substantially mitigate this mismatch across model families. Stage-wise K2-Horizon runs indicate that vanilla distillation can transfer a teacher-preferred termination token when the student already gives it enough sampling support, but later re-inflation and other residual length dynamics remain even after correction. Template and grader comparisons further show that the severity of length inflation depends on evaluation context, while some downstream performance gains can still transfer.

The studied model families differ in their declared EOS token configurations across training stages. Qwen and Llama base checkpoints declare a single end-of-text token, while their post-trained checkpoints also recognize additional conversational termination tokens. Gemma and K2-Horizon show cases where declared stopping sets are shared across stages, so termination mismatches can persist through learned preferences rather than decoding configuration alone. Post-trained Qwen and Llama checkpoints declare additional conversational termination tokens beyond the single end-of-text token used by their base checkpoints. Gemma and K2-Horizon keep declared stopping tokens consistent across compared stages, but learned termination preferences can still diverge.

The evaluation inspects declared end-of-text token configurations across base and post-trained checkpoints from Qwen, Llama, Gemma, and K2-Horizon. Qwen and Llama base models declare a single end-of-text token, while their post-trained versions add conversational termination tokens, creating a configuration-level shift. Gemma and K2-Horizon keep declared stopping tokens consistent across stages, but their learned termination preferences can still diverge, showing that stopping mismatches may arise from training behavior rather than decoding setup alone.