Technical Summary: ARKD – Adaptive Reinforcement Learning-Guided Bidirectional KL Divergence Distillation

1. Problem Statement

Knowledge Distillation (KD) is a primary technique for compressing Large Language Models (LLMs) to reduce inference costs while maintaining performance. However, traditional KD methods relying on a single Kullback-Leibler (KL) divergence objective face a fundamental trade-off in open-ended text generation:

• Forward KL (FKL): Encourages "mode-covering" behavior, where the student attempts to cover all modes of the teacher distribution. This often leads to overestimating probabilities in low-confidence, long-tail regions, degrading generation quality.
• Reverse KL (RKL): Encourages "mode-seeking" behavior, focusing on dominant modes and ignoring low-probability regions. While this prevents long-tail overestimation, it compromises diversity and coverage, potentially leading to mode collapse.

Existing static interpolation methods (e.g., fixed 0.5:0.5 weighting) fail to capture the dynamic nature of the student-teacher gap as training evolves. A static coefficient cannot adapt to the shifting requirements of distribution alignment, which necessitate different balances of FKL and RKL at different training stages.

2. Methodology: ARKD

The authors propose ARKD (Adaptive Reinforcement Learning-Guided Bidirectional KL Divergence Distillation), a framework that dynamically assigns weights to FKL and RKL using a reinforcement learning (RL) agent.

Theoretical Foundation

The paper first establishes the theoretical consistency of FKL and RKL. Analysis shows that both objectives converge to the same solution ($q_\theta = p$) when gradients are zero. This confirms that both divergences are valid for distribution matching but exhibit complementary behaviors during the optimization trajectory, justifying their joint use.

RL-Based Framework

ARKD reframes the weight assignment problem as a sequential decision-making process (Markov Decision Process):

• State ($s$): A 6-dimensional vector capturing the current training dynamics, including teacher/student entropies ($H_T, H_S$), variances ($\sigma^2_T, \sigma^2_S$), and current FKL/RKL loss values ($L_{FKL}, L_{RKL}$).
• Action ($\alpha$): A policy network ($\pi_\phi$) outputs a scalar weight $\alpha \in (0, 1)$, where $\alpha$ weights the FKL term and $(1-\alpha)$ weights the RKL term.
• Reward ($r$): Defined as the negative distillation loss ($r = -L_{distil}$) for the current batch. The objective is to maximize cumulative reward (minimize total loss) over the training trajectory.
• Optimization:
  • Student Update: Minimizes the weighted distillation loss plus a language modeling loss.
  • Policy Update: Uses the REINFORCE algorithm with an Exponential Moving Average (EMA) baseline and entropy regularization to stabilize training and prevent boundary saturation.

Training Dynamics

The policy network learns to adaptively navigate the trade-off:

1. Exploration Phase: Early in training, the policy explores a broader range of $\alpha$ (e.g., 0.3–0.6), emphasizing FKL to ensure broad coverage and prevent mode collapse.
2. Convergence Phase: As the student distribution approaches the teacher, the policy shifts toward lower $\alpha$ values (RKL-dominant) to focus on precise alignment and avoid long-tail overestimation.
3. Stability Phase: The policy stabilizes at a specific $\alpha$ (e.g., $\approx 0.18$) for final refinement.

3. Key Contributions

1. Theoretical Analysis: The paper provides a theoretical analysis establishing the complementarity and convergence consistency of FKL and RKL, forming the basis for adaptive distillation.
2. ARKD Framework: It introduces a novel RL-based joint optimization framework where a policy network dynamically weights FKL and RKL based on state-dependent characteristics, moving beyond static heuristics.
3. Empirical Validation: Extensive experiments demonstrate that ARKD consistently improves generation quality, distribution fitting, and generalization across multiple benchmarks and model scales.

4. Experimental Results

The authors evaluated ARKD on instruction tuning tasks using GPT-2 (1.5B $\to$ 120M/340M) and LLaMA (13B $\to$ 7B) configurations.

• In-Distribution (ID) Performance: On the DollyEval benchmark, ARKD achieved the highest BertScore and competitive or superior Rouge-L scores compared to baselines (SFT, SeqKD, FKL, RKL, FKL+RKL, and MiniLLM). For example, on GPT-2 1.5B $\to$ 120M, ARKD achieved a Rouge-L of 24.5 and BertScore of 0.815, outperforming the strong MiniLLM baseline.
• Out-of-Distribution (OOD) Generalization: ARKD demonstrated superior generalization on unseen datasets (SelfInst, VicunaEval, S-NI, UnNI), consistently outperforming all baselines across both model families.
• Ablation Studies:
  • vs. Static/Linear Scheduling: ARKD significantly outperformed static weighting (fixed $\alpha$) and linear scheduling strategies (time-based decay/increase), which failed to capture complex alignment dynamics.
  • vs. Greedy Heuristics: ARKD surpassed greedy heuristics (selecting the lower-loss divergence at each step) by 0.4–0.6 points. The greedy approach was found to be myopic, often degenerating to pure RKL, whereas ARKD maintained a beneficial small FKL component through continuous mixing.
  • Trajectory Analysis: The learned $\alpha$ trajectory confirmed an intelligent strategy: early emphasis on FKL for coverage, transitioning to RKL for precision, a pattern not achievable by manual design.

5. Significance and Claims

The paper claims that ARKD advances the state of adaptive knowledge distillation by addressing the limitations of static KL-based methods. By leveraging reinforcement learning to dynamically balance mode-covering and mode-seeking behaviors, ARKD achieves a better distribution alignment that enhances both generation quality and generalization.

The authors position ARKD as a principled approach that validates the necessity of dynamic, state-aware weighting over fixed heuristics. They note that while the current policy relies on a handcrafted state vector and immediate batch rewards, the framework is compatible with broader $f$-divergence families and on-policy paradigms. The work provides a solid foundation for efficient LLM compression by demonstrating that adaptive, RL-guided optimization can consistently outperform static and greedy baselines across diverse tasks and model scales.