基于信任区域的行为融合用于策略蒸馏 Trust-Region Behavior Blending for On-Policy Distillation
通过KL信任区域引导早期rollout行为，提升策略蒸馏效果
前置知识
On-policy Distillation（OPD）
OPD是一种从学生策略本身采样的前缀上进行知识蒸馏的方法，而不是使用固定的离线数据集或教师生成的前缀。它解决了传统离线蒸馏的前缀不匹配问题，因为在推理时学生会生成自己的前缀。OPD使用反向KL散度作为目标函数：\(\mathcal{L}_{OPD}(\theta) = \mathbb{E}_{h \sim P_{\pi_S}}[D_{KL}(\pi_\theta(\cdot | h) \parallel \pi_T(\cdot | h))]\)，其中\(\pi_S\)是学生策略，\(\pi_T\)是教师策略。
本文的核心改进就是在OPD框架上进行早期训练的rollout控制，理解OPD是理解本文的基础。
KL散度（Kullback-Leibler Divergence）
KL散度是衡量两个概率分布之间差异的指标。在本文中有两种KL散度：反向KL \(D_{KL}(\mu \parallel \pi_T)\) 衡量行为策略到教师的距离，正向KL \(D_{KL}(\mu \parallel \pi_S)\) 衡量行为策略到学生的距离。KL散度的数学定义为 \(D_{KL}(p \parallel q) = \sum_x p(x) \log \frac{p(x)}{q(x)}\)。反向KL用于生成任务因为它鼓励模型支持的高概率区域，而正向KL则更倾向于覆盖分布。
本文的核心约束就是基于KL散度构建的信任区域，理解KL散度的方向性和约束机制是理解TRB方法的关键。
信任区域方法（Trust Region Method）
信任区域方法是优化中的一种约束策略，它限制每次更新步长在一个局部邻域内，以防止过度偏离当前策略。在本文中，信任区域定义为 \(D_{KL}(\mu \parallel \pi_S) \leq \varepsilon\)，即行为策略不能偏离学生策略太远。这个约束确保了rollout分布仍然保持学生的自主性，同时允许一定的教师引导。通过调节预算 \(\varepsilon\)，可以控制教师介入的程度。
TRB的名称直接来源于此概念，信任区域是方法的核心技术手段，理解它有助于明白为什么TRB能够在早期训练中获得更高质量的前缀。
行为策略（Behavior Policy）
在强化学习和策略蒸馏中，行为策略是用于生成数据的采样策略，与学生策略分离。学生策略是通过梯度更新优化的目标策略，而行为策略只负责产生训练样本。本文中，\(\mu\) 表示行为策略，\(\pi_S\) 表示学生策略。在标准OPD中，行为策略等于学生策略（\(\mu = \pi_S\)），但在TRB中，\(\mu\) 是一个在学生信任区域内最接近教师的混合策略。
TRB的核心创新就是改变了行为策略的定义，而不是改变目标函数。理解行为策略与学生策略的分离是理解TRB的关键。
研究动机
现有的On-policy Distillation（OPD）方法面临一个核心困境：在训练早期，学生模型的能力还很弱，生成的前缀质量很差，导致教师的监督信号作用在低质量的前缀上，浪费了宝贵的教师指导资源。具体来说，当学生从自己当前的策略rollout时，如果轨迹质量很差，教师在这些轨迹上提供的指导信号就缺乏价值。这就像一个新手在练习错误的招式，高手的纠错也难以发挥效果。作者引用Li et al. (2026)的分析指出，OPD的成功取决于学生访问的状态是否携带兼容且可转移的教师信号。在数学推理等复杂任务中，这个问题尤为突出，因为推理链条一旦偏离，后续的纠正就变得非常困难。
本文的目标是本文的目标是设计一种方法，在OPD训练的早期阶段，通过控制rollout的行为策略来提高前缀质量，同时不破坏OPD的核心优势即在推理分布上进行训练。具体来说，作者希望：在不改变per-prefix的反向KL蒸馏目标的前提下，通过调整行为策略，在早期训练中提供更多的教师引导；随着学生能力的提升，逐渐减少教师介入，最终回到纯学生rollout；在两个数学推理蒸馏设置中验证方法的有效性，超越vanilla OPD和其他基线方法。
与已有工作不同的是，本文的独特切入角度在于聚焦于行为策略这个控制层面，而不是修改蒸馏目标函数本身。与Veto（Jang et al., 2026）改变目标分布、Entropy-Aware OPD（Jin et al., 2026）改变散度度量、TIP（Xu et al., 2026）改变监督位置、SKD（Xu et al., 2025）直接替换token等方法不同，TRB在访问前缀之前就改变了前缀的分布来源。这相当于从源头改善数据质量，而不是在数据生成后进行补救。作者特别强调，TRB只在早期训练中发挥作用，并通过KL预算的衰减机制自动退化为标准OPD，这与其他持续介入的方法形成鲜明对比。
核心方法
TRB的核心思想是在学生策略周围建立一个KL信任区域，在这个区域内选择最接近教师策略的行为策略来生成rollout前缀。直觉上，这就像在新手学习时，高手在旁边实时指导，但又不完全取代新手的动作新手仍然在做自己的决策，但在必要时会得到高手的提示。技术路线上，TRB首先定义了一个约束优化问题：在保持行为策略不超过学生策略KL距离\(\varepsilon\)的约束下，最小化行为策略到教师的KL距离。这个问题有闭式解，形式上是学生策略和教师策略的加权几何平均，权重系数\(\beta\)通过二分搜索确定。其次，TRB引入了warmup机制，KL预算\(\varepsilon\)从初始值\(\varepsilon_0\)线性衰减到0，在K步warmup后完全退化为标准OPD。这种方法既利用了教师指导，又保持了on-policy训练的核心优势。
TRB的核心创新点是将学生策略视为中心，在以学生为中心的信任区域内进行教师引导，而不是直接用教师替换学生行为。具体来说，目标函数是 \(\mu^*(\cdot | h) = \arg\min_\mu D_{KL}(\mu \parallel \pi_T)\)，约束条件是 \(D_{KL}(\mu \parallel \pi_S) \leq \varepsilon\)。这与SKD等方法有本质区别：SKD是当学生采样token不在教师top-K时直接用教师采样替换，这是一种硬替换，而TRB是一种软引导，通过概率分布的平滑融合来保持学生轨迹分布的特性。另一个关键区别是TRB只在warmup阶段激活，而其他方法可能持续介入。这种设计基于一个关键洞察：当学生的轨迹分布仍然与教师不匹配时，TRB最有价值；一旦学生的前缀质量提升，持续的教师指导可能会变得不再有用，甚至可能干扰学习过程。
方法步骤详情
TRB方法包含三个主要步骤。第一步是构建per-prefix行为策略。对于每个前缀h，计算学生策略\(\pi_S(a | h)\)和教师策略\(\pi_T(a | h)\)，然后求解约束优化问题：\(\mu^*(\cdot | h) = \arg\min_\mu D_{KL}(\mu \parallel \pi_T)\)，s.t. \(D_{KL}(\mu \parallel \pi_S) \leq \varepsilon\)，以及概率归一化约束。这个问题的闭式解是 \(\mu_\beta(a | h) = \frac{\pi_S(a | h)^{1-\beta}\pi_T(a | h)^\beta}{Z_\beta(h)}\)，其中\(Z_\beta(h)\)是归一化常数，\(\beta \in [0,1]\)是融合系数。第二步是通过二分搜索确定\(\beta\)。由于\(D_{KL}(\mu_\beta \parallel \pi_S)\)是\(\beta\)的单调非减函数，可以在\([0,1]\)区间内二分搜索找到最大的可行\(\beta^*\)，使得约束条件满足。如果\(\varepsilon = 0\)，则\(\beta^* = 0\)，即\(\mu^* = \pi_S\)；如果教师策略本身就满足约束，则\(\beta^* = 1\)，即\(\mu^* = \pi_T\)。第三步是KL预算的衰减调度。对于warmup长度K，设置\(\varepsilon_k = \varepsilon_0(1 - k/K)\)，\(k \leq K\)。在第k步rollout生成时，使用预算\(\varepsilon_k\)计算\(\beta^*\)，然后从行为策略\(\mu_{\beta^*}\)采样。一旦warmup结束，\(\varepsilon = 0\)，行为策略完全回归到学生策略。
技术新颖性
TRB的技术新颖性体现在三个方面。首先，它明确分离了行为策略（用于数据生成）和学生策略（用于目标优化），并提出了在约束下优化行为策略的框架。这种分离使得在不改变蒸馏目标的前提下控制rollout质量成为可能。其次，TRB提供了信任区域约束的闭式解和高效的求解方法（二分搜索），这使得方法在实际训练中可以高效实现。论文附录E详细推导了这个闭式解，并证明了\(D_{KL}(\mu_\beta \parallel \pi_S)\)的单调性，保证了二分搜索的正确性。第三，TRB引入了预算衰减机制，将warmup与后续训练平滑连接。这种设计基于对训练动态的深刻理解：早期需要更多指导，后期应让学生自主探索。附录F的数学分析表明，在小预算情况下，trust-region方法有良好的局部权衡：稍微偏离学生可以一阶地减少到教师的KL距离，而只需支付二阶的行为KL成本。这意味着早期的教师引导具有高边际价值。
实验结果
本文在两个数学推理蒸馏设置中进行了系统实验：Qwen3-1.7B-Base蒸馏自Qwen3-8B，以及Qwen3-0.6B-Base蒸馏自Qwen3-4B。主要发现如下：首先，TRB在两个设置中都取得了最好的平均得分。在Qwen3-1.7B设置中，TRB的平均得分为33.2%，超过vanilla OPD的32.3%，提升了0.9个百分点；在Qwen3-0.6B设置中，TRB的平均得分为44.4%，超过vanilla OPD的44.0%，提升了0.4个百分点。其次，TRB在大多数单项基准上也表现优异。在Qwen3-1.7B设置中，TRB在MATH500上达到69.7%（vanilla OPD为69.1%），在Olympiad上达到34.3%（vanilla OPD为33.7%），在AIME24上达到44.8%（vanilla OPD为43.0%），在AIME25上达到10.2%（vanilla OPD为8.8%）。在Qwen3-0.6B设置中，TRB在GSM8K上达到53.6%（vanilla OPD为53.1%），在MATH500上达到70.1%（vanilla OPD为69.9%），在Olympiad上达到22.3%（vanilla OPD为21.8%），在AMC上达到31.6%（vanilla OPD为31.1%）。第三，TRB的warmup效果优于其他基线。Figure 2显示，TRB在早期训练中上升速度快于vanilla OPD，但最终的稳定性能更好。相比之下，SFT warmup和fixed-ε blending虽然初期上升快，但最终性能不如TRB。SKD在某些配置下可以接近TRB，但在整体平均上仍然落后。第四，诊断实验表明TRB确实改善了前缀质量。Figure 3显示，在warmup期间，TRB访问的前缀具有更低的教师token-mean熵，说明这些前缀更接近教师的偏好。Figure 4的前缀截断实验显示，在step 0时，TRB前缀相对于vanilla-OPD前缀的相对成功率增益在64到512的截断长度范围内达到+9%到+26%之间，说明TRB在早期确实生成了更高质量的前缀。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 数学推理（Qwen3-1.7B←Qwen3-8B） | 平均pass@1（MATH500 + Olympiad + AMC + AIME24 + AIME25） | 33.2% | 32.3% (Vanilla OPD) | +0.9个百分点 |
| 数学推理（Qwen3-0.6B←Qwen3-4B） | 平均pass@1（GSM8K + MATH500 + Olympiad + AMC） | 44.4% | 44.0% (Vanilla OPD) | +0.4个百分点 |
| MATH500数据集（Qwen3-1.7B设置） | pass@1 | 69.7% | 69.1% (Vanilla OPD) | +0.6个百分点 |
| AIME24数学竞赛（Qwen3-1.7B设置） | pass@1 | 44.8% | 43.0% (Vanilla OPD) | +1.8个百分点 |
| GSM8K（Qwen3-0.6B设置） | pass@1 | 53.6% | 53.1% (Vanilla OPD) | +0.5个百分点 |
局限与改进
作者承认的局限性包括：第一，研究范围局限于两个数学推理OPD设置和Qwen3师生对，不能保证同样的warmup schedule可以直接迁移到其他领域或师生差距设置。这意味着TRB在其他任务（如代码生成、对话系统）上的有效性还需要验证。第二，TRB在warmup期间增加了训练时间成本，因为它需要在线教师解码和师生共驻。虽然总的教师FLOP数量与vanilla OPD相当，但batched teacher pass在wall clock时间上可能比TRB的在线教师解码更快。在本文研究的环境中，这些成本是临时的，因为TRB只在warmup期间使用，训练后会回到普通OPD的运行配置。第三，TRB增加了峰值生成时间内存，因为需要同时保持学生和教师的KV cache。从内存角度看，标准OPD在时间t需要的峰值内存约为\(M_{OPD}^{gen}(t) \approx W_S + KV_S(n_t)\)，而TRB需要\(M_{blend}^{gen}(t) \approx W_S + W_Q + KV_S(n_t) + KV_Q(n_t)\)，其中\(W_S\)和\(W_Q\)分别是学生和教师的权重，\(KV_S(n_t)\)和\(KV_Q(n_t)\)是对应的KV缓存。额外的开销主要是教师权重和教师KV缓存。此外，我观察到的一些额外限制包括：方法引入了两个超参数（初始KL预算\(\varepsilon_0\)和warmup长度K），虽然论文进行了sweep，但在实际应用中可能需要针对不同任务进行调整。TRB的有效性依赖于教师和学生策略之间的差距，如果差距太大或太小，可能需要调整策略。最后，论文的实验主要集中在pass@1指标上，对于其他指标（如生成质量、鲁棒性）的评估较少。
独立分析的弱点
从独立分析的角度，TRB存在几个可以改进的弱点。首先，超参数敏感性是一个潜在问题。论文对\(\varepsilon_0 \in \{0.001, 0.005, 0.01, 0.02, 0.05\}\)和\(K \in \{15, 25, 50\}\)进行了sweep，但不同任务和模型规模的最优配置可能不同。改进方向可以是设计自适应的预算调度机制，根据学生和教师的KL距离动态调整\(\varepsilon\)，或者使用元学习来自动确定最优warmup schedule。其次，TRB目前只考虑了token-level的信任区域约束，序列级别的约束是通过token-level约束的累积自然实现的（如附录G证明的\(D_{KL}(P_\mu \parallel P_S) \leq T\varepsilon\)）。但可能存在更高效的序列级控制策略，例如根据整体轨迹质量而非局部token质量来调整预算。改进方向可以是在轨迹层面设计更精细的信任区域，考虑轨迹的语义连贯性而非仅仅是概率分布的KL距离。第三，TRB的teacher guidance是完全基于当前token的local信息，没有考虑global上下文。在复杂推理任务中，有时需要在early阶段大胆偏离，在later阶段精细调整。改进方向可以是设计层次化的信任区域，在不同层次的子目标上应用不同程度的教师引导。第四，TRB目前主要用于数学推理任务，其在开放域对话、创意写作等任务上的效果未知。这些任务的评估标准更加主观，可能需要调整信任区域的定义。改进方向可以是探索不同任务类型下的最优约束形式，例如对于创意任务可以使用更宽松的约束，对于事实任务可以使用更严格的约束。
未来方向
作者提出的未来研究方向包括：扩展到更多领域和任务类型，验证TRB在代码生成、对话系统、多轮交互等任务上的有效性；探索更复杂的warmup schedule，例如非线性衰减、基于validation指标的动态调整、多阶段的渐进式引导等；研究TRB与其他方法的结合，例如与Veto或TIP结合，同时在前缀生成和目标优化两个层面进行改进。基于本文成果，还可以延伸的研究方向包括：将信任区域方法应用于其他training regime，例如从人类反馈的强化学习（RLHF），通过控制rollout质量来提高奖励信号的可靠性；研究TRB在大规模分布式训练中的实现，特别是在多机多卡环境下如何高效地实现在线教师解码和师生策略的协调；探索TRB在持续学习和迁移学习中的应用，通过动态调整信任区域来适应不断变化的任务分布；研究TRB的理论性质，例如信任区域约束对收敛性、泛化性的影响，以及在非凸优化环境下的理论保证；开发TRB的自动化版本，通过算法选择、超参数优化等技术减少人工调参的工作量。
复现评估
从复现性角度评估，TRB的复现难度中等。论文提供了详细的实验设置和超参数配置。训练使用verl pipeline（Sheng et al., 2024）和SGLang（Zheng et al., 2023）进行rollout生成，使用FSDP2进行分布式训练。实验在8张NVIDIA H100 GPU上运行。论文在附录A中详细列出了训练配置（Table 2）、评估配置（Table 3）和SFT warmup配置（Table 4），包括优化器（AdamW，\((\beta_1, \beta_2) = (0.9, 0.999)\)）、权重衰减（0.01）、梯度裁剪（1.0）、学习率（\(1 \times 10^{-5}\)）、batch size（64）、rollouts per prompt（4）、KL loss类型（reverse KL）、KL top-k support（16 tokens）等关键参数。然而，论文没有提供源代码链接，也没有明确说明是否开源。数据方面，论文使用OpenThoughts3-1.2M语料库中的25,600个训练提示，系统提示为Please reason step by step, and put your final answer within \boxed{}。奖励分配通过math-verify库完成：正确解得1.0，错误解得0.0。评估数据集包括MATH500、AIME24、AIME25、AMC、Olympiad和GSM8K，这些都是公开的基准测试。算力需求方面，在8张H100上的训练时间没有明确说明，但考虑到模型规模（最大8B教师）和训练步数（约400步），计算成本相对较高。总体而言，如果有足够的计算资源和开源的verl/SGLang框架，TRB应该是可以复现的，但缺少源代码会增加复现难度。建议作者在未来工作中开源代码和训练脚本，以促进研究社区的复现和进一步研究。
论文图表