Stable On-Policy Distillation through Adaptive Target Reformulation

TL;DR

This paper addresses instability in on-policy knowledge distillation for language models caused by the inherent asymmetry of KL divergences between teacher and student. It introduces Veto, an objective-level reformulation that constructs a logit-space bridge via a Product of Experts target $Q(y\mid x) \propto P_T(y\mid x) P_S(y\mid x)^{\beta}$, controlled by a single parameter $\beta$. Theoretical analyses show that this bridge stabilizes optimization (Adaptive Gradient Veto), sharpens the student distribution relative to the teacher (via $P_S^*(y\mid x) \propto P_T(y\mid x)^{1/(1-\beta)}$), and connects KD with reinforcement learning through a scaled entropy term in the reverse KL objective. Empirically, Veto yields substantial improvements across reasoning, code generation, and summarization benchmarks, with ablations confirming robustness to data-generation strategies and model heterogeneity. The work provides a versatile, objective-agnostic framework for enhancing stability and generation quality in small-language-model distillation.

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 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.

Figures (5)

• Figure 1: Overview of stable on-policy knowledge distillation.Top (Forward KL): Standard Forward KL triggers explosive gradients when students are ignorant of teacher-preferred tokens (low $P_S$). Veto reformulates the logit-space target to suppress harmful updates to stabilize early optimization. Bottom (Reverse KL): Standard Reverse KL causes mode collapse and diversity loss. By tuning $\beta$, Veto balances generation decisiveness with distributional diversity, allowing the student to capture multiple teacher modes effectively.
• Figure 2: Comparison of gradient magnitudes on ignorant tokens ($P_T > 0.1, P_S < 0.01$). While standard Forward KL triggers pathological gradient explosions exceeding $10^7$, Veto effectively suppresses these spikes within a stable range. This demonstrates Veto’s ability to ensure optimization stability by suppressing harmful updates during early training.
• Figure 3: Effect of different on-policy data generation strategies on student performance for mathematical reasoning. Accuracy (%) is reported along with absolute improvements over the corresponding supervised student baseline for supervised KD, SKD, and on-policy KD under different values of $\beta$.
• Figure 4: Ablation comparing forward and reverse KL divergence objectives. Accuracy (%) is reported for standard on-policy knowledge distillation using each KL loss and for the same settings augmented with the proposed method.
• Figure 5: Ablation comparing forward and reverse KL divergence objectives. Accuracy (%) is reported for standard on-policy knowledge distillation using each KL loss and for the same settings augmented with the proposed method.

Theorems & Definitions (3)

• Theorem 1: Adaptive Gradient Veto
• Theorem 2: Sharpening Effect
• Theorem 3: Bridge to REINFORCE