Learning to Revise Reasoning with Segment-wise On-Policy Distillation

📄 arXiv: 2610.02703v1 📥 PDF

作者: Yuxiang Zhang, Ding Cao, Shuting Cui, Lei Wang, Weijieying Ren, Tianxiang Zhao

分类: cs.AI

发布日期: 2026-10-02


💡 一句话要点

提出段落级蒸馏以改进推理过程中的修正问题

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 在线蒸馏 推理修正 段落级蒸馏 自然语言处理 机器学习

📋 核心要点

  1. 现有的在线蒸馏方法在推理修正方面存在不足,未能有效处理学生生成的劣质推理前缀。
  2. 论文提出段落级在线蒸馏(Seg-OPD),通过重构中间推理步骤来改善后续推理的支持。
  3. 实验表明,Seg-OPD训练的模型在推理准确性上相较于基线方法有显著提升,平均提高5.22%。

📝 摘要(中文)

在本研究中,我们关注段落级的在线蒸馏(Seg-OPD),旨在通过对中间推理步骤的重构来改善后续推理的支持。现有的基于token的在线蒸馏方法未能明确提供一致的推理修正步骤,导致在学生生成劣质推理前缀时,后续教师的监督可能会强化不良推理模式。通过控制推理干预,我们发现用教师重写的段落替换学生段落能显著提高后续推理的准确性。Seg-OPD通过不确定性度量选择学生段落,并获取相应的教师重写,训练学生优先选择教师重写而非其配对的学生段落,同时保留密集的token级监督。实验结果表明,Seg-OPD训练的学生在数学推理和竞争编程任务中表现出更高的修正成功率,平均相对提升5.22%。

🔬 方法详解

问题定义:本研究旨在解决现有基于token的在线蒸馏方法在推理修正中的不足,特别是在学生生成劣质推理前缀时,教师的监督可能会强化不良推理模式。

核心思路:论文提出段落级在线蒸馏(Seg-OPD),通过对中间推理步骤进行重构,提供明确的修正步骤,从而改善后续推理的支持。该方法通过不确定性度量选择学生段落,并获取相应的教师重写,以此训练学生优先选择教师重写。

技术框架:Seg-OPD的整体架构包括选择学生段落、获取教师重写、训练学生模型三个主要模块。首先,根据不确定性度量选择需要重写的学生段落;然后,获取教师的重写版本;最后,训练学生模型以优先选择教师重写。

关键创新:Seg-OPD的主要创新在于通过段落级的重写来显著改善推理过程中的修正能力,与传统的token级蒸馏方法相比,能够更有效地处理推理中的不确定性和错误。

关键设计:在Seg-OPD中,关键参数包括不确定性度量的选择和损失函数的设计,确保学生模型在训练过程中能够有效地学习到教师的重写信息,同时保留密集的token级监督。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

实验结果显示,Seg-OPD训练的学生在数学推理和竞争编程任务中表现优异,修正成功率显著高于基线方法,平均相对提升达到5.22%。这一结果表明,段落级蒸馏在推理准确性上具有明显优势。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理中的推理任务、教育领域的智能辅导系统以及编程任务的自动化评估。通过改进推理过程中的修正能力,Seg-OPD有助于提升模型在复杂任务中的表现,未来可能对智能助手和自动化学习系统产生深远影响。

📄 摘要(原文)

On-policy distillation (OPD) improves large language model reasoning by training students on their own rollouts with dense token-wise supervision from the teacher. However, token-wise OPD does not explicitly provide a coherent alternative reasoning step showing how the student's step could be revised to improve subsequent reasoning. Furthermore, this paradigm can become less effective when the student produces a degenerate reasoning prefix, as subsequent teacher supervision remains conditioned on that prefix and may reinforce poor reasoning patterns. In this work, we focus on learning reasoning revision with segment-wise OPD to rework intermediate reasoning steps and better support subsequent reasoning. Through controlled reasoning interventions, we find that replacing student segments with teacher redrafts improves subsequent reasoning accuracy. Therefore, we address the problem of turning teacher redrafts into explicit supervision for learning to revise reasoning. We propose Segment-wise On-Policy Distillation (Seg-OPD), which selects student segments based on an uncertainty metric and obtains corresponding teacher redrafts. Seg-OPD trains the student to prefer teacher redrafts over their paired student segments while retaining dense token-wise OPD supervision. Extensive experiments on mathematical reasoning and competitive programming tasks show that Seg-OPD-trained students achieve higher revision success rates than baselines. Seg-OPD consistently outperforms the compared state-of-the-art baselines in reasoning accuracy with an average relative improvement of 5.22% across diverse models and tasks. Code is available at https://anonymous.4open.science/r/Seg-OPD.