Hybrid Policy Distillation for LLMs

📄 arXiv: 2604.20244v1 📥 PDF

作者: Wenhong Zhu, Ruobing Xie, Rui Wang, Pengfei Liu

分类: cs.CL, cs.AI

发布日期: 2026-04-22

备注: WIP

🔗 代码/项目: GITHUB


💡 一句话要点

提出混合策略蒸馏(HPD)方法,提升LLM在长文本生成和代码任务上的性能。

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 知识蒸馏 大型语言模型 混合策略 模式覆盖 模式搜索 长文本生成 代码生成

📋 核心要点

  1. 现有知识蒸馏方法在发散方向、优化策略和数据机制的选择上存在相互制约,缺乏统一视角。
  2. 提出混合策略蒸馏(HPD),结合前向和反向KL散度的优势,平衡模式覆盖和模式搜索,并结合离线和在线策略数据。
  3. 实验表明,HPD在长文本数学推理、对话和代码生成任务上,提升了优化稳定性、计算效率和最终性能。

📝 摘要(中文)

知识蒸馏(KD)是压缩大型语言模型(LLM)的有效范例,其有效性取决于发散方向、优化策略和数据机制的相互交织的选择。本文分解了现有KD方法的设计,并提出了一个统一的视角,建立了它们之间的联系,将KD重新定义为token级别的重新加权的对数似然目标。此外,本文提出了混合策略蒸馏(HPD),它集成了前向和反向KL的互补优势,以平衡模式覆盖和模式搜索,并将离策略数据与轻量级的近似在线策略采样相结合。在长文本生成的数学推理以及短文本生成的对话和代码任务上验证了HPD,证明了跨不同模型系列和规模的优化稳定性、计算效率和最终性能的提高。

🔬 方法详解

问题定义:现有知识蒸馏方法在压缩大型语言模型时,面临着发散方向(如前向KL和反向KL)、优化策略和数据机制选择上的复杂权衡。不同的选择会导致不同的性能表现,且缺乏一个统一的框架来理解和连接这些方法。现有方法可能在模式覆盖或模式搜索方面存在不足,影响最终模型的生成质量和泛化能力。

核心思路:本文的核心思路是将知识蒸馏重新定义为token级别的重新加权的对数似然目标,从而建立不同知识蒸馏方法之间的联系。通过结合前向KL和反向KL的优势,HPD旨在平衡模式覆盖(避免生成过于集中的结果)和模式搜索(确保生成结果的多样性)。同时,利用离策略数据和轻量级的在线策略采样,提高训练效率和稳定性。

技术框架:HPD的整体框架包括以下几个主要步骤:1) 使用大型教师模型生成离策略数据;2) 使用轻量级的学生模型进行在线策略采样,生成近似的在线数据;3) 将知识蒸馏目标重新定义为token级别的重新加权的对数似然目标,其中权重由前向KL和反向KL的混合策略决定;4) 使用优化算法(如Adam)训练学生模型。

关键创新:HPD的关键创新在于:1) 提出了一个统一的知识蒸馏视角,将不同的方法联系起来;2) 结合了前向KL和反向KL的优势,平衡了模式覆盖和模式搜索;3) 采用了离策略数据和轻量级的在线策略采样,提高了训练效率。与现有方法的本质区别在于,HPD不是简单地选择前向KL或反向KL,而是根据具体任务和模型特性,动态地调整两者的权重。

关键设计:HPD的关键设计包括:1) 前向KL和反向KL的权重比例,需要根据具体任务进行调整;2) 在线策略采样的频率和数量,需要在计算效率和数据质量之间进行权衡;3) 损失函数的设计,需要确保学生模型能够有效地学习教师模型的知识,同时避免过拟合。

📊 实验亮点

实验结果表明,HPD在长文本数学推理、对话和代码生成任务上均取得了显著的性能提升。例如,在数学推理任务上,HPD优于传统的知识蒸馏方法,并在对话和代码生成任务上展现出更高的生成质量和效率。具体性能数据请参考论文原文。

🎯 应用场景

该研究成果可应用于各种需要压缩大型语言模型的场景,例如移动设备上的自然语言处理、低资源环境下的模型部署、以及需要快速推理的应用。通过知识蒸馏,可以将大型模型的知识迁移到小型模型,从而降低计算成本和延迟,提高用户体验。此外,该方法还可以用于提升模型的鲁棒性和泛化能力。

📄 摘要(原文)

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.