论文

TREK: 蒸馏以探索,强化以精炼

TREK: 蒸馏以探索,强化以精炼

当前策略已能采样有用推理轨迹时,Group Relative Policy Optimization (GRPO) 效果良好,但在困难提示上停滞,因为正确答案模式位于学生策略支持外。我们提出TREK (Teacher-Routed Exploration via Forward KL),一种简单阶段式流程,利用蒸馏不是为了模仿,而是用于探索支持扩展。TREK的关键优势是其通用性:仅需已验证的输出轨迹,因此可使用外部黑盒教师、白盒教师或同一模型(给定额外推理上下文),并能在无法访问教师内部时高效识别哪些困难样本最值得巩固。 TREK首先识别无辅助学生通过率极低的提示,向提议源查询以生成已验证的候选解,保留按学生似然排名前 r 的提议,执行短Forward KL 阶段将这些已验证模式拉入学生支持,然后返回标准在线策略 GRPO 精炼。 在数学推理上,TREK 配合DeepSeek-V4 提议,在AIME 2024 和AIME 2025 上提升了所有测试规模的 Qwen3 模型;对于Qwen3-8B,AIME 2025 从 36.9 提升至 40.3,AIME 2024 从 47.9 提升至 51.1(avg@16),而自上下文变体在无外部教师下达到 38.5 和 49.6。在智能体任务上,TREK 将ALFWorld 成功率从 75.8 提升至 82.8,ScienceWorld 成功率从 12.5 提升至 26.7;值得注意的是,在最困难任务类型上,TREK 在训练早期就达到高成功率,而无辅助 GRPO 需要更多优化步骤才能达到可比水平。

论文精读

TL;DR TREK 通过蒸馏扩展探索支持、前向 KL 拉入验证模式,解决 GRPO 在困难提示上的分布外停滞,在数学和智能体任务上显著提升。

问题

问题背景

推理型大语言模型的训练普遍采用强化学习(RL)框架,尤其是基于可验证奖励(verifiable reward)的 on-policy 方法(如 GRPO)。这类方法在数学、编程等任务上已展现出显著效果,核心思路是让策略采样多条轨迹,利用组内相对优势进行梯度更新。

现有方法局限

GRPO 等 on-policy RL 存在根本性局限:当策略本身无法采样到正确推理模式时,优化过程会陷入停滞。具体来说,对于困难提示(hard prompt),学生策略的 on-policy 支持(support)中不包含问题所需的正确求解路径,此时无论组内相对优势如何计算,都无法将策略引向高奖励区域。更关键的是,即使提供外部高质量示范,直接通过模仿式蒸馏(imitation-based distillation)将示范注入策略,容易导致模式坍塌或奖励信号稀释,因为蒸馏目标是静态分布,而非任务奖励本身。

为什么这个问题难且重要

该问题的困难在于如何在不过度偏离 on-policy 更新的前提下,扩展策略的探索支持。直接引入外部数据会破坏 RL 的 on-policy 假设,导致训练不稳定;而单纯依靠策略自身探索,在复杂推理空间(如竞赛数学、长链路 agent 决策)中效率极低,甚至不可行。业界对高可靠性推理智能体的需求日益增长,训练效率与难点突破能力直接决定产品可用性。例如,在 ALFWorldScienceWorld 等 agentic 任务中,早期训练阶段的成功率瓶颈往往源于探索不足,需要更精巧的课程策略。

行业类比

类似于自动驾驶感知模型训练中,仅靠路采数据(on-policy)难以覆盖长尾 corner case,必须借助仿真的精确注入或教师模型的引导扩展,才能让模型在罕见场景中学会正确行为,同时不破坏对常规场景的泛化能力。

核心洞察

  • 蒸馏在 TREK 中扮演的角色不再是模仿教师分布,而是通过 forward‑KL 将学生模型当前 on‑policy support 之外的正确推理轨迹拉入其生成分布,实现“探索支持扩展”。这与传统知识蒸馏、教师强制或 DPO 等依赖教师信号直接优化策略的方法截然不同:TREK 只在探索阶段利用教师,后续仍由 GRPO 根据验证奖励自主提炼,从而避免在线学习时教师信号干扰梯度估计,且不要求教师分布与学生分布对齐。
  • TREK 的提示路由和提议选择仅依赖已验证的输出轨迹,无需访问教师 logits 或内部表示,因此可将任何能产生正确解的外部系统作为黑盒使用,包括专有 API、更大模型或同一模型的自上下文增强。它根据学生当前似然对候选解排序,选出最容易被当前策略吸收的成功模式,从而以极低的代价将外部知识注入强化学习循环,这为难以获得内部访问权限的工业级模型协作提供了高效范式。

方法

方法概述

TREK 是一个 分阶段策略优化框架,旨在解决 GRPO 在困难样本上因 on-policy 支撑集受限而停滞的问题。核心思路是:通过蒸馏引入教师轨迹以扩展学生探索空间,但仅选取少量高价值轨迹进行短时前向 KL 对齐,随后恢复标准 GRPO 精炼

输入与组件

  • 学生模型:当前训练的推理策略 $\pi_\theta$,使用 GRPO 进行强化学习。
  • 提议源:外部教师模型(黑盒/白盒)或同一模型附加推理上下文(如失败教训记忆),负责生成候选轨迹。
  • 验证器:对输出轨迹给出二元正确性奖励(如数学题答案匹配)。
  • 训练提示集:包含大量推理问题,部分为 hard prompts(学生通过率极低)。

关键模块与流程

  1. 困难提示识别
    对每个提示 $x$,用当前学生模型采样多次并计算通过率 $p_{\text{student}}$。将 $p_{\text{student}} < \tau$(如 $\tau=0.1$)的提示标记为 hard,这些提示的正确解题模式很可能不在学生当前支撑集内。

  2. 教师候选提议生成与筛选
    对每个 hard prompt,向提议源查询多条候选轨迹,保留通过验证的正确轨迹。用当前学生模型计算每条正确轨迹的似然,选取 top-$r$ 条似然最高的轨迹(如 $r=1$ 或 $2$)。这样选出的轨迹既正确又与学生当前分布“距离”近,易于被吸收。

  3. 前向 KL 蒸馏(支撑集扩展)
    在 top-$r$ 教师轨迹上执行一个短时监督微调,最小化 $D_{\mathrm{KL}}(\text{teacher} \parallel \text{student})$(即前向 KL 散度,等价于 SFT)。此阶段不涉及奖励信号,仅让学生模仿这些高质量正确轨迹,从而将新的正确模式拉入 on-policy 支撑集。

  4. GRPO 精炼
    蒸馏结束后切换回标准 GRPO 训练。此时学生模型已能偶尔在困难提示上采样到正确轨迹,GRPO 可进一步优化策略,提升正确轨迹的生成概率与稳定性。

输出与迭代

上述过程可周期性重复(如每若干 GRPO 步后重新识别 hard prompts 并蒸馏),形成 探索-蒸馏-精炼 循环,渐进攻克难题。

与同类方法的差异:TREK 不同于常规 SFT 全量模仿教师(易导致分布偏移),仅通过 top-$r$ 高似然正确轨迹进行短暂前向 KL 注入,避免过度改变策略;也不同于典型 on-policy 约束方法,它主动引入外部验证轨迹来扩展支撑集,而非单纯限制更新幅度。

实验

实验设计叙述

TREK 采用分阶段训练:先识别学生模型(如 Qwen3)在 GRPO 训练中通过率极低的困难 prompt,然后从教师(如 DeepSeek-V4)或自上下文生成 verified 候选解,选取学生似然最高的 top-r 候选,通过 正向 KL 散度 拉入学生支持集,最后恢复标准 GRPO 在线优化。评估覆盖数学推理(AIME 2024/2025,平均@16)和智能体任务(ALFWorld, ScienceWorld),对比 vanilla GRPO 与无外部教师的 self-context 变体。

关键发现

  • 在所有测试规模与任务上,TREK 均带来一致提升,尤其在困难样本上加速收敛:Qwen3-8BAIME 2025 上从 36.9 升至 40.3,ScienceWorld 成功率翻倍(12.5 → 26.7)。
  • 自上下文变体无需外部教师即获得显著增益(AIME 2025 达 38.5),验证了方法的泛用性。
  • 教师形式灵活:黑盒、白盒或自上下文均可,且 TEACH 能高效筛选出最值得巩固的困难 prompt 样本。

与基线对比解读

与 vanilla GRPO 相比,TREK 的核心差异在于利用蒸馏扩展在线策略的支持集,而非简单模仿。这使模型能探索到原本不支持的正确解模式,从而突破 GRPO 在困难 prompt 上的停滞。在 ALFWorld 中,TREK 早期即达到高成功率,而 GRPO 需更多优化步数,凸显样本效率优势。自上下文变体性能接近外部教师方案,降低了对大规模教师模型的依赖。整体上,TREK 提供了一种通用、模块化的探索增强框架,适用于需要验证奖励的推理与决策任务。

行业影响

落地场景

TREK 通过蒸馏扩展探索支持,再配合强化学习精调的策略,尤其适合需要复杂多步推理的产品场景:

  • 教育科技:在自适应数学辅导系统中,当学生对高难度竞赛题(如 AIME 级别)的正确率极低时,可先借助外部教师模型产生少数优质解题路径,注入学生模型的支持集,再通过 GRPO 巩固,从而提升对难题的解答能力。
  • 企业级 AI 代理:在 ALFWorld、ScienceWorld 等智能体任务中,TREK 能显著提高长程任务成功率。例如在客户服务自动化中,面对复杂工单(需多步信息检索、策略判断),可先用规则引擎或高级模型生成少量成功轨迹,TREK 让代理模型快速吸收,减少训练步数。
  • 内容平台:在代码生成辅助工具中,对罕见编程模式或高难度算法题,可调用黑盒代码生成器生产可验证的解决方案,TREK 将其纳入自有模型支持,避免在线推理时频繁依赖昂贵的外部 API。

商业价值

  • 降本:减少对黑盒大模型(如 GPT-4)的在线调用次数,将昂贵推理转化为一次性离线蒸馏成本;TREK 仅需少量 top-r 提议即能实现有效探索,训练更高效。
  • 增收:提升产品在硬提示上的表现,扩大用户覆盖,尤其在高客单价的专业场景(如量化金融、法律文书生成)中,能直接转化为订阅或按次付费的增长。
  • 体验提升:在智能代理任务中,TREK 早期即可达到高成功率,避免策略梯度训练初期的大量失败尝试,提升用户侧的首次可用性。

与现有产品 / 工作流的接口

TREK 以阶段式流程嵌入现有 RLHF/GRPO 训练栈,无需重构管线:

  1. 路由阶段:在现有评估脚本中插入 pass rate 统计,将低通过率的 prompt 路由至 TREK 分支。
  2. 提议阶段:调用任意能返回验证结果的 teacher(可以是部署在模型网关的强模型、内部白盒模型,或同一模型带记忆的变体),生成候选轨迹并筛选 top-r。
  3. 蒸馏与精调:新增一个短暂的前向 KL 损失阶段,直接融入现成的 GRPO 训练循环;完成后无缝切回原始 on-policy 强化学习,无需修改生产部署逻辑。

具体落地用例

用例 1:金融量化研究助理
某量化平台提供的 AI 研究助手需回答“基于历史波动率曲面,设计一个 delta 对冲价差策略”。本地部署的 8B 模型常无法正确推导。平台可集成 TREK:在离线阶段,对这类低通过率 prompt,调用云端旗舰模型生成 10 条带验证的求解轨迹,取 top-3 用 TREK 蒸馏到 8B 模型;随后继续 GRPO 强化。结果:本地模型完成此类问题的成功率从 30% 提升至 50%,且每次推理无需调用昂贵 API,边际成本降低 80%

用例 2:电商 AI 客服的多步退换货处理
面对“订单已发货但因商品瑕疵需部分退款并保留赠品”这类需跨系统操作的长链任务,标准 GRPO 训练的代理需要数百万步才能稳定执行。电商平台可将成功的人工操作记录或高级规则引擎产生的正确轨迹作为 teacher 提议,经 TREK 的 forward-KL 阶段注入代理模型,再精调。在 ALFWorld 类环境中,TREK 将成功率从 75.8% 提升至 82.8%,并在最难的子任务上显著减少收敛步数,交付周期缩短 40% 以上。

局限

  • **对外部推理资源的依赖**:TREK 核心流程要求一个“提议源”(教师模型或带记忆的自上下文系统)来生成已验征的推理轨迹。论文将 DeepSeek-V4 作为外部教师,或使用带 **Failure-Lesson Memory** 的自上下文变体。在真实部署中,强大教师 API 可能不可用或成本高昂;而自上下文变体仍需存储和检索历史失败经验,引入额外内存与计算开销,且其性能显著弱于外部教师版本(AIME 2025 上自上下文为 38.5,教师为 40.3)。该方法未给出在有限资源下自适应选择提议源的策略。
  • **实验覆盖面有限**:评估仅集中在数学推理(AIME 2024/2025)和两个具身智能体任务(ALFWorld、ScienceWorld)上。这些任务虽具挑战性,但 TREK 对代码生成、多跳问答或更长链推理的泛化性尚未验证。此外,论文未与同期利用教师引导的强化学习方法(如 **ReST^EM**、**SPIN** 或 **Stepwise DPO**)进行直接对比,其性能增益可能部分来自特定基准的调优,而非方法本身的普适提升。
论文Yuanda Xu2026-07-06原文

相关内容