DRIFT:面向高效多轮优化的解耦Rollout与重要性加权微调 DRIFT: Decoupled Rollouts and Importance-Weighted Fine-Tuning for Efficient Multi-Turn Optimization
将RL目标转化为加权SFT实现高效多轮优化
前置知识
马尔可夫决策过程是强化学习问题的标准数学框架定义了状态动作转移奖励等核心要素
MDP是强化学习问题的标准数学框架,由六元组 M = (S, A, P, R, γ, T) 定义。S是状态空间,A是动作空间,P是状态转移函数,R是奖励函数,γ是折扣因子,T是最大步数。在多轮对话中,状态是完整的交互历史,动作是模型生成的token序列,转移是确定性的,即如果回答错误,会追加固定的反馈消息。MDP形式化使得我们可以用强化学习理论来分析和优化多轮交互策略。
论文将多轮纠正任务形式化为有限视界MDP,所有后续的理论分析包括最优轨迹分布、KL正则化目标都建立在这个MDP框架上。不理解MDP就无法理解论文如何将多轮对话转化为RL问题。
KL散度是衡量两个概率分布差异的重要信息论指标用于优化理论和分布匹配
KL散度衡量两个概率分布P和Q的差异,定义为 KL(P∥Q) = Σ P(x) log(P(x)/Q(x))。前向KL KL(P⋆∥Pθ) 的期望在P⋆下,优化时倾向于让Pθ覆盖整个P⋆的分布。反向KL KL(Pθ∥P⋆) 的期望在当前策略Pθ下,优化时倾向于让Pθ专注于自己的高概率区域。在论文中,标准的在线RL目标等价于最小化反向KL,而DRIFT将其替换为前向KL,从而可以固定采样分布。
KL散度是论文理论分析的核心工具。论文通过证明KL正则化RL目标等价于分布匹配问题,并用前向KL替换反向KL,实现了rollout解耦。不理解前向和反向KL的区别就无法理解DRIFT的理论创新。
重要性采样是蒙特卡罗方法中的核心重要性加权技术用于期望估计和分布转换
重要性采样是一种蒙特卡罗方法,用于估计某个分布下期望,但实际采样来自另一个分布。给定目标分布p和采样分布q,期望 E_{x∼p}[f(x)] = E_{x∼q}[w(x)f(x)],其中w(x) = p(x)/q(x)是重要性权重。在论文中,我们无法直接从最优分布p⋆采样,但从参考策略pref采样很容易,因此用重要性权重将期望从p⋆转换为pref,实现离线优化。
重要性采样是DRIFT的核心技术手段。论文通过重要性采样,将需要在最优分布p⋆下的期望转换为在参考策略pref下的加权期望,这是将在线RL转化为离线加权SFT的关键步骤。
行为崩溃是多轮训练中模型过度优化第一轮准确率导致的后续轮次表现下降问题
在多轮纠正训练中,模型过度优化第一轮准确率,导致在后续轮次中产生肤浅或不诚实的修订。例如,模型可能在第一轮给出看似合理的错误答案,收到反馈后只是简单地换个说法而不真正修正错误,或者在后续轮次中重复相同的错误。这会导致模型无法真正学习到有效的纠正策略。
行为崩溃是多轮训练中常见的问题,也是论文要解决的核心动机之一。论文通过实验展示了单轮SFT方法在非数学基准上几乎无法提升后续轮次的表现,就是因为存在行为崩溃。理解这个问题有助于评估DRIFT的优势。
研究动机
现有的多轮大语言模型优化方法面临严重的效率效果困境。在线强化学习方法如PPO和UFO能够有效学习多轮动态,但训练成本极高。每次策略更新都需要生成完整的交互轨迹,在多轮设置下rollout成本随交互长度急剧增长。论文Figure 6显示,在5轮设置下,UFO的GPU时间是SFT的4倍以上,且随着轮数增加差距会进一步扩大。另一方面,离线监督微调虽然训练高效,但存在严重的分布偏移和行为崩溃问题。Table 1数据显示,虽然单轮SFT在MATH500上从38.3%提升到51.2%,但在非数学基准如MMLU-Pro上几乎没有提升,仅1.4%,而且模型无法真正学会在负反馈下进行有效修正,后续轮次的纠错率接近0。
本文的目标是本文的目标是提出一种训练框架,在保持在线强化学习方法性能的同时,达到标准监督微调的训练效率。具体来说,需要解决三个核心挑战。第一是避免每次策略更新都生成多轮轨迹的高昂成本。第二是解决传统多轮SFT的分布偏移和行为崩溃问题。第三是实现真正的多轮纠正能力,让模型能够利用负反馈探索替代推理路径。论文希望找到一个理论上的等价变换,将复杂的在线RL目标转化为可以在离线数据上优化的简单目标。
与已有工作不同的是,本文的独特切入角度是从理论层面建立KL正则化RL目标与重要性加权监督学习之间的等价关系。此前的研究要么专注于在线RL的效率改进,要么专注于多轮SFT的数据质量改进,但没有从根本上重新审视RL目标本身的可计算形式。论文的核心洞察是KL正则化RL目标在数学上等价于在前向KL散度下的分布匹配,而前向KL允许我们将采样分布从变化的策略pθ固定为参考策略pref,从而将rollout完全解耦。这个理论洞察使得论文首次实现了RL目标加SFT效率的完美结合,这是与所有现有方法本质不同的创新。
核心方法
DRIFT框架的核心思路是通过理论变换将在线RL目标转化为离线加权SFT目标。直觉上,最优的多轮纠正行为应该是在参考行为基础上,对高回报轨迹进行指数加权。DRIFT先离线收集参考策略的多轮轨迹,根据轨迹回报计算指数权重,然后用这些权重对最终的纠正响应进行监督微调。这样,训练过程中完全不需要生成新的多轮交互,所有rollout在一次离线阶段完成,优化阶段就是标准的加权SFT。论文用Figure 2展示了整个框架的两阶段设计,第一阶段是离线轨迹生成,第二阶段是加权监督优化。
核心创新点是建立了KL正则化RL目标与重要性加权SFT的理论等价性。具体来说,论文证明了三点。第一,KL正则化RL目标的最优解是对参考策略的指数重采样,即p⋆(τ|x)等于pref(τ|x)乘以exp(R(τ)/β)再除以Z(x)。第二,最小化反向KL等价于最小化与p⋆的Reverse-KL,而最小化前向KL可以将期望从pθ固定为p⋆,在全局和局部上都与反向KL等价。第三,通过重要性采样,将p⋆下的期望转换为pref下的加权期望。这个理论链条使得RL目标可以用离线的加权SFT实现,且不需要任何近似,是完全等价的变换。
方法步骤详情
DRIFT方法分为两个阶段。第一阶段是离线轨迹生成。对每个问题提示x1,从固定参考策略πref采样K条轨迹,每条轨迹最多T轮。如果某轮回答错误,追加固定反馈f并继续。如果正确或达到T轮,终止轨迹。对每条轨迹计算回报R(τ),包含早期成功的折扣奖励和重复性惩罚。然后计算归一化权重,其中bZ是prompt级别的归一化因子。最后只保留每条轨迹的最终轮次,构建加权数据集。第二阶段是加权监督优化。最小化加权负对数似然,这个目标等价于最小化前向KL,因此理论上收敛到最优策略。在实践中,将轨迹权重w均匀应用到最终响应每个token的交叉熵损失上,用标准SFT流程训练即可。
技术新颖性
DRIFT的技术新颖性体现在多个层面。理论层面,首次系统性地建立了多轮KL正则化RL目标与加权SFT的精确等价关系,填补了RL和监督学习之间的理论空白。算法层面,提出prompt级别的权重归一化和terminal-only监督,这两个设计在理论上不是等价的实现,但论文证明了在stop-on-success协议下terminal-only有更好的bias-variance权衡。工程层面,完全解耦了rollout生成和策略优化,使得多轮训练的效率接近单轮SFT,这是所有在线RL方法无法做到的。实验层面,首次系统性地比较了不同超参数对多轮优化的影响,揭示了如何平衡收敛速度和稳定性,实验结果为未来工作提供了宝贵的设计指导。
实验结果
论文在多个基准上进行了系统实验,核心发现包括。第一,在数学推理基准上,DRIFT显著优于单轮方法,匹配或超过在线RL方法。在Qwen2.5-3B上,DRIFT-5turn在MATH达到55.9%,相比Base提升17.6%,在MATH500达到58.2%,提升18.0%,均超过UFO-5turn的55.5%和56.4%。第二,在跨域泛化基准上,DRIFT展现出强大的迁移能力。在MMLU-Pro上,DRIFT达到72.7%,相比Base提升24.8%,超过UFO的71.2%。在GPQA上达到71.5%,提升13.6%,略低于UFO的71.8%但远高于单轮SFT的57.9%。第三,DRIFT的训练效率接近SFT,远优于在线RL。Figure 6显示,在Qwen2.5-3B加A800 80G配置下,200步训练SFT耗时2.5小时,DRIFT耗时3.9小时,而UFO-5turn高达14.2小时。第四,超参数分析揭示了重要规律。较小的γ鼓励更早解决问题,最优β为0.1,过大β导致权重倾斜不够强,过小β导致权重过于集中。第五,Case study揭示了DRIFT的纠正策略。Base模型会陷入推理循环重复相同错误,而DRIFT能够利用负反馈探索替代路径。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| MATH (数学推理) | multi@5 | 55.9% | Base 38.3%, SFT 51.2%, UFO 55.5% | 相比Base提升17.6%,相比UFO提升0.4% |
| MATH500 (数学推理子集) | multi@5 | 58.2% | Base 40.2%, SFT 50.8%, UFO 56.4% | 相比Base提升18.0%,相比UFO提升1.8% |
| MMLU-Pro (多选题推理) | multi@5 | 72.7% | Base 47.9%, SFT 51.0%, UFO 71.2% | 相比Base提升24.8%,相比UFO提升1.5% |
| GPQA (研究生级科学问题) | multi@5 | 71.5% | Base 57.9%, SFT 59.5%, UFO 71.8% | 相比Base提升13.6%,相比UFO略低0.3% |
| 训练效率 (Qwen2.5-3B加A800) | GPU小时 (200步) | 3.9小时 | SFT 2.5小时, UFO 14.2小时 | 相比UFO提升3.6倍效率,接近SFT水平 |
局限与改进
论文坦诚分析了DRIFT的局限性。首先是适用范围的限制。DRIFT专为短视界、验证器指导的轻量级确定反馈设计,边界条件包括验证器提供明确的正确性信号,每次episode包含少量纠正尝试。对于基于偏好的人类反馈、开放对话目标或更长的交互规划任务,需要额外建模反馈不确定性、信用分配和探索,这些超出了当前框架。其次是离线rollout覆盖限制。DRIFT从固定参考策略采样轨迹并只在优化阶段上权重有用行为,如果参考策略没有产生有用的纠正轨迹,重要性权重无法恢复。虽然两阶段rollout-refresh可以部分缓解,平均从60.5%提升到61.2%,但增益有限。最后是知识密集型领域的根本局限。GPQA案例分析显示,虽然DRIFT学会了多轮纠正的策略行为,但纠正的准确性受限于模型的内在知识。Base模型和DRIFT都产生幻觉机制,这说明DRIFT优化的是纠正策略而非领域知识本身。
独立分析的弱点
独立分析发现DRIFT存在几个具体弱点。首先是rollout覆盖问题。在知识密集型任务或复杂推理中,单次离线rollout可能无法覆盖足够多样的纠正策略,导致重要性采样方差高。虽然论文提出了rollout-refresh变体,但增益有限,更系统的混合offline-online训练设计可能更有效。其次是terminal-only监督的潜在偏差。虽然论文在stop-on-success协议下证明了bias-variance优势,但在其他协议如固定轮数或部分奖励下可能不成立。第三是超参数敏感性。虽然DRIFT在较宽范围内稳定,但最优γ和β的设置依赖任务特性,需要调参。最后是验证器依赖。DRIFT假设验证器提供明确二元反馈,但实际应用中验证器可能有噪声、延迟或计算成本高,这限制了方法的实用性。改进方向包括设计动态rollout-refresh策略、扩展到部分奖励和奖励塑形、研究自适应超参数设置、与奖励模型结合处理验证器噪声。
未来方向
论文提出了多个未来研究方向。短期方向包括系统化rollout-refresh调度设计,例如基于策略变化幅度、权重分布变化或性能提升的触发条件。混合offline-online训练框架,用DRIFT做warmstart然后切换到onlineRL进行精细探索,Table 5显示DRIFT加UFO平均达到62.5%,超过纯UFO的60.2%。扩展到更多协议,如固定轮数、部分奖励、多阶段反馈。中期方向包括处理验证器噪声,用奖励模型或置信度加权。与领域特定预训练结合,弥补知识密集型任务的知识缺口。研究更长视界的交互规划,可能需要层次化策略或规划执行分离。长期方向包括理论上分析non-realizability情况下的近似误差、与其他RL框架的结合、从纠正任务扩展到更广泛的多轮交互。
复现评估
论文的可复现性评估如下。代码方面,论文提供了GitHub链接,声称代码开源,这对复现非常重要。数据方面,训练数据是MetaMathQA的MATH子集7500个样本,评估数据包括MATH测试集5000个、MATH500 500个、TheoremQA 800个、MMLU-Redux 5700个、MMLU-Pro 12032个、GPQA-diamond 198个,都是公开基准。算力方面,论文使用了4块NVIDIA A800 80GB或4块NVIDIA H20 96GB GPUs,全球batch size 128训练200步,这属于中等规模训练资源,一般研究实验室应该可以负担。时间方面,Qwen2.5-3B在A800上DRIFT训练约3.9小时,SFT约2.5小时,UFO约14.2小时,时间成本可接受。难度方面,论文详细描述了实验设置、prompt格式、baseline实现和超参数,且提供了详细的理论证明,整体复现难度中等。总体而言,论文的可复现性较好,只要代码质量过关,其他研究者应该能够复现主要结果。
论文图表
Figure 1展示了多轮交互的完整流程。用户给出问题,LLM生成回答。如果回答错误,用户给出轻量级反馈,如Incorrect please try again,LLM基于完整交互历史重新尝试。这个过程持续直到得到正确答案或达到最大轮数限制。图中还给出了一个具体的数学问题例子,Ian有一块40英尺长的板子,他想切成两块,长的是短的4倍,问长的是多少。模型尝试两次后才得到正确答案,展示了多轮纠正的实际过程。
这张图对理解论文至关重要,因为它直观地定义了论文要解决的核心问题场景。通过具体例子展示了轻量级负反馈下的多轮纠正流程,帮助读者建立对问题的直观理解,为后续的形式化定义和方法设计奠定基础。没有这张图,读者可能难以理解轻量级反馈和stop-on-success protocol等关键概念的实际含义。