以下报告对研究论文《Trust Region On-Policy Distillation》进行了详细分析。
2. 这项工作如何融入更广泛的研究领域
大型推理模型(LRM)的发展推动了数学、代码生成和智能体任务等领域的进步。然而,LRM 巨大的推理成本促使人们开发资源效率更高的小型推理模型(SRM)。知识蒸馏是一种常用的技术,用于将知识从更大、更强大的“教师”模型传递到更小的“学生”模型。
历史上,生成式语言模型的知识蒸馏主要采用离策略方法,例如序列级知识蒸馏。这些方法通常训练学生模型模仿教师生成的响应。离策略蒸馏的一个已知局限性是“暴露偏差”,即学生模型接受的是教师生成的轨迹训练,但却基于自身生成的轨迹进行推理,这会导致结果不匹配,尤其是在冗长的多步骤推理过程中。
策略蒸馏(OPD)是一种通过直接在学生模型自身生成的轨迹上训练学生模型来缓解暴露偏差的方法。该方法旨在提高学生模型在技能风险管理(SRM)中的效率。尽管OPD具有潜力,但现有的OPD方法经常遇到训练不稳定的问题。当学生模型生成的轨迹与教师模型的分布显著偏离时,就会出现这种不稳定,导致监督信号不可靠,策略梯度可能出现错误,进而导致优化失败。此外,面向推理的OPD通常面临内存限制,无法进行全词汇表监督,并且需要使用KL散度估计器(例如K1估计器),这可能会进一步降低监督信号的可靠性。
以往解决这些问题的尝试包括奖励裁剪(如 REOPOLD 中的技术)和基于熵的词元过滤等方法。然而,这些方法存在一些局限性,例如可能在去除异常梯度的同时移除信息丰富的监督信息,或者无法持续提升性能。本文旨在通过解决师生分布严重不匹配情况下监督可靠性的根本问题,建立一个更稳健、更稳定的 OPD 框架,专门针对推理任务进行优化。它借鉴并试图克服先前 OPD 和语言模型知识蒸馏方法的局限性。
3. 主要目标和动机
本研究的主要目标是提高策略蒸馏(OPD)在推理任务中的可靠性和稳定性,尤其是在训练小型推理模型(SRM)方面。这一总体目标是由现有OPD方法中固有的几个具体挑战所驱动的:
-
解决分布不匹配导致的训练不稳定性: 核心动机在于观察到当学生模型和教师模型分布出现显著差异时,OPD训练会出现不稳定性。在这种情况下,学生模型生成的标记可能落入教师模型提供低置信度或不可靠监督的区域。这会导致错误的策略梯度,从而破坏优化过程,并可能导致训练崩溃。本文旨在开发一种机制来识别和管理这些不可靠的监督信号。
-
缓解策略梯度异常值: K1 估计器常用于面向推理的优化决策(OPD)中,通过避免全词汇表计算来降低内存开销,但它极易受到显著策略梯度异常值的影响。当教师为学生采样轨迹分配极低的概率时,基于 K1 的策略梯度可能变为极负值,进一步加剧优化的不稳定性。因此,一个关键目标是在不丢失有用信息的情况下抑制这些异常值。
-
提升代际学习者(SoG)的质量: 由于最优解设计(OPD)完全依赖于从学习者的策略中采样得到的轨迹,能力有限的学习者模型可能会生成低质量的响应,尤其是在处理复杂的推理问题时。这些低质量的轨迹会限制有效的优化空间,并阻碍学习者获得高质量、信息丰富的指导。本研究旨在引导学习者生成更高质量、更易于有效提炼的轨迹。
-
克服现有策略的局限性: 本文指出了当前处理不可靠监督的方法存在不足。简单的奖励裁剪虽然试图降低错误信号的影响,但可能会无意中移除有用的监督信息,从而导致性能瓶颈。同样,基于熵感知的令牌选择并未持续提升 OPD 的性能,这表明丢弃“普通”令牌可能适得其反。本文旨在开发更复杂、更有效的信用分配策略。
-
建立统一基准: 鉴于现有面向推理的 OPD 研究通常在不同的实验配置下进行,因此,建立一个统一的基准是关键目标。该基准将有助于从三个角度系统地研究监督可靠性挑战:多领域评估(数学、代码、STEM)、多样化的 OPD 策略(传统方法和最新方法)以及内存高效的 KL 估计(K1 和 top-k)。这种比较分析旨在推动更优方法的开发。
总而言之,本研究旨在开发一种稳健的策略蒸馏(OPD)框架——信任区域策略蒸馏(TrOPD),该框架明确地考虑并管理监督可靠性。通过这种方式,本研究力求实现稳定的优化,提高蒸馏出的小型推理模型的质量,并最终提升小型推理模型在各种推理任务中的能力。
4. 方法论和途径
信任域策略蒸馏(TrOPD)方法旨在解决策略蒸馏(OPD)方法因不可靠的监督信号(尤其是在推理任务中)而导致的不稳定性问题。该方法融合了信任域优化、异常值估计和离策略指导的原理。
4.1 问题界定
本文首先在语言模型蒸馏的背景下对该问题进行表述。基于策略的蒸馏(OPD)使用反向KL散度(RKL)目标函数训练“代际学生”(SoG):$D_{KL}(pi_S || pi_T) = E_{x sim pi_S}[log rac{pi_S(x)}{pi_T(x)}]$。该目标函数的梯度自然地采用策略梯度形式,奖励学生生成在教师策略下概率较高的序列。RKL具有模式搜索的特性,惩罚落入教师策略低概率区域的学生输出。对于推理模型,直接计算全词汇表RKL会带来很高的内存开销(O(n · k),其中n为序列长度,k为词汇表大小)。为了规避这个问题,最近的方法采用了 K1 估计器来无偏估计 KL 散度,从而得到目标 $J_{KD} = -E_{x sim pi_S}[log rac{pi_S}{pi_T}]$。
然而,K1 估计器存在两个瓶颈:
- 显著的策略梯度异常值: 当从 $pi_S(x)$ 中采样的轨迹的 $pi_T(x) 接近 0 时,基于 K1 的策略梯度可能会变得非常负,从而导致不稳定的优化。
- 低质量学生群体(SoG): 如果学生产生的回答质量较低,则有效优化空间受到限制,从而限制了信息丰富的监督。
4.2 门诊基线基准测试
在提出 TrOPD 之前,作者建立了一个统一的基准来评估现有的 OPD 方法,重点关注它们在内存约束下基于 K1 估计器的标记级 OPD 的有效性。
- 差异性评估: 他们考察了前 k 个词元的前向 KL (FKL) 算法、基于 K1 的 RKL 算法以及广义 JSD 目标。他们发现,由于近似存在偏差,在小词汇子集上单独使用 FKL 算法效果不佳。
- 令牌过滤和奖励裁剪: 本文评估了现有的策略,例如基于熵的令牌过滤(例如,GRPO、EOPD 中的令牌过滤)和奖励裁剪(例如,REOPOLD 中的奖励裁剪)。基准测试结果表明,考虑熵的选择通常会降低性能,而简单的奖励裁剪带来的增益有限且不稳定,并且会引入超参数调优的复杂性。这些发现表明,需要一种更精细的方法。
4.3 信任区域政策学习(TrOPD核心)
受信任区域策略优化(TRPO)的启发,TrOPD 引入了一种机制,仅在策略梯度可靠的区域进行优化。
- 自适应信任区域: 对于每个学生生成的标记 $x sim pi_S$,TrOPD 定义了一个被分类到“信任区域”的概率,$P_{trust}(x) = min( rac{pi_T(x)}{pi_S(x)}, 1)$。该概率的动机源于推测性解码,并确保优化主要发生在教师模型与学生解码一致的地方,从而使 K1 估计器可靠。
- 异常值掩蔽(初步探索): 首先,作者研究了异常值掩蔽方法,即忽略奖励幅度超过阈值的令牌。结果表明,与原始 OPD 或简单裁剪相比,该方法能够维持更高的策略熵和更低的梯度范数,从而提高稳定性。
- 异常值估计: TrOPD 并非仅仅采用掩码方法,而是引入了一个辅助的前向 KL (FKL) 目标函数来处理异常值区域。当存在显著的分布不匹配(异常值区域)时,反向 KL 算法会变得不可靠。因此,对于这些区域,TrOPD 从教师视角出发,使用 Top-k FKL 估计器计算蒸馏信号:$J^{FKL} x = -M { ar{x}} sum_{v in V^{(k)} T} pi {T(v)} log rac{pi_{T(v)}}{pi_{S(v)}}$,其中 $M_{ ar{x}}$ 是异常值掩码。这种方法既保留了信息丰富的监督信息,又不会干扰信任区域中可靠的梯度。
- 统一的策略内目标: 令牌 $x$ 的策略内目标定义为信任域的 RKL 和异常域的 FKL 的组合:$J^{On} x = -M_x log rac{pi_S}{pi_T} - M { ar{x}} sum_{v in V^{(k)} T} pi {T,v} log rac{pi_{T,v}}{pi_{S,v}}$。
4.4 非政策性信托区域指导
为了进一步指导学生模型,特别是对于学生生成的前缀质量可能较低的复杂推理任务,TrOPD 加入了非策略指导:
- 机制: 蒸馏轨迹由教师生成的前缀 ($x[:l]$) 和学生生成的后续部分 ($x[l:]$) 组成。
- 目标: 对于离策略前缀,应用前向 KL 目标函数,使用 K1 估计器:$J_x = - eta I[x sim pi_T] log rac{pi_T}{pi_S}$。这鼓励从教师的高质量轨迹中进行模仿学习。
- 退火: 使用余弦调度逐渐将最大离策略轨迹长度退火至零,从最初的教师重度指导过渡到训练结束时的完全策略学生生成。
4.5 统一优化
总体 TrOPD 目标 $J^{TrOPD}_x$ 结合了策略内信任域 RKL 项、策略内异常值 FKL 项和策略外指导 FKL 项。每个组成部分根据样本是学生在信任域内生成、学生在异常值域内生成还是教师生成(策略外指导)而应用。内存复杂度得到控制,RKL 和策略外 FKL(K1 估计器)的内存复杂度为 O(n),异常值 FKL(Top-k 估计器)的内存复杂度为 O(nk)。
4.6 实现细节
实验在两种环境下进行:
- 单域蒸馏: DeepSeek-R1-Distill-Qwen-1.5B 学生模型与 Skywork-OR1-Math-7B 教师模型进行数学推理。
- 多领域蒸馏: DeepSeek-Distilled-Qwen-1.5B 与 Skywork-OR1-7B 结合,以及 Qwen3-SFT-1.7B 与 Qwen3-Nemotron-4B 结合,涵盖数学、代码和科学领域。训练进行 200 步,采用固定学习率($5 imes 10^{-6}$),基于 FKL 的方法的 top-k 值为 64,离策略指导的 eta 值为 0.001。提示批大小为 128,每个提示进行 4 次 rollout,最大生成长度为 8096 个 token。
4.7 基准评估
模型评估依据如下:
- 数学推理: AIME 2024、AIME 2025、AMC 2023(32 次评估的平均准确率)。
- STEM 推理和指令遵循: GPQA Diamond、MMLU-Redux v2、IFBench。
- 代码生成: LiveCodeBench v6。
5. 主要发现和结果
信任区域策略蒸馏 (TrOPD) 的实验评估表明,在各种推理基准测试中,其性能优于现有的策略蒸馏 (OPD) 基线。
整体性能: TrOPD 的性能始终优于目前最先进的 OPD 基准模型,包括 OPD、EOPD 和 REOPOLD。在数学推理、代码生成、指令执行和 STEM 基准测试中(图 1),TrOPD 均表现出显著的改进:
- 数学(AIME 25):比 OPD 高 3.34 分(44.06 对 40.72)。
- Code(LiveCodeBench):比 OPD 高 4.00 分(36.00 对 32.00)。
- 指导(IFBench):比 OPD 高 5.11 分(42.18 对 37.07)。
- STEM(GPQA dia.):与 OPD 相比 +6.18 分(35.98 对 29.80)。
单领域提炼(表 3): 以 DeepSeek-R1-Distill-Qwen-1.5B 为学生,Skywork-OR1-Math-7B 为教师,与基线 OPD 相比,TrOPD 将平均数学推理表现提高了 3.06 分,将一般领域任务提高了 2.63 分。
- 与使用奖励裁剪的 REOPOLD 相比,TrOPD 在数学任务中提高了 1.99 分,在一般领域中提高了 1.84 分。
- TrOPD 的性能也显著优于基于熵的标记选择方法,例如 EOPD、Entropy OPD 和 REOPOLD 2Stage,分别提高了 2.63、3.73 和 3.74 个百分点。这些结果表明,考虑异常值的标记选择比基于熵的标记选择更有效。
多领域提炼(表 3 和表 4): TrOPD 在不同的师生配置的多领域环境中均表现出持续的性能提升:
- 当使用 Skywork-OR1-7B 对 DeepSeek-Qwen2.5-1.5B 进行提炼时,TrOPD 比 OPD 平均提高了 4.62 个百分点。
- 当使用 Qwen3-Nemotron-4B 对 Qwen3-SFT-1.7B 进行提炼时,TrOPD 相比 OPD 平均得分提高了 3.44 分。这些结果表明,TrOPD 能够推广到各种推理任务和模型架构中。
消融研究(表 5): 针对数学推理任务的消融研究揭示了 TrOPD 各组成部分的贡献:
- 将 FKL 专门应用于异常值区域(FKL 异常值)比简单地掩盖或裁剪异常值获得了更好的性能,验证了恢复信息监督的方法。
- 纳入非政策指导进一步增强了 TrOPD 变体(Mask、Clip、FKL Outlier)的性能,证实了其补充优势。
- 三种 TrOPD 变体(TrOPD Mask、TrOPD Clip、TrOPD FKL)均优于原版 OPD,其中 TrOPD FKL 的平均改进幅度最高,达到 +3.06 分。
优化稳定性(图 3): 对策略熵和梯度范数(针对 TrOPD 的“屏蔽异常值”组件)的分析表明:
- Mask Outlier 保持了较高的策略熵,表明在训练过程中更好地保留了探索能力。
- 与 OPD 和 Clip Outlier 相比,它实现了更低的梯度范数,这意味着通过减少不可靠梯度的影响,优化更加稳定。
与并行方法的比较: TrOPD 还与一种并行方法 AOPD(非对称策略蒸馏)进行了比较。TrOPD 的性能优于 AOPD。此外,将 TrOPD 与 AOPD 结合使用可进一步提高性能(平均得分从 40.63 提高到 41.67),这表明这两种方法是正交的,可以相互补充。
6. 重要性和潜在影响
信任区域策略蒸馏(TrOPD)的研究在大型语言模型和知识蒸馏领域做出了若干重要贡献,并具有潜在的影响:
-
增强 OPD 的稳定性和可靠性: TrOPD 的主要意义在于它能够解决策略蒸馏 (OPD) 中的一个根本性挑战:在师生分布严重不匹配的情况下,由于不可靠的监督而导致的训练不稳定。通过引入信任域优化和异常值估计,TrOPD 提供了一种机制来过滤或妥善处理错误的策略梯度,从而实现更稳定、更鲁棒的训练。
-
小型推理模型 (SRM) 性能提升: TrOPD 在各种数学推理、代码生成、指令跟踪和 STEM 基准测试中均展现出持续的性能提升,表明其能够有效提高小型语言模型的推理能力。这对于在资源受限的环境中部署高级推理功能至关重要,因为在这些环境中,大型语言模型 (LLM) 由于推理成本高昂而不切实际。
-
新型学分分配策略: TrOPD 提出了一种新型学分分配策略,能够自适应地确定教师对学生生成词元的监督可靠性。该策略结合了基于解码一致性的自适应信任域和针对异常区域的定向前向 KL 估计,相比以往的简单奖励裁剪或基于熵的过滤等方法,提供了一种更为精细的方案。即使在分布差异较大的区域,该方法也能保留信息丰富的监督信息,避免宝贵知识的丢失。
-
整合非策略性指导以提升质量: 引入非策略性指导(即在学生生成的路径初始阶段注入教师预设的前缀)可以解决学生生成的路径质量低下的问题。这有助于引导学生走向更连贯、更高质量的推理路径,尤其是在训练初期或处理复杂问题时,从而扩大有效的优化空间。
-
统一基准测试与问题识别: 建立面向推理的OPD统一基准测试,并系统地将监督可靠性识别为关键问题,为未来的研究奠定了宝贵的基础。它能够对不同方法进行公平的比较,并突出需要进一步方法论发展的领域。
-
更广泛的应用潜力: 虽然目前的工作侧重于推理任务,但信任域学习和自适应监督可靠性的原理可以扩展到其他形式的知识蒸馏或语言模型的强化学习。与其他 OPD 策略(例如 AOPD)的正交性表明,可以通过结合互补技术来获得更佳的性能。
-
资源高效型人工智能的进步: TrOPD 通过训练能够进行复杂推理的高质量智能资源管理器 (SRM),为实现更易于获取和更可持续地使用先进人工智能这一更广泛的目标做出了贡献。这对于边缘计算、移动应用以及计算和内存资源有限的场景都具有重要意义。
局限性部分指出,该方法缺乏实际部署,且无法与完整的预训练和训练中期阶段进行整合。未来的研究可以探讨,当TrOPD被整合到更全面的训练流程中时,如何进一步提升其在实际应用中的实用性,并提高SRM推理性能的上限。