← 返回信息流
技术博客arXiv cs.CL·2 天前

近未来引导实现策略蒸馏中的推理轨迹对齐

原标题:Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance

速览

传统策略蒸馏(OPD)依赖词元级监督,难以有效弥合学生模型与教师模型的推理轨迹差异。研究提出轨迹感知策略蒸馏(TOPD),通过近未来轨迹信息识别真实分歧状态并分配指导。实验表明,该方法在AIME24和AIME25基准上分别将准确率提升至63.3%和53.3%,效果优于标准OPD。

AI 深度解读

通过近未来引导桥接策略蒸馏中的推理轨迹

来源:arXiv cs.CL 提交日期:2026年5月29日 标题:Bridging Reasoning Trajectories in On-Policy Distillation via Near-Future Guidance

背景

在大型语言模型(LLM)的推理能力提升领域,策略蒸馏(Policy Distillation)是一种主流方法。其中,策略内蒸馏(On-Policy Distillation, OPD)因其能够利用学生模型自身生成的轨迹数据进行训练而备受关注。在 OPD 框架下,学生模型在教师模型的监督下,基于其自身策略采样生成的推理轨迹进行训练。

然而,尽管 OPD 的操作对象是完整的推理轨迹,其底层的信号传递机制却仍然停留在“词元级”(token-level)。传统的 OPD 通过识别高损失(high-loss)词元来发现偏差,并通过局部的反向 KL 散度(reverse-KL)修正来修复这些偏差。这种机制存在一个根本性的局限:它虽然采样的是轨迹,但学习的却是孤立的词元。

这种“轨迹采样但词元学习”的机制无法可靠地将学生模型的推理轨迹向教师模型的轨迹靠拢。研究表明,约 30% 的高损失词元实际上处于低发散区间,这意味着许多被标记为“错误”的词元仅仅是表面形式的失配,而非真正的推理分叉点。此外,即使面对真正发散的词元,孤立的词元级监督也难以进行有效修复,因为推理失败往往表现为短视域内的分布漂移(short-horizon distributional drift),而非单个词元的孤立错误。

核心内容

为了解决上述问题,研究团队提出了一种名为 Trajectory-aware OPD (TOPD) 的新方法。该方法的核心创新在于引入了“近未来引导”(Near-Future Guidance)机制,利用轨迹的后续信息来更准确地识别真正的发散状态,并将指导信号分布到多个未来的词元上,而非仅关注当前词元。

1. 传统 OPD 的局限性分析

  • 高损失词元的误判:在标准 OPD 中,模型倾向于将所有高损失词元视为需要修正的错误。然而,数据显示约 30% 的高损失词元属于“低发散 regime”。这些词元在语义或逻辑上与教师轨迹并无本质冲突,仅仅是表面形式上的不匹配。对这些词元进行强制修正,不仅无效,还可能干扰模型对真正关键推理步骤的学习。
  • 局部修正的失效:推理错误通常不是由单个词元引起的,而是由一系列紧密相关的词元构成的“短视域分布漂移”导致的。仅通过局部反向 KL 散度修正单个词元,无法纠正这种序列层面的结构性偏差。

2. TOPD 机制:近未来引导

TOPD 通过引入对“近未来”轨迹信息的观察,改进了偏差识别和修正过程:

  • 识别真实发散状态:通过查看当前词元之后的几个词元(即近未来),TOPD 能够判断当前的偏离是否会导致最终推理路径的根本性偏离。如果后续轨迹与教师轨迹保持一致,则当前的“高损失”可能只是噪声或表面差异,无需过度修正。
  • 多词元指导分布:一旦识别出真实的发散状态,TOPD 不再局限于修正单个词元,而是将指导信号分布到多个未来的词元上。这种方法能够更好地捕捉推理过程中的序列依赖关系,从而更有效地将学生模型的轨迹拉回教师模型的轨迹。

3. 实验结果

实验数据有力地证明了 TOPD 的有效性:

  • 基准提升:通过抑制非发散的高损失词元,标准 OPD 的平均准确率从 47.8% 提升至 48.2%。这表明仅优化损失词元的筛选机制就能带来显著增益。
  • TOPD 的显著优势:引入 TOPD 后,性能进一步提升至 52.2%。
  • 高难度测试集表现
    • AIME24(美国数学邀请赛 2024 数据集)上,准确率从 60.0% 提升至 63.3%。
    • AIME25(美国数学邀请赛 2025 数据集)上,准确率从 46.7% 大幅提升至 53.3%。

这些结果不仅展示了 TOPD 在整体性能上的优势,更证明了其在复杂推理任务(如数学解题)中的强大能力。

关键要点

  • 问题诊断:现有策略内蒸馏(OPD)虽基于轨迹采样,但仅进行词元级学习,导致约 30% 的高损失词元被误判,且无法有效处理由短视域分布漂移引起的推理失败。
  • 方法创新:提出 Trajectory-aware OPD (TOPD),利用“近未来引导”机制,通过观察后续词元来判断当前偏离的真实性,并将修正信号分布到多个未来词元。
  • 去噪机制:TOPD 能够有效抑制那些仅属于表面形式失配的非发散高损失词元,避免模型在无关紧要的细节上过度拟合。
  • 性能突破
    • 平均准确率从标准 OPD 的 47.8% 提升至 TOPD 的 52.2%。
    • 在 AIME24 上提升 3.3 个百分点(60.0% -> 63.3%)。
    • 在 AIME25 上提升 6.6 个百分点(46.7% -> 53.3%)。
  • 核心洞察:推理轨迹的修正需要从“词元级”转向“轨迹级”,关注短视域内的序列分布变化,而非孤立的词元概率。

意义与影响

这项研究对大型语言模型的推理能力训练具有重要的理论和实践意义:

  1. 重新审视蒸馏信号:它揭示了传统策略蒸馏中“轨迹”与“词元”之间的错位问题,指出单纯依赖词元级损失无法充分捕捉推理过程中的结构性偏差。这为未来设计更高效的蒸馏算法提供了新的视角。
  2. 提升推理效率:通过识别并忽略非发散的高损失词元,TOPD 减少了模型在噪声上的计算浪费,使训练过程更加聚焦于真正的推理瓶颈。这对于资源受限的模型部署尤为重要。
  3. 推动复杂推理任务进展:在 AIME 等高难度数学推理数据集上的显著提升,表明 TOPD 方法能够有效处理需要多步逻辑推导的复杂任务。这可能加速 LLM 在科学计算、代码生成和复杂决策等领域的应用。
  4. 方法论启示:引入“近未来”信息作为监督信号,是一种简单而有效的策略。这种方法可以推广到其他序列生成任务中,帮助模型更好地理解上下文依赖和长期一致性。

总之,TOPD 通过桥接推理轨迹中的断层,为策略内蒸馏提供了一种更精细、更有效的监督机制,有望成为提升 LLM 推理能力的重要技术路径。

查看原文 →arxiv.org