PaperlayerView exact revision on arXiv

Source-linked AI summary

Less is More: Early Stopping Rollout for On-Policy Distillation

Zhou Ziheng, Jiaqi Li, Huacong Tang, Ying Nian Wu, Demetri Terzopoulos

arXiv:2605.27028v1cs.LGcs.AI

TL;DR

OPD can lose corrective teacher supervision when later tokens are conditioned on off-policy student prefixes. The paper proposes ESR, which limits rollouts and distillation to the first N response tokens, and reports broad performance gains with improved efficiency and stability, alongside mechanisms that may explain its effectiveness.

  • Problem

    OPD lacks reliable late-position corrective supervision when the teacher is conditioned on increasingly off-policy student prefixes.

  • Method

    Early Stopping Rollout (ESR) restricts student rollouts and the distillation loss to the first N response tokens.

  • Results

    ESR outperforms full-rollout OPD across tasks, model families, scales, and training regimes, with up to 24× lower wall-clock cost and 4× lower peak training memory.

  • Takeaways & Limitations

    Cascading Alignment and Sub-mode Commitment may explain ESR’s effectiveness, including cases where ESR-trained students exceed the teacher.

  • Takeaways & Limitations

    The experiments focus on instruction-tuned sub-100B open-source models and limited data budgets, so ESR’s behavior at trillion-scale industrial settings remains unclear.

Abstract

from arXiv · show

On-policy distillation has recently emerged as a promising alternative to standard sequence-level imitation, training a student by scoring its own rollouts with a teacher model. However, we observe ``Off-policy Teacher Decay'' problem in this paradigm: for the later tokens, with student's earlier trajectory as context that is off-policy to the teacher, the teacher's ability to produce a corrective score would decay, and may fall back to token-completion behavior learned in the pre-training stage. We empirically verify this problem, and we propose Early Stopping Rollout (ESR) to fix it: a simple yet effective distillation strategy that simply restricts the rollout generation to the first response tokens. We show that ESR both surpasses the full rollout OPD performance across model size, family, tasks and training regime, and exhibit much higher GPU efficiency and training stability, especially under cross model family scenarios. We further investigate the mechanism behind this surprising performance and discovered "Cascading Alignment" and "Sub-mode Commitment" effect of ESR that may explain why it works effectively and even sometimes exceeding the teacher model performance. Besides, we show that this position-based token selection strategy cannot be fully explainable by KL divergence and entropy signals.

1 Introduction

The paper identifies off-policy teacher decay in full-rollout OPD and proposes Early Stopping Rollout (ESR), which restricts distillation to early response tokens. ESR outperforms full-rollout OPD across broad settings while improving efficiency and stability, with cascading alignment and sub-mode commitment offered as explanatory mechanisms.

  • Problem: Full-rollout OPD can suffer off-policy teacher decay because later teacher scores condition on student-generated context that may no longer support corrective guidance.The teacher may instead fall back to token-completion behavior from the off-policy state.
  • Results: ESR consistently outperforms full-rollout OPD across tasks, model families, model scales, and training regimes while being more efficient and stable.The reported settings include math, code, function calling, LoRA, full fine-tuning, students from 1.5B–32B, and teachers from 1.7B–72B.
  • Results: ESR is especially valuable in cross-generation and cross-family settings, where full-rollout OPD frequently becomes unstable or collapses.ESR remains stable across examples such as Qwen generations and Gemma-to-Qwen transfer.
  • Analysis: Cascading Alignment, Sub-mode Commitment, and position-based token selection beyond KL and entropy are investigated as mechanisms behind ESR’s effectiveness.These mechanisms may explain late-token improvement and cases where ESR-trained students exceed the teacher.
  • Method: ESR restricts student rollouts and distillation loss to the first N response tokens through a one-line change.The method is presented as a minimal modification to on-policy distillation.

2 Off-Policy Teacher Decay in OPD

The paper characterizes Off-policy Teacher Decay as a failure mode of OPD: student prefixes increasingly move away from the teacher’s distribution, weakening late-position corrective supervision. A recoverability-gap analysis and MATH-500 experiment show that teacher performance declines as the student prefix grows.

  • Failure mode: Off-policy Teacher Decay occurs when OPD evaluates later tokens after conditioning the teacher on increasingly off-policy student prefixes.This challenges the assumption that the teacher can provide useful corrective supervision at every position.
  • Failure mode: The teacher may leave its natural reasoning state and revert to completing tokens from the student-induced off-policy state.The passage describes this as a fallback behavior rather than continued correction of the student trajectory.
  • Measurement: The recoverability gap ∆decay(t) measures the difference between teacher accuracy from the original prompt and after continuing from a length-t student prefix.A larger gap indicates weaker recovery and a less reliable late-position corrective target.
  • Evidence: On MATH-500, teacher avg@4 accuracy declined from 65.30% unconditionally to 62.70% at N=100 and 51.75% at N=300.At N=300, performance approached the student baseline, suggesting that late-position scores increasingly reflect trajectory continuation.

3 Method: Early Stopping Rollout (ESR)

Standard on-policy reverse-KL distillation scores student-generated responses token by token with the teacher. ESR instead truncates rollouts to the first N response tokens and computes the loss only on that early window, leaving the rest of the training loop unchanged.

  • Standard OPD: Standard on-policy reverse-KL distillation has the student generate a response conditioned on the prompt, then compares student and teacher token distributions.The student is denoted πs and the teacher πt.
  • Cutoff choice: For the MATH-500 Qwen2.5-Math-1.5B → Qwen3-1.7B setup, performance saturates for N ∈[50, 200], and every tested cutoff beats OPD.
  • Early Stopping Rollout: ESR truncates the student rollout to its first N tokens and computes the distillation loss over exactly those tokens.The cutoff satisfies N ≪ T in practice.
  • Early Stopping Rollout: If the student emits EOS before N, the rollout terminates naturally; generation temperature, LoRA modules, optimizer, and scorer remain unchanged.

4 Main Experiments

Experiments evaluate ESR across model pairings, scales, tasks, and training regimes. ESR generally matches or outperforms OPD while providing greater stability and substantially lower training cost, especially in cross-generation and cross-family settings.

  • Overall performance: ESR matches or beats OPD across same-family, cross-generation, and cross-family model pairs, with students from 1.5B–32B and teachers from 1.7B–72B on MATH-500.Table 1 reports peak OPD performance across training and marks teacher-surpassing ESR results.
  • Task and training generalization: ESR generalizes across math, coding, and tool-calling tasks, and across LoRA and full fine-tuning regimes.The evaluated metrics include MATH-500 avg@4, HumanEval pass@1, and BFCL full accuracy.
  • Training-regime comparison: For full fine-tuning, ESR dominates OPD on the Gemma pair by 12.75% avg@4 and +15.40% pass@4, while Qwen OPD scores 58.20 versus ESR 56.20 avg@4.
  • Training stability: In cross-generation and cross-family settings, full-rollout OPD frequently degrades or collapses, whereas ESR does not degrade in the reported experiments.The paper reports more than 10% avg-accuracy improvements many times in these settings, while same-family same-generation gains are smaller.
  • Cutoff sensitivity: ESR reaches comparable performance from N=50 through N=200 in a cross-generation sweep, although the exact cutoff is more sensitive in cross-family settings.
  • Efficiency: 24× wall-clock speedup and approximately 4× lower peak training memory are reported for ESR with N=100.The comparison uses a single A6000, batch size 16, and Qwen3-1.7B student and teacher models; generation falls from 180 s/step to 5 s/step.
  • Efficiency: ESR permits the student and teacher models to fit comfortably on one A6000 GPU, avoiding additional model loading and unloading overhead in the authors’ practice.

5 More Analysis on Why ESR Works

ESR’s effectiveness is linked to cascading alignment from early strategic tokens and sub-mode commitment under reverse KL, while position remains distinct from KL and entropy as a selection signal.

  • Cascading Alignment: 30–40% of the per-position KL beyond the trained [0, N] window drops after ESR training, despite receiving no direct loss.This convergence cascade suggests early-window supervision can align later rollout positions.
  • Cascading Alignment: Early tokens often establish problem framing and strategy, while later tokens mainly execute once the strategy is fixed.In the MATH-500 case study, the first 100 tokens determine the geometric setup and key relationship, whereas later tokens perform algebraic computation.
  • Sub-mode Commitment: ESR-trained students produce median responses of ∼380 tokens, versus ∼1,150 for the teacher and ∼1,530 for full-rollout OPD.The shorter responses are consistent with commitment to a more concise behavioral mode.
  • Sub-mode Commitment: ESR shifts choices toward the teacher’s top 2–5 alternatives while increasing top-1 confidence: 47.4% versus 44.6%, 41.9% versus 45.7%, and 0.79 versus 0.77.These pairs compare ESR with full-rollout OPD for top 2–5 membership, top-1 agreement, and top-1 probability, respectively.
  • Sub-mode Commitment: Reverse KL can penalize unsupported student mass without preventing concentration on one supported teacher mode, enabling commitment to a better sub-mode.The schematic contrasts a verbose teacher mode with a concise correct mode and describes late-position OPD as reverting toward the teacher’s average.
  • Position Versus Scalar Signals: Top-KL and entropy-based selectors underperform ESR, while the largest 100 KL tokens account for around 93% of trajectory loss.Thus, high signal magnitude does not identify the most effective supervision positions, and position cannot be reduced to KL or entropy alone.

6 Related Work

The paper situates ESR among distillation, token-importance, and concurrent prefix-supervision studies, emphasizing that positional selection differs from scalar saliency and depends on the starting student setting.

  • Knowledge Distillation for Language Models: Prior LLM distillation work studies reverse KL, student-generated rollouts, alternative divergences, and adaptive sampling choices.These works establish the broader design space in which ESR changes rollout supervision.
  • Token-Level Importance in Distillation and Reasoning: Reasoning studies identify high-entropy forking tokens, planning anchors, and functionally important positions as contributors to downstream reasoning.These findings motivate examining token-level structure rather than treating all reasoning tokens identically.
  • Token-Level Importance in Distillation and Reasoning: Concurrent token-selection methods use teacher verification, adaptive temperature, entropy filtering, or preference ranking, but the paper reports these scalar criteria underperform ESR.The authors interpret position as a load-bearing supervision axis rather than merely a proxy for token saliency.
  • Concurrent Work: A concurrent prefix-OPD study finds prefix supervision does not surpass full-trajectory OPD when reasoning is bootstrapped from a base model.The present work instead studies math-SFT students that already reason and reports better ESR performance than full rollout.

7 Conclusion

The conclusion presents ESR as a minimal prefix-supervision modification that improves distillation performance, efficiency, and stability, while identifying mechanisms and boundaries for its applicability.

  • Conclusion: ESR constrains on-policy rollouts to the first N response tokens and outperforms full-rollout OPD across tasks, scales, and training regimes.The modification is described as a one-line change with benefits across performance, efficiency, and stability.
  • Conclusion: Off Policy Teacher Decay is identified as the root problem ESR mitigates, while Cascading Alignment explains later-token improvement without direct late-token training.The conclusion also names Sub-mode Commitment as a mechanism behind occasional performance beyond the teacher.
  • Limitations: The experiments are limited to instruction-tuned students with basic thinking ability, small open-source models under 100B, specific tasks, and limited data budgets.The paper has not tested trillion-scale capacity improvement, multimodality, or long-horizon tasks, so ESR’s broader scope remains unclear.

A Training Efficiency

The supplied efficiency passages identify a single-A6000 per-step measurement setting and a memory constraint for an 8B teacher configuration.

  • Training Efficiency: Training efficiency is measured by per-step wall-clock time on a single A6000 48GB GPU with a Qwen2.5-Math-1.5B LoRA student and batch size 16.This specifies the hardware and training configuration for the efficiency breakdown.
  • Training Efficiency: An 8B teacher combined with vLLM and the student frequently causes out-of-memory errors because 48GB is insufficient.The implementation note marks a practical memory boundary for this configuration.

B.1 Math Results: Per-Step Performance

The primary math experiments report per-step results for LoRA with n=1 across 3,200 problems. The MATH-500 baseline is 50.95% avg@4.

  • Table 6 reports per-step results for the primary math experiments.The experiments use LoRA, n=1, and 3,200 problems.
  • The primary math experiments use LoRA with n=1 on 3,200 problems.
  • 50.95% avg@4 is the MATH-500 baseline.

B.2 Math Results: Full Per-Step Trajectories (n=1, 3,200 problems, bs=16)

The paper reports complete results for MATH-500, coding, and function calling under LoRA. These evaluations use different task-specific metrics, including pass@1 and accuracy-related measures.

  • MATH-500: Table 7 reports complete MATH-500 results for LoRA with n=1, batch size 16, and 3,200 problems.Best results per configuration are bolded.
  • Coding: Table 8 reports complete LoRA coding results on HumanEval using pass@1.
  • Function calling: Table 9 reports LoRA function-calling results on BFCL using name accuracy, full accuracy, and parse rate.Best full_acc is bolded.

C Token Classification Methodology

The methodology classifies tokens into six categories, including planning, structural, numerical, operator, LaTeX, and continuation tokens. Additional tables analyze KL by category and position, identify high-KL tokens, and document asset licensing.

  • Token categories: Tokens are classified into six categories for the analysis.The supplied passages specify planning, structural, math_number, math_operator, math_latex, and continuation categories.
  • Token categories: Planning tokens include reasoning keywords such as “To,” “First,” “Therefore,” and “Thus.”
  • Token categories: Structural tokens comprise punctuation, whitespace, and formatting tokens.
  • Token categories: Math-number tokens are digits, math-operator tokens are arithmetic operators, and math-LaTeX tokens are LaTeX delimiters.The listed arithmetic operators are +, −, ×, /, and =; the listed delimiters are \( and \[.
  • Token categories: Continuation tokens include all tokens not assigned to the other categories.
  • Supplementary analyses and assets: Tables analyze mean KL by token category and position range, list the 20 highest-KL tokens, and document the paper’s licensed assets.The highest-KL-token table uses a minimum of 50 occurrences across 10,000 trajectories, and asset licenses were verified at submission.
Loading 2605.27028v1…