Entropy-Aware On-Policy Distillation of Language Models

TL;DR

提出Entropy-Aware On-Policy Distillation,有效平衡模式追求与多样性,提升数学推理任务表现。

cs.LG 🔴 高级 2026-03-07 42 次浏览
Woogyeol Jin Taywon Min Yongjin Yang Dennis Wei Yi Zhou Swanand Ravindra Kadhe Nathalie Baracaldo Kimin Lee
知识蒸馏 强化学习 自然语言处理 模型压缩 不确定性

核心发现

方法论

该方法结合反向KL和正向KL损失,通过动态调节在低高熵区域的应用,利用算法如PPO优化,增强模型在高不确定性状态下的输出多样性。具体实现包括对教师模型的输出熵进行判断,选择性地引入正向KL,确保学生模型既能高效学习,又能保持输出的多样性。实验中采用六个数学推理基准,验证了该方法在Pass@8指标上的提升,平均提升达+1.37至+5.05不等。该框架在保持训练效率的同时,有效缓解了逆KL的模式偏向问题,增强了模型的表达能力。

关键结果

  • 在六个数学推理任务中,Qwen3-0.6B模型通过EOPD实现Pass@8提升+1.37,Qwen3-1.7B模型提升+2.39,Qwen3-4B模型提升+5.05,显著优于传统逆KL蒸馏方法。
  • 在高熵区域,EOPD保持了更高的token级别熵,增强了输出的多样性,减少了模式崩溃现象。
  • 在多个外域任务中,EOPD模型展现出优越的泛化能力,超越基线方法,尤其在复杂推理和指令遵循任务中表现突出。

研究意义

该研究突破了传统逆KL偏向模式的局限,通过引入熵感知机制,有效保持了模型输出的多样性和不确定性,推动了知识蒸馏在高复杂度推理任务中的应用。其在数学推理、外域泛化等方面的提升,为大模型压缩与高效推理提供了新思路,具有重要的理论与实践价值。

技术贡献

创新点在于提出动态调节逆正KL结合策略的熵感知机制,利用教师模型的局部不确定性信息,优化学生模型的输出分布。该方法在保持训练效率的基础上,有效缓解了逆KL的模式偏向问题,提升了模型的输出多样性与鲁棒性。技术实现包括基于PPO的策略优化框架,结合教师输出熵阈值调节正向KL的引入,提出了高效的采样与计算策略,显著改善了传统蒸馏方法的局限。

新颖性

本研究首次系统性将教师模型的输出不确定性融入到强化学习式的在策略蒸馏中,通过动态调节逆向与正向KL的结合,解决高熵状态下输出多样性不足的问题。相较于现有的静态蒸馏策略,该方法实现了在保持训练效率的同时,显著提升输出多样性与推理能力,具有较强的创新性。

局限性

  • 该方法依赖于教师模型的熵阈值设定,可能在不同任务或模型规模下需要调优,存在一定的参数敏感性。
  • 在极端高熵状态下,正向KL的引入可能导致训练不稳定或收敛缓慢,需进一步优化调节策略。
  • 计算成本虽优于纯正向KL,但在大规模模型中仍存在一定的资源消耗,未来需优化采样与计算效率。

未来方向

未来可探索自适应熵阈值机制,结合多模态信息提升不确定性估计的准确性。此外,扩展至多任务、多模态场景,验证其在实际应用中的鲁棒性与泛化能力,推动模型压缩与高效推理的深度融合。

AI 总览摘要

在自然语言处理领域,知识蒸馏作为模型压缩与性能提升的重要手段,面临输出多样性不足的挑战。传统方法多采用逆KL损失,偏向于模式追求,导致生成结果趋于单一,难以反映教师模型的真实不确定性。本文提出Entropy-Aware On-Policy Distillation(EOPD),通过动态调节逆向与正向KL的结合,利用教师模型的输出熵信息,有效平衡了模型的学习效率与输出多样性。

具体而言,EOPD在低熵区域采用逆KL,保证训练的稳定性与快速收敛;在高熵区域引入正向KL,增强模型对多样性和不确定性的捕获。这一机制通过策略优化算法(如PPO)实现,结合教师输出的局部熵阈值进行调节。实验结果显示,在六个数学推理基准上,EOPD显著优于传统逆KL蒸馏,Pass@8指标提升达+1.37至+5.05不等,特别在复杂推理和外域任务中表现出更强的泛化能力。

该研究突破了模型蒸馏中输出多样性不足的瓶颈,为大规模语言模型的高效压缩与推理提供了新思路。未来,结合多模态信息和自适应调节机制,有望进一步提升模型的鲁棒性与应用范围,为智能系统的高效推理和知识传递奠定基础。

深度分析

研究背景

近年来,随着大规模预训练语言模型(如GPT、BERT)的崛起,模型压缩与知识蒸馏成为研究热点。传统蒸馏方法多基于离策略训练,利用软标签进行监督,提升模型效率。近年来,基于强化学习的在策略蒸馏(OPD)逐渐兴起,采用逆KL作为优化目标,强调模式追求,提升训练速度。然而,这种方法在高不确定性状态下输出多样性不足,限制了模型推理能力的提升。已有研究如Lu & Lab (2025)和Yang et al. (2025)展示了OPD的潜力,但仍存在输出单一、模式崩溃的问题。为解决这一瓶颈,本文提出结合教师输出不确定性信息的熵感知机制,旨在在保持训练效率的同时,增强输出的多样性和鲁棒性。

核心问题

现有的在策略蒸馏方法主要依赖逆KL损失,导致在高熵状态下输出趋于模式化,缺乏多样性,限制了模型在复杂推理任务中的表现。逆KL的模式偏向性使得学生模型难以捕获教师模型的多模态输出,尤其在关键决策点,模型容易过早收敛,丧失不确定性信息,影响推理的全面性和鲁棒性。此外,逆KL在高熵区域引起训练不稳定,表现为梯度震荡和收敛缓慢,严重制约了其应用范围。解决这一问题,需设计一种机制,既能保持训练效率,又能充分传递教师模型的局部不确定性。

核心创新

本研究的创新点在于提出熵感知机制,将正向KL引入逆KL的训练框架中,动态调节在不同不确定性区域的损失权重。具体包括:

  • �� 设计基于教师输出熵的阈值调节策略,自动识别高不确定性状态;
  • �� 在低熵区域采用逆KL,确保训练稳定快速;
  • �� 在高熵区域引入正向KL,增强输出多样性,避免模式崩溃;
  • �� 利用PPO策略优化框架实现动态调节,兼顾效率与多样性。此机制突破了传统静态蒸馏的局限,显著提升了模型在复杂推理任务中的表现。

方法详解

  • �� 以教师模型输出的条件概率分布为基础,计算每个token的熵值;
  • �� 设定熵阈值τ,根据熵值判断是否引入正向KL;
  • �� 在低熵区域,优化逆KL(即反向KL),保证训练速度和稳定性;
  • �� 在高熵区域,加入正向KL,传递教师的不确定性信息;
  • �� 采用PPO策略优化,利用采样的轨迹和教师的输出概率,动态调节损失函数;
  • �� 结合采样策略,限制在教师的前k个候选词中,提升计算效率;
  • �� 通过多轮训练,模型逐步学习到既能快速收敛,又能保持输出多样性的能力。

实验设计

采用六个数学推理基准(如MATH500、AIME24/25)验证方法效果。训练使用Qwen3-8B作为教师,学生模型规模从0.6B到4B不等。对比基线包括传统蒸馏、OPD和GRPO,指标为Pass@8和Avg@8。调节参数包括熵阈值τ和正向KL系数α。通过消融实验验证不同参数设置对输出多样性和推理性能的影响。结果显示,EOPD在保持训练稳定的同时,显著提升了推理准确率,尤其在高难度任务中表现优异。

结果分析

在六个数学推理任务中,EOPD模型的Pass@8平均提升达+1.37至+5.05,显著优于逆KL蒸馏。模型在高熵区域保持更高的token级别熵,减少模式崩溃,增强输出多样性。外域任务中,模型展现出更强的泛化能力,特别是在复杂推理和指令遵循方面。Pass@k性能随k值增加而差距扩大,表明EOPD更好地探索多样推理路径,提升了模型的解决问题能力。

应用场景

该方法可用于大规模语言模型的高效压缩,适合在资源有限环境中部署高性能推理系统。特别适合数学推理、复杂决策支持和多模态任务,能显著提升模型的推理多样性和鲁棒性。未来可结合多任务学习和多模态信息,推动智能系统在教育、科研和工业中的应用。

局限与展望

该方法依赖于教师模型的熵阈值设定,调节参数需在不同任务间调整,存在一定的调优成本。在极端高熵状态下,正向KL引入可能引起训练不稳定,需进一步优化调节策略。计算成本虽优于纯正向KL,但在超大模型中仍有资源消耗,未来需探索更高效的采样与优化方案。

通俗解读 非专业人士也能看懂

想象你在教一个学生画画。老师(教师模型)画出一幅画,里面有很多细节和可能的变化。你(学生)试图模仿老师,但只专注于最明显的部分,忽略了那些有多种可能的细节。为了让学生画得更像老师,老师告诉他哪些部分可以有多种变化(高不确定性),哪些部分可以直接模仿(低不确定性)。这样,学生既能快速学会基本技巧,又能理解复杂的变化。这个方法就像在训练AI模型一样,根据不同的情况调整学习策略,让它既快又能表现出丰富的内容。

简单解释 像给14岁少年讲一样

想象你在玩一个游戏,游戏里有很多不同的路径可以走。有时候,你知道下一步怎么走(低不确定性),可以快速做决定;但有时候,路上有很多可能的选择(高不确定性),你需要考虑更多可能性。以前的AI学习方法就像只看最明显的路,忽略了其他可能的路径,结果变得很单一。现在,这个新方法就像是让AI在知道路很清楚时快点走,在路不清楚时多考虑一些可能的选择。这样,AI就能学得更聪明,走的路也更丰富,能应对各种复杂的情况。它就像是你在玩冒险游戏,既能快速完成任务,又能探索更多隐藏的秘密!

原文摘要

On-policy distillation is a promising approach for transferring knowledge between language models, where a student learns from dense token-level signals along its own trajectories. This framework typically uses reverse KL divergence, encouraging the student to match the teacher's high-confidence predictions. However, we show that the mode-seeking property of reverse KL reduces generation diversity and yields unstable learning signals when the teacher distribution has high entropy. To address this, we introduce Entropy-Aware On-Policy Distillation. Our key idea is augmenting the standard reverse KL objective with forward KL when teacher entropy is high, capturing the full range of plausible outputs while retaining precise imitation elsewhere. It balances mode-seeking precision with mode-covering robustness without sacrificing on-policy training efficiency. Experiments show that our method maintains generation diversity (sustained token-level entropy) and improves student-teacher alignment (lower forward KL on high-entropy tokens). Across six math reasoning benchmarks, this yields Pass@8 accuracy gains of +1.37 for Qwen3-0.6B-Base, +2.39 for Qwen3-1.7B-Base, and +5.05 for Qwen3-4B-Base compared to baseline on-policy distillation methods. These results demonstrate that accounting for teacher uncertainty is essential for maintaining diversity and achieving effective knowledge transfer.

cs.LG cs.CL