PaperlayerView exact revision on arXiv

Source-linked AI summary

Stable On-Policy Distillation through Adaptive Target Reformulation

Ijun Jang, Jewon Yeom, Juan Yeo, Hyunggyu Lim, Taesup Kim

arXiv:2601.07155v3cs.LGcs.AI

TL;DR

On-policy knowledge distillation struggles with teacher-student distribution gaps that cause unstable optimization and diversity collapse. Veto constructs an intermediate logit-space target, and experiments show consistent gains over supervised fine-tuning and existing on-policy baselines across reasoning, code generation, and summarization.

  • Problem

    On-policy knowledge distillation lacks stable handling of large teacher-student distribution gaps, causing forward-KL gradient explosion and reverse-KL mode collapse.

  • Method

    Veto reformulates the distillation objective by constructing an intermediate logit-space target that suppresses harmful low-confidence updates and controls the diversity-decisiveness trade-off.

  • Results

    Across reasoning, code generation, and summarization tasks, Veto consistently outperforms supervised fine-tuning and existing on-policy baselines, reaching 39.9% accuracy with a +4.8% gain in on-policy distillation.

  • Takeaways & Limitations

    Veto offers one objective-level reformulation for stabilizing forward-KL distillation while balancing performance and diversity in reverse-KL regimes.

Abstract

from arXiv · show

Knowledge distillation (KD) is a widely adopted technique for transferring knowledge from large language models to smaller student models; however, conventional supervised KD often suffers from a distribution mismatch between training and inference. While on-policy KD approaches attempt to mitigate this issue by learning directly from student-generated outputs, they frequently encounter training instabilities because the distributional gap between the novice student and the expert teacher is often too wide to bridge directly. These challenges manifest as pathological gradients in forward KL objectives or diversity collapse in reverse KL regimes. To address these limitations, we propose Veto, an objective-level reformulation that constructs a geometric bridge in the logit space. Unlike prior methods that mix data samples, Veto creates an intermediate target distribution that promotes alignment between the teacher and the student. By introducing a tunable parameter beta, Veto serves as an Adaptive Gradient Veto that stabilizes optimization by suppressing harmful gradients on low-confidence tokens, while simultaneously acting as a Decisiveness Knob to balance reward-driven performance with output diversity. Extensive experiments across various reasoning and generation tasks demonstrate that Veto consistently outperforms supervised fine-tuning and existing on-policy baselines.

1 Introduction

Traditional supervised KD suffers from exposure bias because teacher trajectories differ from student-generated outputs, while on-policy methods face instability when novice students directly match sharp expert distributions. Veto addresses these issues through objective-level logit-space interpolation that stabilizes optimization without changing model architectures.

  • Motivation: Supervised KD suffers from exposure bias because teacher-provided trajectories mismatch student-generated outputs, degrading autoregressive performance, especially in long-horizon generation [Ranzato et al., 2015; Bengio et al., 2015].On-policy KD instead learns from student-generated outputs to align training with inference-time behavior.
  • Motivation: Data-level mixing does not resolve the objective-level gap: sharp teacher distributions create a steep optimization cliff for novice students.The failure modes arise from divergence geometry rather than model architecture or data-generation strategy.
  • Failure modes: Forward KL can produce gradient explosions on teacher-preferred tokens that students assign near-zero probability, whereas reverse KL can cause mode collapse and diversity loss.These contrasting instabilities motivate a unified target reformulation.
  • Veto: Veto constructs a geometric bridge by interpolating teacher and student distributions in logit space, emphasizing agreement and vetoing harmful low-confidence updates.The reformulation improves stability without modifying architectures.
  • Contributions: A single parameterized target unifies forward- and reverse-KL treatment, with theory linking it to adaptive gradient suppression and entropy-regularized policy gradients.Experiments cover reasoning, code generation, and summarization, reporting consistent gains over supervised and on-policy baselines.

2 Related Works

Knowledge distillation compresses LLMs by transferring knowledge from high-capacity teachers to efficient students, but supervised methods face training–inference mismatch. On-policy methods address this mismatch through student-generated outputs, while subsequent objectives improve behavior but largely overlook early optimization stability.

  • Supervised Knowledge Distillation: Knowledge distillation transfers knowledge from high-capacity teachers to efficient students, with supervised approaches mimicking teacher distributions or sequence outputs on fixed datasets.These approaches include distribution-level training on fixed data and sequence-level imitation.
  • On-Policy Knowledge Distillation: On-policy KD aligns training with student inference by learning from student-generated outputs, progressing from mixed fixed and self-generated sequences in ImitKD (Lin et al., 2020) to fully self-generated distillation in GKD (Agarwal et al., 2024).GKD reports that learning from self-generated mistakes can substantially improve performance.
  • On-Policy Objectives: MiniLLM (Gu et al., 2023) and f-distill (Wen et al., 2023) use reverse KL and f-divergences to encourage mode-seeking behavior, while SKD (Xu et al., 2025) uses interleaved sampling to correct low-quality generations.SKD’s interleaved sampling is intended to improve feedback quality.
  • Open Challenge: Despite these advances, existing on-policy approaches largely overlook the stability of the optimization objective, particularly during early training.This limitation concerns the objective itself rather than only the choice of training distribution or sampling strategy.

3 Preliminaries

Knowledge distillation transfers a high-capacity teacher’s predictive capabilities to a smaller student, but supervised KD can cause exposure bias because training uses fixed teacher trajectories. On-policy KD instead samples from the student’s policy to align training with inference, although KL-direction choices create numerical challenges when early outputs are unstable.

  • Supervised KD minimizes divergence against fixed teacher trajectories, exposing the student to distribution shifts it faces during autoregressive inference.This exposure bias arises because the student is not trained on its own inference-time states.
  • On-policy KD samples sequences from the student policy, aligning training with inference and providing feedback on mistakes in likely test-time regions (Agarwal et al., 2024; Xu et al., 2025).
  • Forward KL is zero-avoiding, whereas reverse KL is mode-seeking; both can create numerical challenges in on-policy distillation when early student outputs are highly unstable.Forward KL encourages non-zero probability for teacher-likely tokens, while reverse KL concentrates probability mass on primary modes.

4 Methodology

Veto reformulates on-policy distillation with an intermediate logit-space target that geometrically bridges teacher quality and student confidence. Its adaptive parameter β stabilizes optimization, sharpens decisions, and connects reverse-KL distillation to reinforcement learning while controlling diversity.

  • Analysis I: Stability via Adaptive Veto: In forward KL, Veto prevents pathological gradient explosions by incorporating student uncertainty, causing the loss to converge to 0 when β > 0.The student-dependent term decays faster than log PS diverges, gating updates on tokens about which the student is ignorant.
  • Veto Objective: Veto addresses severe teacher–student distribution mismatch by constructing an intermediate target distribution Q as a geometric bridge rather than matching the teacher directly.Q(y|x) ∝ exp(zT(y|x) + β·zS(y|x)) forms a Product of Experts, assigning high probability only to tokens supported by both teacher quality and student confidence.
  • Analysis I: Stability via Adaptive Veto: Veto sharpens the student toward the teacher through implicit temperature scaling T = 1 − β, encouraging more decisive outputs when 0 ≤ β < 1.The sharpening effect makes the student more decisive than the teacher rather than merely reproducing its uncertainty.
  • Analysis II: Bridge to Reinforcement Learning: In reverse KL, Veto acts as a Decisiveness Knob whose gradient is equivalent to REINFORCE with scaled entropy regularization.This bridges knowledge distillation and reinforcement learning while exposing β as a control over mode-seeking behavior and diversity.
  • The Spectrum of Decisiveness: As β increases from 0 toward 1, the objective moves from distribution matching toward reward-seeking, reducing entropy regularization and potentially collapsing onto the highest-reward mode.Intermediate β values retain a diversity budget proportional to (1 − β), whereas β → 1 yields the REINFORCE objective.

5 Experiments

Across mathematical reasoning, code generation, and dialogue summarization, Veto improves student performance over supervised and on-policy distillation baselines. Ablations indicate that it stabilizes both KL objectives, supports different data-generation strategies, and benefits from adaptive β scheduling and broader model-family applicability.

  • Main Results: Veto improves GSM8K accuracy from 30.7% to 39.9% over supervised student fine-tuning and surpasses SKD and on-policy KD by up to 6.3%.The evaluation uses answer accuracy on the 1,319-instance GSM8K test set.
  • Main Results: Veto raises HumanEval Pass@1 from 22.9 to 29.0 (+6.1) and Pass@10 from 35.3 to 37.7 (+2.4) versus on-policy KD.Code generation is evaluated with pass@k metrics on HumanEval.
  • Main Results: Veto achieves the highest DialogSum win-rate, improving from 54.3 to 56.5 (+2.2) over supervised KD and on-policy KD and from 53.6% to 56.5% over SKD.The gains occur with only 1K student-training examples, indicating effectiveness under limited student data.
  • Data Generation Setting: Across supervised KD, SKD, and on-policy KD, Veto improves mathematical-reasoning accuracy, including 34.3% for supervised KD, a 0.9% gain over its standard baseline.Figure 3 evaluates Veto across multiple on-policy data-generation strategies.
  • Adaptability of KL Loss: Veto improves both forward and reverse KL distillation: it suppresses pathological gradients on ignorant tokens under forward KL and raises reverse-KL accuracy by 2.3% (37.9% →40.2%).These results support Veto as an objective-agnostic reformulation, acting as an Adaptive Gradient Veto for stability and a Decisiveness Knob against premature mode collapse.
  • β Scheduling and Generalizability: A higher β prioritizes optimization stability, whereas a lower β more closely follows the teacher after the student gains proficiency; linear decay outperforms fixed scheduling, with initial β = 0.8 best.The model-family extension also reports consistent improvements for Gemma2-2B-IT distilled from Gemma2-9B-IT on GSM8K at β = 0.3.

6 Conclusion

Veto is an objective-level reformulation for on-policy knowledge distillation that improves optimization stability by constructing a geometric target distribution in logit space. It suppresses pathological forward-KL gradients while balancing reward-driven performance and distributional diversity in reverse-KL regimes.

  • 6 Conclusion: Veto improves on-policy KD optimization stability by constructing a geometric target distribution in logit space that emphasizes teacher–student agreement.The reformulation operates at the objective level rather than by mixing data samples.
  • 6 Conclusion: Veto suppresses the pathological gradients typically encountered in forward KL objectives by using the geometric target to promote teacher–student agreement.
  • 6 Conclusion: In reverse KL regimes, Veto acts as a Decisiveness Knob that balances reward-driven performance with distributional diversity.

A Detailed Mathematical Proofs · A.1 Proof of Theorem 1 (Adaptive Gradient Veto)

The proof contrasts standard forward KD with Veto’s β > 0 formulation, showing that Veto addresses divergent behavior as the student probability for a token approaches zero. The resulting mechanism effectively vetoes updates for tokens the student is ignorant of.

  • A.1 Proof of Theorem 1 (Adaptive Gradient Veto): Standard forward KD with β = 0 has loss L ≈ −c log P_S(y), whose loss and gradient diverge as P_S(y) → 0.
  • A Detailed Mathematical Proofs: Thus, the theorem’s adaptive-gradient interpretation is tied to suppressing updates on low-student-probability tokens rather than allowing the standard forward-KD divergence.
  • A.1 Proof of Theorem 1 (Adaptive Gradient Veto): Veto with β > 0 substitutes a product-of-experts target Q ∝ P_T P_S^β into the loss.
  • A.1 Proof of Theorem 1 (Adaptive Gradient Veto): The proof evaluates Veto’s limiting behavior as the student probability P_S(y) approaches zero from above.
  • A.1 Proof of Theorem 1 (Adaptive Gradient Veto): Applying L’Hôpital’s rule is part of the limiting analysis for the logarithmic term.
  • A.1 Proof of Theorem 1 (Adaptive Gradient Veto): The proof concludes that Veto effectively vetoes updates for tokens where the student is ignorant.

A.2 Proof of Theorem 2 (Sharpening Effect)

The proof shows that forward KL is minimized when the student matches the target distribution, and that for 0 < β < 1 this target induces a sharpened teacher distribution.

  • A.2 Proof of Theorem 2 (Sharpening Effect): Forward KL divergence D_KL(Q∥P_S) is minimized when the student distribution P_S matches the target Q.The proof applies this minimization condition after substituting the target definition from Eq. (4).
  • A.2 Proof of Theorem 2 (Sharpening Effect): For 0 < β < 1, γ = 1/(1−β) > 1, so the student converges to a sharpened version of the teacher distribution.The exponent exceeds one precisely in the stated β range, establishing the theorem’s sharpening effect.

A.3 Proof of Theorem 3 (Bridge to REINFORCE)

The proof starts from the reverse KL objective with the teacher distribution Q treated as fixed during the gradient step. Its reformulation separates a scaled student entropy cost from a teacher-log-probability reward signal.

  • The reverse KL objective is J(θ) = D_KL(P_S∥Q), with Q treated as a fixed target for the gradient step.
  • The reformulated term (1−β) log P_S(y) acts as a scaled entropy cost, while log P_T(y) serves as the reward signal.
Loading 2601.07155v3…