Stable On-Policy Distillation through Adaptive Target Reformulation

Jul 16, 2026

Upload by Anda

Knowledge distillation (KD) transfers capabilities from large language models to smaller students, but conventional supervised KD can create a training–inference distribution mismatch. On-policy KD learns from student-generated outputs, yet optimization may become unstable when the novice policy diverges sharply from the expert teacher, leading to pathological gradients or diversity collapse. Veto stabilizes objective-level training by reformulating the target distribution in logit space, using a tunable β to veto harmful gradients and balance decisiveness and diversity. Experiments on reasoning and generation tasks show consistent gains over baselines, with code provided.

LanguageEnglish

CategoryResearch & Report

FormatPDF

Pages11

Reading Time28 min

Abstract

1 Introduction

Gradient Forward KL

Veto

Mode Collapsing (Reverse KL)

Figure 1: Overview of stable on-policy knowledge distillation

Stable On-Policy Distillation through Adaptive Target Reformulation Ijun Jang Jewon Yeom Juan Yeo Hyunggyu Lim Taesup Kim† Graduate School of Data Science, Seoul National University {ijun0824, jewon0908, juanyeo, sjrksjek, [taesup.kim}@snu.ac.kr](taesup.kim}@snu.ac.kr) arXiv :2601 .07 155v 3 [ cs .LG] 11 Jul 2026 Abstract 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 β , Veto serves asan Adaptive Gradient Veto that stabilizes optimization by suppressing harmful gradientson low-confidence tokens, while simultaneously acting as a Decisiveness Knob to balance reward-driven performance with output diversity. Experiments across reasoning and generation tasks demonstrate that Veto consistently outperforms supervised fine-tuning and existing baselines. The code is available at [https://github.com/jjun-0824/Veto](https://github.com/jjun-0824/Veto). 1 Introduction Knowledge distillation (KD) (Bucilu et al., 2006 ; Hinton et al., 2015) is widely used for transferring capabilities from proprietary models to efficient open-source counterparts and facilitating model self-improvement (Xu et al., 2024) . Beyond model compression, KD has recently become an important mechanism for improving large language †Corresponding author. models (LLMs) through self-training and alignment, where a strong teacher supervises a weaker or partially trained student. However, traditional supervised KD (Sanh et al., 2019 ; Kim and Rush, 2016) suffers from exposure bias (Ranzato et al., 2015 ; Bengio et al., 2015): the mismatch between teacher-provided trajectories and the student’s selfgenerated outputs leads to degraded performance in autoregressive tasks (Zhang et al., 2019 ; Arora et al., 2022), especially in long-horizon generation and reasoning settings where errors compound overtime. On-policy knowledge distillation addresses this limitation by aligning training with inference-time behavior, learning directly from student-generated outputs (Agarwal et al., 2024 ; Lin et al., 2020) . This enables the student to receive feedback on its own predictions and adapt within regions of the probability space it is likely to visit at test time. Building on this idea, recent methods attempt to bridge the teacher-student gap using reverse KL divergences (Gu et al., 2023 ; Wen et al., 2023) or interleaved sampling strategies (Xu et al., 2025) . While these approaches try to bridge the gap by mixing teacher and student tokens at the data level, they largely overlook a critical challenge: the stability of the optimization objective itself. Even with mixed data, forcing a novice student to match an expert’s sharp distribution creates a steep optimization cliff. In practice, early-stage student policies are highly noisy and often assign near-zero probability to teacher-preferred tokens. As illustrated in Table 1, this leads to unreliable teacher feedback and pathological gradients. Standard forward KL objectives suffer from gradient explosion on such ignorant tokens (Agarwal et al., 2024), while reverse KL objectives, although numerically

FAQs

What problem does this work address in on-policy knowledge distillation?

Conventional on-policy KD can become unstable because the distribution gap between an early-stage student and the expert teacher is often too large, causing pathological gradients (forward KL) or diversity collapse (reverse KL).

How does Veto improve stability?

Veto performs an objective-level reformulation by constructing an intermediate target distribution in logit space that bridges teacher and student agreement regions, suppressing harmful updates on low-confidence tokens.

What role does the parameter β play in Veto?

β acts both as an Adaptive Gradient Veto to prevent gradient explosion in forward KL settings and as a Decisiveness Knob that balances output decisiveness with distributional diversity in reverse KL regimes.

Jul 16, 2026

Upload by Anda