Step Rejection Fine-Tuning: A Practical Distillation Recipe
作者: Igor Slinko, Ilia Zavidnyi, Egor Bogomolov, Yaroslav Zharov
分类: cs.LG, cs.AI, cs.CL, cs.SE
发布日期: 2026-05-11
💡 一句话要点
提出步级拒绝微调(SRFT)方法,通过细粒度损失掩码提升LLM智能体在复杂任务中的表现
🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)
关键词: 大语言模型 智能体训练 拒绝微调 损失掩码 代码生成 知识蒸馏
📋 核心要点
- 现有拒绝微调(RFT)方法直接丢弃失败轨迹,导致模型无法从包含部分正确逻辑的复杂任务数据中学习,造成了宝贵训练样本的浪费。
- SRFT通过引入评论员模型对轨迹中的每一步进行评估,识别并掩码错误步骤的损失,使模型能够保留上下文并学习从错误中恢复的策略。
- 实验表明,SRFT在SWE-bench Verified基准上实现了32.2%的解决率,相较于传统RFT方法,性能提升幅度显著,证明了细粒度反馈的有效性。
📝 摘要(中文)
拒绝微调(RFT)是训练大语言模型(LLM)智能体的标准方法,其核心在于从训练集中剔除失败的轨迹。在SWE-bench任务中,这意味着过滤掉无法通过测试的补丁提交。然而,这种方法丢弃了大量未解决的轨迹,尽管这些轨迹在复杂任务中占比很高,且往往包含部分正确的推理步骤。本文提出了步级拒绝微调(SRFT),这是一种利用这些未解决轨迹的实用方法。我们利用一个评论员(Critic)LLM来评估轨迹中每一步的正确性。在训练过程中,我们对错误步骤的损失进行掩码处理,同时将其保留在上下文窗口中。这种方式确保模型学习从错误中恢复,而不是重复错误。在SWE-bench Verified上的评估显示,RFT通过排除未解决轨迹将解决率提高了2.4%,而SRFT通过过滤而非完全丢弃这些轨迹,将解决率提升了3.7%,总解决率达到32.2%。
🔬 方法详解
问题定义:在LLM智能体训练中,RFT方法简单粗暴地剔除失败轨迹,忽略了这些轨迹中可能存在的正确中间步骤,导致模型在处理复杂任务时难以从错误中学习并进行自我修正。
核心思路:SRFT的核心思想是将“轨迹级”的过滤细化为“步级”的过滤。通过保留包含错误步骤的轨迹,但仅对错误步骤的损失进行掩码(Masking),模型既能利用上下文信息,又不会被错误的动作引导,从而学习到更鲁棒的纠错能力。
技术框架:该方法包含两个主要阶段:首先,利用一个经过训练的评论员(Critic)LLM对模型生成的轨迹进行逐步评估,标记出错误步骤;其次,在微调阶段,将这些轨迹输入模型,通过损失函数掩码机制,仅对被标记为正确的步骤计算梯度,错误步骤仅作为上下文输入。
关键创新:SRFT的本质创新在于将“拒绝”从样本维度下沉到“步骤”维度。这种方法打破了“全盘否定”的局限,实现了对训练数据的最大化利用,使模型能够区分哪些步骤是有效的,哪些是导致失败的诱因。
关键设计:关键技术细节在于损失函数的动态掩码机制。在训练过程中,模型接收完整的轨迹序列,但通过掩码向量(Mask Vector)将错误步骤的Loss权重设为0,确保模型在反向传播时仅优化正确的推理路径,从而实现高效的知识蒸馏。
📊 实验亮点
在SWE-bench Verified数据集上,SRFT方法表现优异。相比于传统RFT方法2.4%的提升,SRFT实现了3.7%的解决率增长,最终达到32.2%的解决率。实验证明,通过细粒度损失掩码保留未完全成功轨迹,能有效提升模型在复杂编程任务中的推理与纠错能力。
🎯 应用场景
该方法主要应用于需要多步推理的复杂任务场景,如软件工程自动化(代码生成与修复)、复杂逻辑规划及科学计算。通过提升模型从错误中恢复的能力,SRFT能显著增强智能体在长序列任务中的可靠性,具有极高的工业应用价值。
📄 摘要(原文)
Rejection Fine-Tuning (RFT) is a standard method for training LLM agents, where unsuccessful trajectories are discarded from the training set. In the context of SWE-bench tasks, this corresponds to filtering out runs where the submitted patch does not pass the tests. However, this approach discards unresolved trajectories, even though they form a large portion of all trajectories for hard tasks and even then may be partially correct. In this work, we propose Step Rejection Fine-Tuning (SRFT) - a practical way to leverage these unresolved trajectories. For this, we employ a critic LLM to assess the correctness of each step in a trajectory. Consequently, during training, we mask the loss for erroneous steps while retaining them in the context window. This way we ensure the model learns to recover from errors without reproducing them. Evaluation on SWE-bench Verified shows that while RFT improves the resolution rate by 2.4% by excluding unresolved trajectories, SRFT improves it by 3.7% by filtering them instead of discarding completely, reaching the total resolution rate of 32.2%.