PaperlayerView exact revision on arXiv

Source-linked AI summary

Hybrid Policy Distillation for LLMs

Wenhong Zhu, Ruobing Xie, Rui Wang, Pengfei Liu

arXiv:2604.20244v2cs.CLcs.AI

TL;DR

LLM distillation is shaped by coupled choices of divergence, optimization, and data regime, complicating effective compression. This paper unifies these choices and introduces HPD, which consistently matches or outperforms standard SFT + OPD pipelines with less cost across tasks.

  • Problem

    LLM distillation lacks a unified account of how divergence direction, optimization strategy, and data regime jointly govern its behavior.

  • Method

    HPD combines a unified reweighted log-likelihood view with token-level reverse-KL signals and lightweight prefilling-time student sampling.

  • Results

    HPD consistently matches or outperforms standard SFT + OPD pipelines with less cost across math reasoning, dialogue, and code tasks.

  • Takeaways & Limitations

    HPD provides an efficient, stable distillation approach and a strong initialization for subsequent OPD.

  • Takeaways & Limitations

    Monte Carlo estimation of the reverse-KL divergence is unbiased but suffers from high variance.

Abstract

from arXiv · show

Knowledge distillation (KD) is a powerful paradigm for compressing large language models (LLMs), whose effectiveness depends on intertwined choices of divergence direction, optimization strategy, and data regime. We break down the design of existing KD methods and present a unified view that establishes connections between them, reformulating KD as a reweighted log-likelihood objective at the token level. We further propose Hybrid Policy Distillation (HPD), which integrates the complementary advantages of forward and reverse KL to balance mode coverage and mode-seeking, and combines off-policy data with lightweight, approximate on-policy sampling. We validate HPD on long-generation math reasoning as well as short-generation dialogue and code tasks, demonstrating improved optimization stability, computational efficiency, and final performance across diverse model families and scales. The code related to this work is available at https://github.com/zwhong714/Hybrid-Policy-Distillation.

1. Introduction

LLMs’ capabilities create substantial computational and storage costs, motivating knowledge distillation for practical deployment. The paper unifies distillation objectives and proposes HPD to combine forward/reverse divergence with on-/off-policy data while preserving efficient supervision.

  • Motivation: LLM performance gains from large models and high-quality data incur substantial computational and storage overhead, making compression important for practical deployment.
  • Background: Knowledge distillation transfers knowledge from powerful teacher models to lightweight students through black-box SFT or white-box distribution-level objectives.
  • Design challenges: Effective LLM distillation depends on coupled choices including divergence direction, with forward KL supporting mode coverage and reverse KL providing complementary inductive bias.
  • Unified formulation: The paper formalizes SFT, FKLD, and RKLD through a unified reweighted log-likelihood view spanning divergence directions, optimization strategies, and data regimes.
  • Hybrid Policy Distillation: HPD combines divergence directions and on-/off-policy regimes through token-level reward signals, using reverse-k1 discrepancy and prefilling-time student sampling without full-vocabulary matching or sequence rollouts.
  • Empirical findings: Experiments on math reasoning, dialogue, and code tasks show HPD improves optimization stability and efficiency, matches or outperforms standard SFT + OPD with less cost, and strengthens OPD initialization.

2. Problem Formulation

The paper formulates autoregressive language modeling as sequential decision-making over expert trajectories from an offline dataset, with an LLM policy generating next tokens conditioned on ground-truth prefixes. It identifies a capacity-mismatch problem: multimodal teacher or offline distributions can exceed the modes representable by a limited student, diluting probability mass across modes.

  • Auto-regressive Language Models: Autoregressive generation is modeled as a sequential decision process over trajectories of states and expert actions.Each state represents the ground-truth prefix, and each expert action is a token from a finite vocabulary.
  • Auto-regressive Language Models: A pretrained LLM defines a parameterized policy qθ(at | st) that factorizes the likelihood of an expert trajectory.The policy conditions next-token generation on the corresponding trajectory state.
  • Auto-regressive Language Models: Teacher forcing pretrains the model by minimizing expected negative log-likelihood over the offline dataset D.This objective uses expert trajectories as the training signal.
  • Limited Expressivity of the Student Model: Multiple valid continuations and label smoothing create soft supervision, while complex output spaces can make teacher or offline distributions highly multimodal.A capacity-limited student may represent fewer modes, causing distillation to dilute probability mass across them.

3. Preliminary

This section frames SFT as sparse-supervision distillation, contrasts forward and reverse KL’s coverage and mode-seeking behaviors, and motivates Monte Carlo approximations for otherwise intractable vocabulary-level divergence computation.

  • SFT and KD: SFT supplies supervision only for the target action, whereas KD uses the teacher’s full predictive distribution for denser learning signals.SFT therefore ignores alternative plausible actions, while KD provides broader token-level information.
  • Forward and Reverse KL: Forward KL encourages coverage of all teacher-likely behaviors but can produce overly smooth distributions when the student lacks sufficient capacity.It heavily penalizes missing modes, motivating broad mode coverage at the risk of smoothing.
  • Forward and Reverse KL: Reverse KL focuses the student on the teacher’s dominant modes, potentially underrepresenting valid low-probability outputs and causing unstable training from high-variance log-ratio gradients.The instability arises when the student is poorly aligned with the teacher and the log-ratio is unbounded.
  • Monte Carlo Estimation: Exact KL computation is often intractable because it sums over the full vocabulary, sequence length, and batch size, so Monte Carlo sampling is used instead.The simplest K1 estimator is unbiased for DKL(qθ∥p) but has high variance because its log-ratio is negative for many samples.
  • Monte Carlo Estimation: K1 variants can serve as token-level reward penalties or explicit loss terms, and OPD uses negative K1 on student-sampled tokens as a reward for efficient distillation.The passage also attributes an unbiased gradient estimator to this approach.

4. Our Method: HPD

HPD formulates distillation as token-level reweighted likelihood optimization and combines forward- and reverse-KL signals through asymmetric expert and student-sampled token updates. Its masking and reweighting mechanisms reinforce underestimated expert tokens while suppressing unreasonable non-expert behavior.

  • Reweighted Likelihood View: SFT and KL distillation are unified as reweighted likelihood optimization over a data source and token-level student–teacher discrepancies.The data source may be on-policy samples from the student or off-policy samples from a fixed dataset or teacher policy.
  • Optimization Effects: Positive weights increase sampled-token likelihood, whereas negative weights suppress the token and redistribute probability mass across the predictive distribution.The redistribution follows the current probabilities of alternative tokens, so token-level objectives induce gradients over the full vocabulary.
  • Hybrid Policy: HPD samples a non-expert token from the student under the offline prefix solely to identify and suppress unreasonable behavior.Positive sampled-token reverse-KL gaps are masked, preventing reinforcement of non-expert tokens while limiting unreasonable sampling.
  • Intuitive Explanation of HPD: When the expert is underestimated and the sampled non-expert is overestimated, HPD strengthens the expert-token weight so suppressed probability mass is redirected toward the expert action.The expert token is reinforced with a forward-KL weight when k1 > 0, and this weight is further strengthened when k′1 < 0.

5. Experiments

Experiments evaluate HPD on long-generation mathematical reasoning and short-generation dialogue and code tasks using separate offline and on-policy settings. Across these settings, HPD improves optimization stability, policy alignment, robustness, and distillation performance over the reported baselines.

  • Experimental Setup: Experiments cover long-generation mathematical reasoning and short-generation dialogue and code tasks, with offline and on-policy data considered separately.Reasoning uses Qwen2.5 and LLaMA3 model families, while dialogue and coding use Instruct models with Ultrafeedback and WizardCoder data, respectively.
  • Long-Generation Reasoning: HPD maintains stable entropy, reduces the student–teacher distribution gap, and steadily improves generation performance, unlike SFT’s rapid overfitting and entropy collapse.HPD also aligns student behavior with the teacher across training and inference-time entropy dynamics.
  • Long-Generation Reasoning: 41.0% improves Qwen 2.5 3B from 28.25 to 39.83, while 77.9% improves LLaMA 3 3B from 19.43 to 34.56.HPD consistently outperforms baseline distillation methods across both model families and enables 3B students to approach larger-model reasoning capabilities.
  • Short-Generation Dialogue: HPD generally achieves the highest dialogue scores on MT-1T and MT-2T while leading key alignment metrics including AE-LC, AE-WR, and Arena-WR.The results particularly show stronger preservation of multi-turn conversational coherence and contextual understanding.
  • Short-Generation Code: HPD achieves the best average coding performance across both model families and shows more stable improvements than KD and JSD across tasks and families.The results support a more balanced and reliable distillation objective, even when HPD does not lead every individual DS-Coder benchmark.
  • On-Policy Reasoning: HPD consistently achieves higher task performance during on-policy training, with more stable advantage estimates and lower KL divergence to the teacher.When combined with OPD, HPD further amplifies gains and attains the highest average scores over corresponding baselines.

6. Ablation Study

The ablation study isolates HPD’s two components—student sampling and the Reinforce operation—and finds that both are important for sustained, stable optimization without adding hyperparameters.

  • Ablation Design: HPD’s ablation removes student sampling and the Reinforce operation separately to disentangle their individual contributions, without introducing additional hyperparameters.The two components are student sampling from the student’s preferred actions and reinforcing expert tokens when unreasonable actions are suppressed.
  • Effectiveness of Student Sampling: Without student sampling, performance converges rapidly but plateaus, whereas student sampling enables exploration of diverse trajectories and continued policy refinement.Direct optimization toward the teacher distribution limits exploration and causes premature convergence.
  • Necessity of the Reinforce Operation: Removing the Reinforce operation slows KL loss reduction, while reinforcing expert tokens after suppressing unreasonable actions provides a more stable optimization signal.The operation accelerates alignment with the teacher distribution and produces consistent performance improvements.

7. In-Depth Analysis

HPD reduces distillation overhead by making token-level decisions on only two actions instead of matching full-vocabulary or top-k distributions. It also improves Qwen3-0.6B performance when distilled from a much stronger teacher, with larger gains under a 38K-token evaluation cap.

  • Efficiency: HPD makes token-level decisions only on the expert token and a sampled student token, avoiding full-vocabulary KL computation and top-k distribution matching.This design uses teacher-distribution information while limiting the comparison to two token-level actions.
  • Efficiency: 278 GiB is the additional distribution-matching buffer required by HPD, reduced from the TiB scale while keeping model-state memory close to other white-box KD methods.HPD’s wall-clock cost is also closer to SFT than to AKL or ABKD.
  • Stronger-teacher distillation: Under a 20K training upper bound, HPD improves Qwen3-0.6B when distilled from Qwen3-30B-A3B-Thinking-2507, testing a larger teacher–student gap than the main experiments.The stronger-teacher study uses a 20K upper bound, compared with 8K in the main experiments.
  • Stronger-teacher distillation: 39.03% and 28.22% are HPD’s relative gains on AIME-24 and AIME-25, respectively, when the evaluation cap is extended to 38K tokens.AIME-24 improves from 11.35 to 15.78, while AIME-25 improves from 18.85 to 24.17.
  • Stronger-teacher distillation: 8.36% and 12.42% are HPD’s relative gains on AIME-24 and AIME-25, respectively, under the 20K setting.AIME-24 improves from 11.25 to 12.19, and AIME-25 improves from 18.44 to 20.73.

8. Broader Impacts

HPD is presented as a versatile foundation for building strong small models and supporting subsequent preference-alignment stages. Iterative DPO and self-distillation further transfer and scale performance while retaining the original SFT data.

  • HPD is presented as broadly versatile, with diverse applications suggesting potential for future use.
  • The pipeline alternates DPO training with self-distillation, updates the teacher from the previous DPO stage, and retains the original SFT dataset.This differs from approaches that replace SFT data with responses generated by the DPO-trained model.
  • Iterative DPO on the HPD model transfers teacher performance without degradation and continues increasing alignment performance, though gains become limited across epochs.The results also support strong-to-weak distillation as an effective starting point.

9. Related Work

Prior work on LLM knowledge distillation includes off-policy approaches that use teacher generations or divergence losses, and on-policy approaches that focus students on selected teacher modes. Representative methods address overestimation of low-probability teacher regions or contrast teacher- and student-generated responses.

  • Off-policy Distillation: Off-policy distillation compresses neural networks by enabling smaller students to match larger teachers, using teacher generations as supervised fine-tuning data or divergence-based distribution alignment.These approaches apply when the teacher model is accessible.
  • On-policy Distillation: On-policy distillation encourages students to concentrate on a limited set of modes in the teacher’s distribution.MiniLLM and GKD use RKLD or JSD to mitigate students’ overestimation of low-probability teacher regions.
  • On-policy Distillation: DistillLLM-2 uses contrastive distillation to increase the likelihood of teacher-generated responses while decreasing the likelihood of student-generated responses.The supplied passage identifies this as a contrastive formulation within on-policy distillation.

10. Conclusion … F.2. Teacher models

The paper unifies LLM knowledge distillation through a reweighted log-likelihood view and introduces HPD, combining reverse-KL signals with prefilling-time student sampling. It also analyzes estimator gradients, positions the method as efficient off-policy initialization, and documents experimental baselines and training configurations.

  • 10. Conclusion: HPD uses reverse-KL discrepancy signals and prefilling-time student sampling to decide which expert actions to reinforce or suppress.This token-level method is developed from the unified reweighted log-likelihood perspective on distillation.
  • A. Positioning of Our Work: The method distills on off-policy data with lightweight approximate on-policy next-token sampling, avoiding full-sequence rollouts while supporting later training.Full rollouts, as used in OPD or reinforcement learning, are acknowledged as necessary for further performance enhancement.
  • B. Limitation: The approach assumes teacher and student models share a tokenizer and uses token-level estimators because full KL divergence is intractable.Shared tokenization is identified as a common assumption in white-box knowledge distillation.
  • B.1. Gradient Analysis with K1 reward: Using the K1 estimator as a reward gives an unbiased gradient estimate for the KLD objective, whereas inserting K1 directly into the loss does not.The analysis compares score-function treatment of the estimator with direct loss-based differentiation.
  • C. Gradient of Jensen-Shannon Divergence: The JSD gradient yields a token-level weight, wJSD(at|st), analogous to reward weights in SFT, FKLD, and RKLD.The gradient with respect to θ depends only on the term involving the student distribution.
  • D. Gradient contribution: The gradient contribution is derived by applying the chain rule and evaluating softmax derivatives in two cases before combining them into Eq. (20).The resulting expression is rewritten in terms of softmax probabilities.
  • E. Baselines: The evaluated baselines include SFT, FKLD-based KD, teacher-sequence SeqKD, RKLD-based MiniLLM, and JSD-based GKD across fixed, teacher-generated, and student-generated data.MiniLLM extends RKLD to student-generated data, while GKD is evaluated on ground-truth, teacher-generated, and on-policy settings.
  • F.2. Teacher models: Teacher-model experiments use Qwen-2.5-7B and Llama-3.1-8B with SFT on OpenR1-Math-8192, while DAPO trains with clip-higher 0.28 on DAPO-MATH-17k.The learning-rate setup fine-tunes the corresponding base model for 10 epochs on 40k data, with cosine decay to one-tenth of the initial rate, maximum context length 10240, and batch size 256.

F.3. Evaluation Benchmark. · G. Off-policy Data for Personalization · G.1. Learning Rate

The evaluation covers in-domain mathematics benchmarks and the out-of-domain GPQA task, using standardized long-generation sampling settings. Personalization experiments fine-tune the base model on UltraFeedback with a specified learning rate, context length, batch size, and epoch count.

  • F.3. Evaluation Benchmark.: The in-domain evaluation includes AIME24, AIME25, AMC, MATH-500, and OlympidBench.These benchmarks comprise the paper’s in-domain task suite.
  • F.3. Evaluation Benchmark.: The out-of-domain evaluation uses GPQA.GPQA is identified as the out-of-domain task.
  • F.3. Evaluation Benchmark.: All evaluations generate responses with a maximum sequence length of 10,240, top-p sampling of 0.95, and temperature 0.7.These settings are applied across the reported evaluations.
  • G. Off-policy Data for Personalization: The personalization setup uses off-policy UltraFeedback data with 4,096-token teacher-generated responses.The data is used to fine-tune the corresponding base model.
  • G.1. Learning Rate: The base model is fine-tuned for 10 epochs with a learning rate of 5 × 10−6.These are the reported training duration and learning-rate settings.
  • G.1. Learning Rate: The fine-tuning configuration uses a maximum context length of 4,096 and a batch size of 256.Both values are specified for the UltraFeedback fine-tuning run.

G.2. Evaluation Benchmark. … K. Case Study

The paper evaluates alignment across three benchmarks with fixed generation and judging settings, details training and data-integrity procedures, and uses a case study to compare distilled models’ recommendations. The case study finds that SFT, KD, HPD, and RKLD recommend Park Hyatt Istanbul, while HPD produces especially detailed, teacher-like responses.

  • G.2. Evaluation Benchmark.: The evaluation uses MT-Bench, AlpacaEval, and Arena-Hard, with Qwen3-30B-A3-Instruct-2507 serving as the judge model.Responses use maximum sequence length 4,096, top-p 0.95, and temperature 0.7; judge evaluations use temperature 0.0.
  • H. Off-policy Data for Coding; H.1. Learning Rate: Coding models are fine-tuned for 10 epochs on WizardCode teacher-generated data using learning rate 5 × 10−6, maximum context length 4,096, and batch size 256.The WizardCode data contains 4,096-token teacher-generated examples.
  • I. On-policy Data for Reasoning: Qwen and LLaMA use learning rate 3 × 10−6 with on-policy student rollouts and batch size 256, corresponding to 64 prompts with four repeated rollouts per prompt.Evaluation follows the procedure described in Section F.3.
  • J. Broder Impact: In Section 8, Qwen2.5-7B-Base is selected as teacher; base models receive SFT and are then aligned with DPO implemented in LLama-Factory.Detailed DPO configurations are summarized in Table 12.
  • J.1. System Prompt: The reasoning task uses a step-by-step system prompt requiring the final answer within \boxed{}, whereas alignment and code tasks use “You are a helpful assistant.”
  • J.2. Data Decontamination: Decontamination combines light and aggressive normalization, exact-match detection, 6-gram candidates, and SequenceMatcher near-duplicate analysis; AIME-24, AIME-25, AMC, GPQA, and Minerva show zero overlap.Only MATH-500 and OlympiadBench contain flagged rows, reported in Table 13 for transparency.
  • J.3. Confidence Intervals: Arena-Hard alignment results are reported with 95% confidence intervals, while coding tasks use greedy decoding.The confidence-interval results appear in Table 14.
  • K. Case Study: The case study shows SFT, KD, HPD, and RKLD recommending Park Hyatt Istanbul, while JSD gives a different recommendation and RKLD also lists alternatives.HPD tends to provide more detailed information and thoughtful, user-tailored greetings resembling the teacher model.
Loading 2604.20244v2…