ARKD: Adaptive Reinforcement Learning-Guided Bidirectional KL Divergence Distillation for Text Generation

2026-06-29

AI Summary

This work addresses the limitations of conventional knowledge distillation, which relies solely on unidirectional KL divergence and struggles to simultaneously model dominant and long-tail probability distributions, thereby constraining generation quality and generalization. To overcome this, the authors propose a reinforcement learning–based adaptive bidirectional KL distillation framework that, for the first time, incorporates a policy network into the KL weighting mechanism. This policy dynamically blends forward and reverse KL divergences based on the evolving teacher–student distribution alignment and leverages immediate reward signals to jointly optimize both dominant-mode fidelity and long-tail coverage. Evaluated across multiple benchmarks, the method consistently outperforms existing approaches, achieving gains of 0.4–0.6 points over greedy heuristic baselines in Rouge-L and BertScore metrics.

Abstract

Knowledge distillation (KD) is a key technique for compressing Large Language Models (LLMs), yet methods relying on a single KL objective often fail to balance primary distribution fitting with long-tail probability modeling, limiting both generation quality and generalization. To address this, we analyze the complementary roles of forward and reverse KL divergence (FKL/RKL) in distribution alignment from theoretical and empirical perspectives. We then propose a reinforcement-learning-based adaptive KL-weighted distillation framework, in which a policy network dynamically assigns weights to FKL and RKL based on teacher-student distributional characteristics, guided by immediate reward signals to achieve dual alignment on principal and long-tail modes. Extensive experiments demonstrate consistent improvements across Rouge-L and BertScore metrics, surpassing greedy heuristics by 0.4-0.6 points and outperforming other baseline methods on diverse benchmarks.

Problem

• Knowledge Distillation
• KL Divergence
• Text Generation
• Large Language Models
• Distribution Alignment

Innovation

• Adaptive KL Divergence
• Reinforcement Learning
• Knowledge Distillation
• Bidirectional KL
• Text Generation