Hybrid Policy Distillation for LLMs
Abstract
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 behaviors, 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.
Lay Summary
Large AI models can solve complex problems, but they are expensive to use. A common solution is to train a smaller model to learn from a larger and more capable one. We propose Hybrid Policy Distillation (HPD), a new training method for this process. When learning from a larger model, a small model does not make all mistakes equally: it may overlook some good decisions while being overly confident in some bad ones. HPD identifies these situations and adjusts training accordingly, encouraging the small model to pay more attention to useful behaviors and less attention to misleading ones. Experiments show that HPD helps smaller models achieve stronger reasoning abilities while remaining efficient to train and deploy. This can make advanced AI systems more affordable and accessible in real-world applications.