修正强化学习中的优势估计偏差，提升学生模型蒸馏效果
OPD+: Rethinking the Advantage Design for On-Policy Distillation
- 基于f-散度构建通用优化框架，重新审视优势设计合理性
- 证明停梯度操作会导致奖励目标与梯度估计偏差
- 提出OPD+方法，在推理和工具使用任务上优于基线
在策略蒸馏（OPD）中，能力强大的教师语言模型可将知识迁移至基础学生模型，其目标可被表述为基于学生生成轨迹的强化学习形式。然而，尽管优势奖励依赖于学生模型的概率，现有方法通常采用停梯度设计以保证稳定性，这使得优势估计的有效性存疑。本文基于学生与教师之间的f-散度，提出一个通用优化框架，从数学上重新审视该设计空间的合理性。我们证明，一般的停梯度操作会导致一般散度函数下奖励目标及其对应梯度的偏差。为此，我们提出OPD+，即对原OPD的修正版本，该方法在性能上优于基线KL蒸馏，并支持多种f-散度的选择。我们在数学推理和工具使用基准上验证了本方法的有效性。
原文摘要 · Abstract (English)
On-policy distillation (OPD) is a widely used technique to transfer capabilities from capable teacher language models to the base student models, and can be formulated in a reinforcement learning style objective using student generated rollouts. Yet, despite the divergence reward being dependent on student model likelihood, existing works usually adopt a stop gradient design primarily for stability, which makes the resulting advantage estimation questionable. In this work, we provide a generic optimization framework based on f-divergence between the student and teacher, and mathematically revisit whether such design space is valid. We prove that general stop-gradient operation would lead to biased estimates of the reward objective and corresponding gradient for general divergence functions. We propose OPD+, the corrected version of OPD that demonstrates improved performance over the baseline KL approach and also supports the choice of various f-divergence. We validate our findings on mathematical reasoning and tool-use benchmarks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。