Hybrid Policy Distillation for LLMs

1. Problem Statement

Large Language Models (LLMs) achieve superior performance but incur massive computational and storage costs. Knowledge Distillation (KD) is a primary method for compressing these models, yet existing approaches face significant trade-offs governed by three intertwined design axes:

1. Divergence Direction:
   • Forward KL (FKLD): Promotes "mode coverage" (ensuring the student covers all teacher modes) but often leads to over-smoothed, low-confidence predictions.
   • Reverse KL (RKLD): Promotes "mode-seeking" (focusing on high-probability teacher modes) but can cause training instability and collapse when the student-teacher gap is large.
2. Optimization Strategy: Whether the divergence is treated as a direct loss or a reward signal, leading to different gradient dynamics.
3. Data Regime:
   • Off-policy: Uses static teacher data (efficient but suffers from train-inference mismatch).
   • On-policy: Samples from the student policy (better alignment but computationally expensive due to full rollouts).

Current methods typically choose one direction (FKLD or RKLD) and one regime, failing to balance the complementary benefits of both or to efficiently handle the trade-off between stability and performance.

2. Methodology: Hybrid Policy Distillation (HPD)

The authors propose Hybrid Policy Distillation (HPD), a unified framework that integrates the advantages of FKLD and RKLD while combining off-policy data with lightweight on-policy sampling.

A. Unified Theoretical View

The paper first reformulates KD as a reweighted log-likelihood objective at the token level:
$$\mathcal{L}(\theta) = \min_{\theta} -\mathbb{E}_{(s_t, a_t) \sim D_\pi} [w(a_t | s_t) \log q_\theta(a_t | s_t)]$$
Here, $w(a_t | s_t)$ acts as a reweighting signal. The authors analyze how different methods (SFT, FKLD, RKLD) correspond to specific choices of this weight. They demonstrate that negative weights (characteristic of Reverse KL) actively suppress the sampled token and redistribute probability mass to other tokens, whereas positive weights reinforce the token.

B. The HPD Algorithm

HPD operates on an asymmetric token-level design for each state $s_t$ in an offline dataset:

1. Expert Token Processing (Ground Truth/Teacher Output):

   • Computes a negative Reverse KL ($k_1$) estimator for the expert token $a^*_t$: $k_1 = q_\theta(a^*_t|s_t) [\log p(a^*_t|s_t) - \log q_\theta(a^*_t|s_t)]$.
   • Logic:
     • If $k_1 > 0$ (Student underestimates the expert): Apply a Forward KL weight ($p(a^*_t|s_t) + k_1$) to reinforce the token.
     • If $k_1 \le 0$ (Student overestimates): Mask the Forward KL component to prevent conflicting gradients, relying only on the negative signal.

2. Student-Sampled Token Processing (On-policy):

   • The student samples a non-expert token $a_t \sim q_\theta(\cdot|s_t)$ where $a_t \neq a^*_t$.
   • Computes the $k'_1$ gap for this sampled token.
   • Logic:
     • If $k'_1 < 0$ (Student overestimates a non-expert token): Apply a negative weight ($k'_1$) to suppress this unreasonable behavior.
     • Crucially, if the student overestimates a non-expert token ($k'_1 < 0$) and underestimates the expert ($k_1 > 0$), the weight for the expert token is doubled ($2p(a^*_t|s_t) + k_1$). This explicitly redirects the probability mass suppressed from the bad sample toward the correct expert action.

3. Loss Function:
   The final objective combines these weighted terms:
   $$\mathcal{L}_{HPD} = \mathbb{E} [ -w^*_t \log q_\theta(a^*_t | s_t) - w_t \log q_\theta(a_t | s_t) ]$$
   This approach retains the computational efficiency of one-hot supervision (SFT) while incorporating the dense distributional guidance of KL divergence.

3. Key Contributions

1. Unified View: Established a reweighted log-likelihood framework that connects SFT, FKLD, and RKLD, clarifying how divergence direction and data regimes interact.
2. HPD Algorithm: Proposed a novel method that hybridizes Forward and Reverse KL signals. It uses negative KL estimators as reward signals to suppress bad student samples and dynamically reinforce expert tokens, effectively balancing mode coverage and mode-seeking.
3. Efficiency: Unlike pure On-Policy Distillation (OPD) which requires expensive full-sequence rollouts, HPD uses lightweight, single-token sampling from the student, significantly reducing computational overhead.
4. Strong Initialization: Demonstrated that HPD provides a superior initialization for subsequent On-Policy Distillation or Direct Preference Optimization (DPO), leading to further performance gains.

4. Experimental Results

The authors validated HPD across three domains: Math Reasoning (Long-generation), Dialogue (Short-generation), and Code Generation.

• Math Reasoning (Off-policy):

  • Tested on Qwen2.5 and LLaMA3 families (e.g., 7B Teacher $\to$ 1.5B/3B Student).
  • Results: HPD significantly outperformed SFT, SeqKD, RKLD, and JSD. Notably, HPD improved Qwen2.5 3B reasoning capabilities by 41% and LLaMA3 3B by 78% compared to standard baselines.
  • Dynamics: HPD maintained stable training entropy (avoiding collapse) and progressively reduced the KL divergence gap between student and teacher, unlike SFT which stagnated.

• Dialogue & Code (Off-policy):

  • Dialogue: HPD achieved the highest scores on multi-turn dialogue benchmarks (MT-Bench) and alignment metrics (AlpacaEval, Arena-Hard), showing superior coherence and context retention.
  • Code: HPD achieved the best average pass@1 scores on HumanEval and MBPP, demonstrating more stable improvements across model families compared to the high variance seen in KD and JSD.

• On-Policy & DPO Integration:

  • When used as an initialization for OPD, HPD + OPD outperformed SFT + OPD baselines.
  • In DPO alignment tasks, models initialized with HPD showed larger performance gains ($\Delta$) after DPO compared to SFT or RKLD initializations, proving HPD creates a better "strong-to-weak" foundation.

5. Significance

• Bridging the Gap: HPD resolves the long-standing trade-off between the stability of SFT/Forward KL and the alignment precision of Reverse KL/On-Policy methods.
• Scalability: By avoiding full rollouts and using a lightweight token-level sampling strategy, HPD makes high-quality distillation feasible for resource-constrained environments.
• Generalizability: The method is model-agnostic and effective across diverse tasks (reasoning, dialogue, coding) and model scales, suggesting it is a robust paradigm for the next generation of efficient LLM deployment.
• Future Impact: The paper suggests HPD could be applied earlier in the training pipeline (e.g., mid-training or pre-training) to further enhance LLM efficiency.