Entropy-Aware On-Policy Distillation of Language Models
Entropy-Aware On-Policy Distillation balances mode-seeking and diversity, boosting math reasoning accuracy by +1.37 to +5.05 on benchmarks.
Key Findings
Methodology
The approach combines reverse KL and forward KL losses, dynamically adjusting their application based on token-level entropy estimated from teacher outputs. Using algorithms like PPO, it selectively applies reverse KL in low-entropy regions for efficiency and introduces forward KL in high-entropy regions to preserve output diversity. The method involves computing teacher output entropy, setting a threshold τ, and using a policy optimization framework to balance the two divergences. Experiments across six math reasoning datasets demonstrate that this dynamic mixture improves student-teacher alignment and maintains token-level entropy, leading to significant performance gains over baseline distillation methods.
Key Results
- On six math reasoning benchmarks, Qwen3-0.6B with EOPD achieved a Pass@8 increase of +1.37, Qwen3-1.7B improved by +2.39, and Qwen3-4B by +5.05, outperforming traditional reverse KL distillation.
- EOPD retains higher token-level entropy in high-uncertainty regions, reducing mode collapse and enhancing output diversity.
- In out-of-domain tasks, EOPD models show better generalization, especially in complex reasoning and instruction-following tasks, surpassing baseline methods.
Significance
This work addresses a core limitation in knowledge distillation—loss of output diversity due to mode-seeking objectives—by integrating uncertainty-aware mechanisms. It enhances the capacity of compressed language models to perform complex reasoning and generalize beyond training data. The method offers a practical solution for deploying efficient yet robust models in real-world applications, bridging the gap between model size and performance in tasks requiring nuanced understanding.
Technical Contribution
The main technical innovation is the entropy-adaptive combination of reverse and forward KL divergences within a reinforcement learning framework, guided by teacher output entropy. This allows the student model to learn efficiently in confident regions while capturing uncertainty in ambiguous contexts. The approach leverages PPO for stable policy updates and top-k sampling for computational efficiency, providing a new paradigm for uncertainty-aware distillation that balances exploration and exploitation. It extends existing distillation frameworks with a dynamic, data-driven adjustment mechanism, offering theoretical and empirical improvements.
Novelty
This is the first systematic integration of teacher output uncertainty into on-policy distillation via a dynamic, entropy-aware KL mixture. Unlike prior static approaches, it adaptively modulates the divergence type based on local uncertainty, effectively preserving the teacher’s multimodal distribution. This innovation significantly improves output diversity and reasoning robustness, setting a new standard in model compression techniques.
Limitations
- The method relies on setting an appropriate entropy threshold τ, which may require tuning for different tasks or models, potentially limiting out-of-the-box applicability.
- In extremely high-uncertainty scenarios, the positive impact of forward KL may cause training instability or slow convergence, necessitating further refinement.
- Although more efficient than pure forward KL, the approach still incurs additional computational overhead, especially in large-scale models, which could be optimized further.
Future Work
Future research could focus on adaptive thresholding mechanisms for entropy, multi-modal extensions, and applying this framework to other tasks like dialogue or multimodal reasoning. Additionally, exploring more scalable sampling strategies and theoretical guarantees for convergence and stability could further enhance practical deployment.
AI Executive Summary
In the rapidly evolving field of natural language processing, knowledge distillation has become a vital technique for compressing large language models (LLMs) into smaller, more deployable versions. Traditional distillation methods primarily rely on off-policy training with forward KL divergence, which tends to produce models with limited output diversity, especially problematic in tasks requiring nuanced reasoning or multiple plausible outputs. Recent advances introduced on-policy distillation, which trains models by generating samples on their own trajectories and matching teacher distributions via reverse KL. While efficient, this approach suffers from a mode-seeking bias, causing the student to ignore the teacher’s inherent uncertainty and reducing diversity.
This paper introduces Entropy-Aware On-Policy Distillation (EOPD), a novel framework that dynamically combines reverse KL and forward KL losses based on the teacher’s output entropy. The core idea is to leverage the stability and efficiency of reverse KL in low-uncertainty regions, while incorporating forward KL in high-uncertainty areas to preserve the teacher’s multimodal distribution. This adaptive mixture is achieved through a reinforcement learning approach, specifically using PPO, with a threshold parameter τ to determine when to switch between the two divergences. The method also employs top-k sampling to improve computational efficiency.
Experimental results across six mathematical reasoning benchmarks demonstrate that EOPD significantly outperforms baseline methods, with Pass@8 improvements ranging from +1.37 to +5.05. Notably, the models trained with EOPD maintain higher token-level entropy in ambiguous regions, effectively capturing the teacher’s uncertainty and avoiding premature mode collapse. These improvements translate into better generalization on out-of-domain tasks and more robust reasoning capabilities.
The significance of this work lies in its ability to balance training stability, efficiency, and output diversity, addressing a fundamental challenge in knowledge distillation. By explicitly modeling teacher uncertainty, EOPD offers a pathway toward more expressive and reliable compressed language models. Future directions include refining adaptive thresholds, extending to multimodal tasks, and further optimizing computational costs, promising a broad impact on AI model deployment and reasoning performance.
Deep Analysis
Background
近年来，随着大规模预训练语言模型（如GPT、BERT）的广泛应用，模型压缩与知识蒸馏成为研究焦点。传统蒸馏方法多基于离策略训练，利用软标签提升效率，但在复杂推理任务中表现出输出多样性不足的问题。近年来，基于强化学习的在策略蒸馏（OPD）逐步兴起，采用逆KL作为目标，强调模式追求，提升训练速度。然而，这种方法在高不确定性状态下容易导致输出趋于单一，限制了模型的推理能力。为解决这一瓶颈，本文提出结合教师输出不确定性信息的熵感知机制，旨在在保证训练效率的同时，增强输出的多样性和鲁棒性。
Core Problem
现有逆KL驱动的在策略蒸馏在高熵状态下表现出输出模式单一、崩溃风险高，难以有效传递教师模型的多模态信息。这导致模型在复杂推理任务中表现受限，尤其在关键决策点，模型容易过早收敛，丧失多样性和不确定性信息。传统方法缺乏对教师输出不确定性的动态调节机制，无法兼顾训练效率与输出丰富性。因此，亟需一种机制，既能保证训练稳定，又能充分传递教师的多模态信息。
Innovation
本研究创新在于提出熵感知机制，将正向KL引入逆KL框架中，根据教师模型输出的局部熵动态调节两者的权重。具体包括：
- �� 设计基于教师输出熵的阈值τ，自动识别高不确定性区域；
- �� 在低熵区域采用逆KL，确保训练稳定快速；
- �� 在高熵区域引入正向KL，增强输出多样性，避免模式崩溃；
- �� 利用PPO策略实现动态调节，结合采样的轨迹和教师输出概率，优化模型参数。这一机制突破了静态蒸馏的局限，有效提升模型在复杂推理中的表现。
Methodology
- �� 计算教师模型每个token的条件熵，作为不确定性指标；
- �� 设定熵阈值τ，判断当前状态是否引入正向KL；
- �� 在低熵状态下，优化逆KL，快速收敛；
- �� 在高熵状态下，加入正向KL，传递不确定性信息；
- �� 采用PPO策略，利用轨迹采样和教师输出概率，动态调节损失函数权重；
- �� 通过top-k采样限制候选词集合，提高计算效率；
- �� 多轮训练中，模型逐步学习在不同不确定性状态下的输出分布，兼顾效率和多样性。
Experiments
采用六个数学推理基准（如MATH500、AIME24/25）验证方法效果。训练使用Qwen3-8B作为教师，学生模型从0.6B到4B不等。对比基线包括传统蒸馏、OPD和GRPO，指标为Pass@8和Avg@8。调节参数包括熵阈值τ和正向KL系数α。通过消融实验验证不同参数对输出多样性和推理性能的影响。结果显示，EOPD在保持训练稳定的同时，显著提升推理准确率，尤其在复杂推理任务中表现优异。
Results
在六个数学推理任务中，EOPD模型的Pass@8平均提升达+1.37至+5.05，显著优于逆KL蒸馏。模型在高熵区域保持更高的token级别熵，减少模式崩溃，增强输出多样性。外域任务中，模型展现出更强的泛化能力，特别是在复杂推理和指令遵循方面。Pass@k性能随k值增加而差距扩大，表明EOPD更好地探索多样推理路径，提升了模型的解决问题能力。
Applications
该方法适用于大规模语言模型的高效压缩，特别在数学推理、复杂决策和多模态任务中，能显著提升模型的推理多样性和鲁棒性。未来可结合多任务学习和多模态信息，推动智能系统在教育、科研和工业中的应用，提升模型的泛化能力和实际表现。
Limitations & Outlook
该方法依赖于教师模型的熵阈值设定，调节参数需在不同任务间调优，存在一定的调节成本。在极端高熵状态下，正向KL引入可能引起训练不稳定或收敛缓慢，需进一步优化调节策略。虽然效率优于纯正向KL，但在大规模模型中仍有资源消耗，未来需探索更高效的采样与优化策略。
Plain Language Accessible to non-experts
想象你在教一个学生画画。老师（教师模型）画出一幅画，里面有很多细节和可能的变化。你（学生）试图模仿老师，但只专注于最明显的部分，忽略了那些有多种可能的细节。为了让学生画得更像老师，老师告诉他哪些部分可以有多种变化（高不确定性），哪些部分可以直接模仿（低不确定性）。这样，学生既能快速学会基本技巧，又能理解复杂的变化。这个方法就像在训练AI模型一样，根据不同的情况调整学习策略，让它既快又能表现出丰富的内容。
ELI14 Explained like you're 14
想象你在玩一个游戏，游戏里有很多不同的路径可以走。有时候，你知道下一步怎么走（低不确定性），可以快速做决定；但有时候，路上有很多可能的选择（高不确定性），你需要考虑更多可能性。以前的AI学习方法就像只看最明显的路，忽略了其他可能的路径，结果变得很单一。现在，这个新方法就像是让AI在知道路很清楚时快点走，在路不清楚时多考虑一些可能的选择。这样，AI就能学得更聪明，走的路也更丰富，能应对各种复杂的情况。它就像是你在玩冒险游戏，既能快速完成任务，又能探索更多隐藏的秘密！
Abstract
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.