🤖 machine learning

Trust-Region Behavior Blending for On-Policy Distillation

本文提出了信任区域行为混合（Trust-Region Behavior Blending, TRB），这是一种用于在策略蒸馏中的预热方法，通过在 KL 信任区域内将学生策略与教师策略进行混合，随后逐渐退火回纯学生采样，从而在训练早期稳定过程，进而提升在数学推理任务中的表现。

原作者： Daniil Plyusov, Alexey Gorbatovski, Alexey Malakhov, Nikita Balagansky, Boris Shaposhnikov, Daria Korotyshova, Daniil Gavrilov

发布于 2026-06-01

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Daniil Plyusov, Alexey Gorbatovski, Alexey Malakhov, Nikita Balagansky, Boris Shaposhnikov, Daria Korotyshova, Daniil Gavrilov

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

大局观：教导学生如何思考

想象一下，你正在试图通过让一名年轻学生（学生 AI）学习一位博学教授（教师 AI）来教导他们如何解决复杂的数学问题。

在旧的方法中（离线蒸馏/Offline Distillation），你会给学生一本写满了教授完美答案的教科书。学生记住了这些答案。但问题在于：当学生参加真正的考试时，他们必须从头开始生成自己的答案。由于他们从未练习过生成自己的步骤，他们在初期会感到困惑并犯错。

为了解决这个问题，研究人员开发了在策略蒸馏（On-Policy Distillation, OPD）。与其使用教科书，不如让学生逐步生成自己的答案，并由教授进行实时纠正。这种方式更好，因为它匹配了实际测试的条件。

然而，OPD 有一个缺陷： 在最开始阶段，学生非常弱。如果你让他们立即生成自己的答案，他们可能会写出一段极其糟糕、毫无意义的第一句话。如果教授试图去纠正一个糟糕的句子，这就像是在试图修理一座还没盖好的房子。此时的“信号”太弱，根本无法发挥作用。

解决方案：置信区域行为混合（Trust-Region Behavior Blending, TRB）

作者提出了一种名为 置信区域行为混合（TRB） 的新方法。你可以将其理解为一种**“带有智能向导的辅助轮”**教学法。

1. “置信区域”（安全气泡）

想象学生正在田野中行走。**学生策略（Student Policy）**是学生自然想走的路径。**教师策略（Teacher Policy）**是教授会走的路径。

如果学生偏离航线太远，教授的建议就失去了意义。TRB 创建了一个**“置信区域”**——即围绕学生当前路径的一个安全气泡。

规则： 允许学生向教授的路径迈进一小步，但不能离自己的自然风格太远。
目标： 找到一个既尽可能接近教授路径，又仍保持在学生安全气泡内的“最接近版本”。

2. “混合”（平滑融合）

TRB 并没有强迫学生完美复制教授（这太难了），也没有让他们漫无目的地游荡（这太乱了），而是将两者进行混合。

它创造了一条“混合型”路径，这条路径看起来大部分像学生，但又包含了足够的教授智慧，使学生能够保持在一条可以被学习的轨道上。
这就像舞蹈教练指导初学者：“脚往这里迈（像我一样），但保持重心在那里（像你一样）。”

3. “预热”（撤掉辅助轮）

TRB 最重要的部分是它是暂时的。

初期阶段： “置信区域”较宽。学生会得到教授的大量帮助，以确保最初的几步是高质量的。
淡出过程： 随着训练的进行，“置信区域”逐渐缩小。教授开始退后。
结束阶段： 最终，该区域完全消失。学生现在可以独立行走，但他们在起步阶段已经学会了正确的习惯。

为什么这种方法效果更好（实验结果）

论文在数学推理任务（如解决代数或几何问题）上测试了这种方法，并使用了不同规模的 AI 模型。

对比对象： 他们将 TRB 与以下方法进行了比较：
- Vanilla OPD（基础 OPD）： 让学生立即独自游荡。
- SKD： 让教授偶尔接管并强迫学生说出特定的词汇。
- SFT Warmup（监督微调预热）： 在开始之前进行一段标准的监督学习。
结果： TRB 平均表现胜出。
- 它帮助学生在初始阶段拥有比“Vanilla OPD”更高质量的步骤。
- 它不依赖于教授完全接管（像 SKD 那样），因为完全接管有时会让学生对自己的身份产生困惑。
- 它比仅仅“预热”温度或进行一段短时间的标准课程效果更好。

权衡（代价）

这其中存在一个小代价。因为 TRB 要求教授在训练初期与学生同时处于“在线”状态（同时进行思考和解码），所以它在运行训练时需要更多的计算能力和时间。然而，由于这只发生在短期的“预热”阶段，这种额外的成本是暂时的，而最终的结果是一个更聪明的学生。

总结类比

旧方法（离线）： 给一个从未去过的城市的学生一张地图。他们背下了街道，但在自己开车时却迷路了。
OPD（在策略）： 让学生开车，同时有一个副驾驶进行纠正。但如果学生开始把车开进湖里，副驾驶的纠正也就毫无意义了。
TRB： 副驾驶在最初的几分钟内轻轻引导方向盘，以确保车辆保持在道路上（在置信区域内），从而确保学生学会正确的驾驶“感觉”。一旦学生稳住了，副驾驶就会放手，让学生能够完美地独立驾驶。

技术摘要：用于在策略蒸馏的置信区域行为混合 (Trust-Region Behavior Blending)

问题陈述

在策略蒸馏 (On-Policy Distillation, OPD) 旨在解决离线知识蒸馏中固有的前缀不匹配问题，即通过在从其当前策略采样的前缀上训练学生模型，同时匹配一个更强的教师模型。然而，这种方法在训练初期引入了关键的脆弱性：较弱的学生模型会生成低质量的前缀。当学生的轨迹分布较差时，教师的监督被应用于低质量或“不可用”的前设，从而限制了学习信号的有效性。

以往试图稳定这一过程的方法通常涉及将采集过程转为离策略（例如注入教师 token 或改变目标分布），但这会破坏目标训练分布。相反，纯学生采样虽然保留了分布，但受限于较差的初始质量。核心挑战在于如何在不永久改变采集过程为离策略、也不改变基本的前缀逐一蒸馏目标的前提下，提高早期采样（rollout）的质量。

方法论：置信区域行为混合 (TRB)

作者提出了 置信区域行为理解 (Trust-Region Behavior Blending, TRB)，这是一种旨在控制早期采样阶段行为策略的预热（warmup）方法。TRB 的原理是寻找一个尽可能接近教师、但严格约束在以当前学生策略为中心的“置信区域”内的行为策略。

1. 逐前缀优化
在每个生成前缀 $h$ 处，TRB 通过求解一个约束优化问题来定义行为策略 $\mu^*(\cdot | h)$ 。其目标是在相对于学生策略 $\pi_S$ 的 KL 散度约束下，最小化与教师策略 $\pi_T$ 的 KL 散度：

$\mu^*(\cdot | h) = \arg \min_{\mu} D_{KL}(\mu \parallel \pi_T) \quad \text{s.t.} \quad D_{KL}(\mu \parallel \pi_S) \le \epsilon$

这里， $\epsilon$ 代表允许相对于学生的局部偏差。这种公式确保了采样分布被拉向教师支持的 token，但仍被限制在以学生为中心的置信区域内。

2. 闭式解
该优化产生了一个闭式解，其中行为策略是学生策略与教师策略之间的几何插值：
$\mu_\beta(a | h) = \frac{\pi_S(a | h)^{1-\beta} \pi_T(a | h)^\beta}{Z_\beta(h)}$
系数 $\beta^*(h)$ 是满足 KL 约束的最大可行值，且 $\beta^* \in [0, 1]$ 。如果教师已经在置信区域内， $\mu^*$ 将变为教师；如果 $\epsilon=0$ ，则退化为学生。

3. 退火预热
TRB 仅应用于初始预热阶段。KL 预算 $\epsilon$ 在时长 $K$ 内线性退火至零：
$\epsilon_k = \epsilon_0 \left(1 - \frac{k}{K}\right)$
这使得在训练初期可以通过强力的教师引导来提高前缀质量，同时在学生改进后逐渐回归到纯学生采样。至关重要的是，逐前缀的反向 KL 蒸馏损失保持不变；改变的仅是数据采集策略。

核心贡献

一种新颖的预热策略： TRB 引入了一种通过在显式的以学生为中心的 KL 约束下优化教师引导策略来控制早期采样行为的方法，而不是简单地注入 token 或改变损失函数。
保留在策略特性： 与持久混合教师和学生策略或改变目标分布的方法不同，TRB 确保训练最终会返回到纯在策略采样，从而保留了主训练阶段的目标分布。
理论公式化： 论文为平衡教师引导与学生稳定性的行为策略提供了一个闭式解，并由 KL 散度关于插值系数的单调性提供了理论依据。

实验结果

作者在两个数学推理蒸馏设置中评估了 TRB：

从 Qwen3-8B 蒸馏 Qwen3-1.7B-Base。
从 Qwen3-4B 蒸馏 Qwen3-0.6B-Base。

性能表现：

TRB 在两种设置中均取得了最强的平均性能，优于 vanilla OPD、固定 $\epsilon$ 混合、Veto、SKD（交错教师注入）、温度预热以及 SFT 预热。
在 Qwen3-1.7B 设置中，TRB 实现了 33.2% 的平均 pass@1，而 vanilla OPD 为 32.3%。
在 Qwen3-0.6B 设置中，TRB 实现了 44.4%，而 vanilla OPD 为 44.0%。

分析：

早期训练动态： 轨迹分析（图 2）显示，虽然某些方法（如 SKD）初始上升更快，但它们不一定能维持最高的最终性能。TRB 提供了一个“快速启动”，从而带来了更优的最终结果。
熵与前缀质量： 在预热阶段，TRB 降低了已访问前缀上教师 token 的平均熵，表明了更高质量的引导。Step-0 探测（图 4）表明，与 vanilla OPD 前缀相比，在 TRB 下生成的序列在由教师或学生继续生成时，具有更高的成功率，这表明 TRB 成功地将早期状态分布转向了更具可学习性的轨迹。

意义与主张

论文声称 TRB 通过在不永久偏离在策略目标的前提下提高初始数据分布的质量，有效地解决了在策略蒸馏早期的脆弱性问题。

范围有限： 作者明确指出，其研究范围限定在两个使用 Qwen3-Base 对及基于正确性的评估协议的数学推理设置中。他们并未声称特定的预热调度可以原封不动地迁移到其他领域或教师-学生差距较大的场景。
成本权衡： 论文承认，由于需要在线教师解码和师生共存（co-residency），TRB 在预热阶段增加了训练时间成本。然而，他们指出这些成本是暂时的，因为训练在预热期结束后会回归到标准的、更高效的 OPD 运行模式。
核心洞察： 结果表明，当学生访问的前缀与教师显著不匹配时，教师引导的离策略行为最为有益。一旦这种不匹配得到解决，持续的离策略引导可能会带来收益递减，这支持了将行为策略退火回学生的策略。