Stable On-Policy Distillation through Adaptive Target Reformulation

AI summary

Overview

Research area: Machine learning — knowledge distillation for large language models, specifically on-policy distillation and the stability of the training objective.

Technical level: Intermediate. The paper assumes familiarity with Kullback-Leibler (KL) divergence, logits, and policy-gradient style objectives, though the core idea is explained through an intuitive "consensus filter" framing.

Scope: The paper proposes Veto, an objective-level reformulation of on-policy knowledge distillation that builds an intermediate target distribution in logit space to stabilize training, and validates it on mathematical reasoning (GSM8K), code generation (HumanEval), and dialogue summarization (DialogSum).

What This Paper Is About

On-policy knowledge distillation trains a small student model on text it generates itself, with a larger teacher model scoring those generations. This avoids the training/inference mismatch of conventional supervised distillation, but it is unstable: a novice student assigns near-zero probability to tokens the teacher prefers, which makes forward KL gradients explode, while reverse KL objectives can collapse the student onto a single mode and lose diversity. Veto addresses this by changing the target of the objective rather than the data or the architecture — it constructs an intermediate distribution that only rewards tokens supported by both the teacher and the student.

Key Contributions

1. An objective-level reformulation. Veto constructs a geometric bridge between teacher and student by defining the target distribution as Q(y|x) ∝ exp(z_T(y|x) + β · z_S(y|x)) = (1/Z(x)) P_T(y|x) · P_S(y|x)^β, a Product of Experts that mixes distributions in logit space instead of mixing data samples. It requires no architectural changes.

2. A single parameter with two roles. The scalar β ≥ 0 acts as an Adaptive Gradient Veto that suppresses pathological gradients on low-confidence tokens in the forward KL regime, and as a Decisiveness Knob that provides controlled entropy regularization in the reverse KL regime.

3. Theoretical analysis. Theorem 1 shows the Veto loss scales as ≈ P_S(y)^β log P_S(y), which converges to 0 as P_S(y) → 0 (via L'Hôpital's rule) where standard forward KD diverges. Theorem 2 shows the student converges to a sharpened teacher, P_S*(y|x) ∝ P_T(y|x)^{1/(1-β)}, with implicit temperature T = 1 − β. Theorem 3 shows the reverse-KL gradient equals a policy-gradient update with advantage A(y) = −log P_T(y) + (1 − β) log P_S(y), i.e. a reward term plus a scaled entropy cost.

4. Empirical validation across three domains. Reasoning, code generation, and summarization experiments show consistent gains over supervised fine-tuning and existing baselines, plus ablations on data generation strategy, KL direction, β scheduling, and a second model family.

Main Findings

• Gradient instability is severe in the baseline. On ignorant tokens (defined as P_T > 0.1, P_S < 0.01), standard Forward KL triggers gradient explosions exceeding 10^7, while Veto suppresses these spikes into a stable range.

• Reasoning gains on GSM8K. Veto reaches 39.9% accuracy versus 30.7% for Student SFT — a 9.2% absolute improvement — and surpasses SKD and on-policy KD by up to 6.3%. For reference, Teacher SFT scores 74.7, Supervised KD 33.4, SKD 33.6, and On-policy KD 35.1.

• Code generation gains on HumanEval. Compared with on-policy KD, Veto improves Pass@1 from 22.9 to 29.0 (+6.1) and Pass@10 from 35.3 to 37.7 (+2.4). Teacher SFT scores 64.7 Pass@1 and 72.2 Pass@10.

• Summarization gains on DialogSum. Veto achieves the highest win-rate among student models at 56.5, up from 54.3 for both Supervised KD and On-policy KD (+2.2), and above SKD's 53.6 (a 2.9% gain). Teacher SFT scores 65.0.

• Domain-specific β values. The authors use β = 0.8 for reasoning, β = 1.0 for code tasks, and β = 0.3 for summarization.

• Gains hold across data generation strategies. Applied to supervised KD, the method reaches 34.3% accuracy (+0.9% absolute over that baseline); in the SKD setting, a 1.3% gain is observed; the largest improvement is on-policy, reaching 39.9% (+4.8% over previous on-policy distillation).

• Veto is objective-agnostic. Applied to forward KL, accuracy rises 4.8% (35.1% → 39.9%) via gradient vetoing; applied to reverse KL, accuracy rises 2.3% (37.9% → 40.2%) via entropy regularization.

• Linear β decay beats a fixed schedule. A linearly decreasing schedule consistently outperforms a constant schedule across all tested values, with an initial β = 0.8 yielding the best performance, as the baseline (β = 0) is shown as a dashed reference line.

• Generalization to another model family. Distilling Gemma2-2B-IT from a supervised-fine-tuned Gemma2-9B-IT teacher on GSM8K with β = 0.3 gives 51.2% accuracy, versus 49.3 for Student SFT, 47.6 for Supervised KD, 47.1 for SKD, and 49.8 for on-policy KD (Teacher SFT: 77.2).

• Qualitative failure modes are documented. Table 1 shows student outputs that produce no answer or wrong logic on GSM8K-style word problems, illustrating the noisy teacher feedback on out-of-distribution student samples that motivates the reformulation.

Methodology in Plain English

Instead of asking the student to match the teacher directly, Veto defines a middle target distribution that combines both models' opinions. In logit space, the target is the teacher's logits plus β times the student's logits, then renormalized. Because this is a product of probabilities, the target is high only when the teacher thinks a token is good and the student is at least somewhat confident about it — the authors call this a "consensus filter." If the student assigns almost no probability to a token, the target probability for that token is pulled down too, which prevents the loss from blowing up.

Training then proceeds in the standard on-policy way: sample prompts, let the student generate trajectories, and minimize KL divergence to the reformulated target — either D_KL(Q ‖ P_θ) in the forward regime or D_KL(P_θ ‖ Q) in the reverse regime. Algorithm 1 adds a linear β decay, β ← β · (1 − i/N), so the veto is strong early (when the student is ignorant) and relaxes later (when the student can follow the teacher's finer details).

Experimental setup: Qwen2-0.5B-IT is the student and Qwen2-7B-IT is the teacher. The teacher is first supervised fine-tuned on task-specific data; the student is then trained on on-policy outputs using the proposed objective with forward KL. All experiments use a learning rate of 1e-5, a warmup ratio of 0.1, a dropout rate of 0.1, three epochs, and 2 H100 GPUs with 80GB of memory. β was selected by grid search and scheduled with linear decay.

Task-specific data: GSM8K uses 7K instances for teacher fine-tuning, 1K for student training pairs, and evaluation on the 1,319-instance test set. Code generation draws from WizardCoder with 10K instances for teacher fine-tuning, 1K for student training, and 1K for evaluation. DialogSum uses 1K instances for student training and the official 1,500-instance test set, with 10K instances for teacher fine-tuning. Summarization quality is measured by win-rate from pairwise comparisons judged by GPT-4o-mini.

Why This Matters

Impact on research. The paper reframes on-policy distillation instability as a problem of objective geometry rather than model architecture or data sampling, and shows that a single scalar parameter can unify the treatment of forward and reverse KL. The connection drawn in Theorem 3 between reverse-KL distillation and an entropy-regularized REINFORCE-style policy gradient gives a concrete interpretation of β along a spectrum from β = 0 (standard reverse KD, mode-covering) through 0 < β < 1 (geometric KD, a middle ground with a diversity budget of (1 − β)) to β → 1 (pure RL, zero entropy regularization, collapse to the highest-reward mode).

Real-world applications.

• Compressing capable proprietary or large models into small, locally deployable models while retaining reasoning ability.
• Improving code assistants built on small models, where Pass@1 and Pass@10 gains matter for practical suggestion quality.
• Dialogue summarization for customer service, meeting notes, or chat logs, where faithful summaries are needed at low inference cost.
• Self-training and alignment pipelines in which a strong model supervises a weaker or partially trained one.

Industry relevance. The approach requires no architectural changes and adds only one hyperparameter, which lowers integration cost. The paper reports gains under supervised KD, SKD, and on-policy KD data generation strategies, meaning teams can adopt it on top of existing distillation recipes. Generalization to the Gemma2 family (Gemma2-9B-IT teacher, Gemma2-2B-IT student) suggests the method is not tied to a single model line. Code is available at the linked GitHub repository.

Future Directions

• Principled β selection. The authors currently select β by grid search and note that optimal values differ sharply by task — 0.8 for reasoning, 1.0 for code, 0.3 for summarization (and 0.3 for Gemma2 reasoning). An automatic or theoretically grounded schedule would remove a tuning burden.

• Extending the RL connection. The paper shows a formal bridge to REINFORCE in the reverse-KL regime; whether the same reformulation can be combined with explicit reward models or full RLHF pipelines is left open.

• Broader model and task coverage. Experiments cover Qwen2 (0.5B student / 7B teacher) and Gemma2 (2B / 9B) on three tasks; generalization to other architectures, scales, and long-horizon or multi-turn settings is not reported.

• Interaction with data-level methods. The authors frame Veto as complementary to interleaved sampling approaches such as SKD, and report gains when combining them, but do not exhaustively characterize how objective-level and data-level interventions interact.

Target Audience

Researchers and practitioners working on knowledge distillation, model compression, or LLM self-training who need a stable training objective for the on-policy setting. It is most useful to readers already comfortable with KL divergence, logits, and policy-gradient objectives, and to engineers who need a low-overhead fix that can be layered onto existing distillation pipelines without changing the model architecture. Readers interested in the theoretical link between distillation and reinforcement learning will also find the Theorem 3 analysis relevant.

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