PaperlayerView exact revision on arXiv

Source-linked AI summary

DRIFT: Decoupled Rollouts and Importance-Weighted Fine-Tuning for Efficient Multi-Turn Optimization

Jian Mu, Tianyi Lin, Chengwei Qin, Zhongxiang Dai, Yao Shu

arXiv:2605.31455v1cs.LGcs.CL

TL;DR

Multi-turn optimization must handle lightweight feedback, but online RL requires costly repeated rollouts while SFT faces distribution shift. DRIFT decouples rollout from optimization through importance-weighted SFT, achieving performance comparable to or better than multi-turn RL baselines across math and general-domain benchmarks with higher training efficiency.

  • Problem

    Multi-turn models must robustly use lightweight negative feedback, while online RL training requires costly full-trajectory rollouts as interaction length grows.

  • Method

    DRIFT samples correction trajectories from a fixed reference policy, assigns return-based importance weights, and optimizes the target model through weighted SFT.

  • Results

    DRIFT achieves performance comparable to or better than strong multi-turn RL baselines across mathematical and general-domain benchmarks while offering substantially higher training efficiency.

  • Takeaways & Limitations

    DRIFT provides a simple alternative to online multi-turn reinforcement learning that avoids repeated training rollouts while retaining effective multi-turn optimization.

  • Takeaways & Limitations

    DRIFT is limited to short-horizon, verifier-guided correction with lightweight deterministic feedback, leaving stochastic human feedback and longer-horizon planning for future work.

Abstract

from arXiv · show

Large language models are increasingly deployed in multi-turn interactive settings where users or environments can iteratively provide lightweight feedback. Unfortunately, optimizing such behavior presents a sharp dilemma in practice: online reinforcement learning is able to effectively address multi-turn dynamics but is prohibitively expensive due to the cost of generating full correction trajectories at every update, whereas offline supervised fine-tuning (SFT) is efficient but suffers from distribution shift and behavioral collapse. To this end, we novelly propose DRIFT (Decoupled Rollouts and Importance-Weighted Fine-Tuning), a framework that operationalizes the theoretical insight that the KL-regularized RL objective is equivalent to importance-weighted supervised learning. DRIFT decouples rollout from optimization by sampling offline interaction trajectories from a fixed reference policy, deriving return-based importance weights, and optimizing the policy via weighted SFT on the resulting dataset. Empirically, we demonstrate that DRIFT matches or exceeds the performance of multi-turn reinforcement learning baselines while maintaining the training efficiency and simplicity of standard supervised fine-tuning. Code is available at https://github.com/2020-qqtcg/DRIFT.

1. Introduction

The introduction frames multi-turn optimization as a tradeoff between efficient but fragile offline SFT and effective but costly online RL. DRIFT addresses this tradeoff by decoupling rollouts from weighted-SFT optimization through a KL-regularized RL equivalence, with comparable or better performance and higher training efficiency.

  • Motivation: Single-turn-trained models are fragile under negative feedback, often repeating errors or degrading during multi-turn interaction.Real-world deployment requires users to iteratively provide feedback to guide the model.
  • Challenge: Offline correction-trajectory SFT is sample-efficient but suffers distribution shift and behavioral collapse, over-optimizing first-turn accuracy instead of meaningful later edits.The cited failure mode is associated with naive SFT.
  • Challenge: Online multi-turn RL addresses distribution issues but requires generating full interaction trajectories for every policy update, making rollout costs scale poorly with interaction length.This makes standard online RL computationally prohibitive for multi-turn reasoning tasks.
  • Method: DRIFT decouples rollout collection from policy optimization by reweighting offline reference-policy trajectories with exponentiated rewards.Its two stages are offline interaction-trajectory collection followed by trajectory-weighted SFT using derived importance weights.
  • Theory and Results: DRIFT’s weighted objective is theoretically equivalent to the KL-regularized RL objective, while experiments report comparable or better performance than online RL baselines and substantially higher training efficiency.The reported results span mathematical and general reasoning benchmarks.

2. Problem Setup

The paper models multi-turn answer correction as a finite-horizon MDP with token-sequence responses and interaction-history states. It optimizes discounted correctness under KL regularization to a reference policy, while noting that online RL requires costly fresh multi-turn rollouts at every update.

  • MDP Formulation: The correction task is formulated as a finite-horizon MDP M = (S, A, P, R, γ, T), where T is the maximum turn budget.The trajectory length satisfies L ≤ T.
  • States and Actions: Actions are token sequences generated by the model, while each state records the full interaction history, including prior responses and feedback.The action at turn t is yt ∈ A, and xt contains the initial prompt, earlier responses, and feedback.
  • Transition Dynamics: Transitions are deterministic: correct responses terminate episodes, whereas incorrect responses before T append fixed lightweight feedback and continue interaction.The fixed-feedback assumption matches evaluation protocols that use consistent, rigorous feedback without stochastic user noise.
  • Objective: The objective discounts binary correctness rewards to penalize delayed success and applies KL regularization controlled by β > 0 relative to πref.The policy πθ is initialized from the reference model, and the objective balances expected return with fidelity to that reference.
  • Computational Challenge: Online RL optimization requires fresh on-policy multi-turn trajectories from pθ at every update, making rollout cost scale with the interaction horizon.This motivates separating rollout generation from optimization in the proposed framework.

3. From KL-Regularized RL to Weighted SFT

This section shows that KL-regularized RL targets an exponentially return-weighted reference distribution and can be optimized through importance-weighted SFT using offline rollouts.

  • Optimal trajectory distribution: The optimal trajectory distribution p⋆(τ | x) is an exponential tilting of the reference distribution, with higher-return trajectories receiving exponentially greater probability mass.The temperature β controls the sharpness of this tilt, and Z(x) is the prompt-dependent partition function.
  • RL as divergence minimization: Maximizing the KL-regularized RL objective is equivalent to minimizing KL(pθ∥p⋆), but estimating its gradient requires fresh rollouts from the current policy pθ.This reverse-KL formulation creates the central online-RL bottleneck because its expectation is taken under pθ.
  • Forward-KL surrogate: Forward-KL provides an exact global surrogate when p⋆ is realizable and remains locally faithful near p⋆ without realizability, sharing the same local second-order geometry.In finite-capacity models, the global Forward-KL and Reverse-KL projections need not coincide.
  • Importance-weighted SFT: Importance sampling rewrites expectations over p⋆ as expectations over the reference policy using trajectory weights w(τ|x) = p⋆(τ|x)/pref(τ|x).This measure change is the core engine of DRIFT and yields a tractable objective.
  • Importance-weighted SFT: DRIFT decouples rollout from optimization by sampling trajectories offline from pref and applying normalized exponential-return weights in a standard supervised fine-tuning loop.For autoregressive language models, trajectory log-probabilities decompose into token-level log-probabilities.

4. The DRIFT Algorithm

DRIFT decouples offline multi-turn rollout generation from policy optimization by converting return-based importance weighting into weighted supervised fine-tuning. Its terminal-only retention is a protocol-specific approximation that can reduce variance under stop-on-success feedback while introducing bias relative to full-trajectory training.

  • Framework Overview: DRIFT samples multi-turn trajectories offline from a fixed reference policy, computes return-based importance weights, and optimizes the policy with weighted SFT.This avoids repeatedly generating rollouts during online policy updates, whose cost grows with the interaction horizon.
  • Offline Trajectory Generation: For each prompt, the deterministic protocol appends fixed lightweight feedback after incorrect responses and stops at success or the maximum turn budget.The resulting trajectories capture recovery from errors under fixed feedback.
  • Offline Trajectory Generation: DRIFT assigns each trajectory a scalar weight based on its shaped return, with prompt-level normalization stabilizing importance-sampling variance.The return prioritizes fewer turns through discounting and penalizes repetitive errors through unique-response counts.
  • Terminal-Step Retention: DRIFT retains only each trajectory’s terminal response and weight, which is a protocol-specific approximation rather than an exact implementation of the full-trajectory objective.This introduces bias by omitting intermediate imitation targets but can reduce variance when rejected intermediate responses produce noisy or misaligned gradients.
  • Weighted Supervised Optimization: The optimization stage minimizes weighted negative log-likelihood, a Monte Carlo approximation of forward KL minimization toward the optimal distribution.Weighted SFT amplifies gradients from desirable reasoning paths and suppresses suboptimal ones.
  • Terminal-Step Retention: Under the same training schedule, terminal supervision reaches higher accuracy and produces a smoother optimization curve than supervision that also weights rejected intermediate responses.The result supports a better practical bias-variance tradeoff for stop-on-success trajectories.

5. Experiments

Experiments evaluate DRIFT against SFT and RL baselines on in-domain and out-of-domain reasoning benchmarks using cumulative multi-turn accuracy. DRIFT improves multi-turn performance while retaining SFT-like efficiency, and its behavior is robust to key hyperparameter and feedback variations.

  • Experimental setup: Experiments train Qwen2.5-3B-Instruct and Llama3.1-8B-Instruct on the MATH subset of MetaMathQA.Evaluations use greedy decoding and distinguish math reasoning benchmarks from out-of-domain general reasoning benchmarks.
  • Metrics: The primary metric is cumulative accuracy with a maximum budget of 5 turns, denoted multi@5.Generation stops when the answer is correct or continues up to k turns; the metric uses a verifier over test samples.
  • Main results: Across most benchmarks, UFO outperforms SFT-based baselines, while DRIFT further improves upon UFO under matched settings.Single-turn training mainly improves first-turn accuracy and yields little gain on non-math benchmarks because it does not condition on negative feedback.
  • Main results: DRIFT matches or surpasses RL-based multi-turn improvements while retaining training efficiency comparable to SFT-based methods.RL approaches show stronger turn-by-turn performance than SFT but are substantially less training-efficient; GPU time includes rollout latency.
  • Ablations: A smaller γ encourages solving in fewer turns, while β = 0.1 performs best by balancing multi-turn gains and stability.Larger β weakens exponential tilting, whereas smaller β concentrates weights excessively; feedback wording has only minor performance impact.

6. Conclusion & Limitations · A. Related Work · B. Additional theoretical analyses

DRIFT targets short-horizon, verifier-guided multi-turn correction by separating trajectory generation from importance-weighted supervised optimization. Related work spans self-correction, correction-trajectory SFT, and KL-regularized policy optimization.

  • 6. Conclusion & Limitations: DRIFT samples correction trajectories once from a frozen reference policy and trains the target model with exponential, return-derived importance weights.This decouples trajectory generation from optimization under a KL-regularized multi-turn objective.
  • 6. Conclusion & Limitations: DRIFT provides a simple and stable alternative to online multi-turn reinforcement learning for lightweight negative feedback.
  • 6. Conclusion & Limitations: DRIFT is designed for short-horizon, verifier-guided correction with lightweight deterministic feedback and only a small number of correction attempts.The verifier supplies an unambiguous correctness signal within the stated evaluation protocol.
  • 6. Conclusion & Limitations: Stochastic or preference-based human feedback, open-ended dialogue objectives, and substantially longer-horizon interactive planning require additional modeling beyond DRIFT’s stated boundary.
  • A. Related Work: LLMs can revise answers from minimal negative feedback, but extended multi-turn contexts may cause behavioral drift or degradation.The related work characterizes self-correction as nontrivial but fragile.
  • A. Related Work: Correction traces and multi-turn demonstrations are commonly used for supervised fine-tuning, while self-training methods distill generated solutions or rationales.These approaches aim to improve revision or problem-solving behavior with limited additional supervision.
  • A. Related Work: RL and RLHF optimize multi-turn outcomes from rollouts and rewards, often using KL regularization to keep policies near a reference model.PPO-style RLHF and sequence-level methods such as GRPO are discussed as KL-regularized policy optimization approaches.

B.1. Estimation Stability and Sample Complexity … C.2. Proof of Thm 2

The section establishes sample-complexity conditions for stable normalized importance weights and proves the variational characterization underlying the framework. It shows that weak regularization can require exponentially many samples, while the theorem proofs establish existence, uniqueness, and the claimed identity under support conditions.

  • B.1. Estimation Stability and Sample Complexity: Bounded returns define finite weight bounds mβ and Mβ, framing concentration of the partition estimate as the basis for gradient-estimator stability.The normalized weights depend on the Monte Carlo estimate bZ(x) in the denominator, where errors can be amplified when the partition function is small.
  • B.1. Estimation Stability and Sample Complexity: The required sample size K ensures |bZ(x) − Z(x)| ≤ ϵ with probability at least 1 − δ, yielding convergence of empirical weights to their theoretical values.Bounding bZ is a prerequisite for stable training, and sufficiently large K ensures convergence of the empirical weights.
  • B.1. Estimation Stability and Sample Complexity: exp(2Rmax/β) bounds the worst-case scaling of K for fixed precision ϵ as β → 0, motivating moderate regularization to balance distribution matching and sample efficiency.In the weak-regularization regime, ∆β grows exponentially and can cause the gradient estimator to be dominated by high-return trajectories.
  • C.1. Proof of Thm 1: The proof of Thm. 1 formulates the KL-regularized variational problem for valid trajectories under the deterministic protocol and restricts feasible policies to the support of pref.If pref(τ) = 0, any feasible p with p(τ) > 0 has infinite KL divergence, so an optimum must assign p(τ) = 0 there.
  • C.1. Proof of Thm 1: The stationary solution p⋆ is well-defined because bounded returns make 0 < Z(x) < ∞.The proof derives the candidate maximizer by differentiating the Lagrangian and imposing normalization.
  • C.1. Proof of Thm 1: Strict concavity of the objective over the feasible region makes p⋆(· | x) the unique global maximizer of the variational problem.The feasible set is convex, expected return is linear, and −β KL(p∥pref) is strictly concave on policies absolutely continuous with respect to pref.
  • C.2. Proof of Thm 2: The proof of Thm. 2 reduces the claim to Pθ, Pref, and P⋆ for a fixed prompt x and handles the case Pθ not absolutely continuous with respect to Pref by infinite KL divergence.Under the convention J(θ) = −∞, the claimed identity holds trivially when Pθ̸ ≪Pref; otherwise the proof assumes Pθ ≪Pref and concludes the theorem.

C.3. Proof of Lemma 3 … D.2. Setup

The paper establishes the KL and importance-weighting results underlying DRIFT, then formalizes gradient and sampling properties and specifies experimental baselines and training settings.

  • C.3. Proof of Lemma 3: Lemma 3 shows that both forward and reverse KL divergences attain their global minimum exactly when Pθ = P⋆, while reverse-KL minimization is equivalent to cross-entropy minimization.The reverse-KL minimum value is 0, and the forward-KL objective differs from cross-entropy only by a θ-independent term.
  • C.4. Proof of Lemma 4: Lemma 4 establishes that forward and reverse KL share the same quadratic local geometry, with their difference controlled by third-order remainder terms.The first-order terms vanish, leaving a shared quadratic term around p⋆ even when p⋆∉Πθ.
  • C.5. Proof of Thm 5: Theorem 5 proves that the KL-regularized objective and importance-weighted negative log-likelihood are identical for every θ and therefore have the same minimizers.The weight is w(τ | x) = p⋆(τ | x) / pref(τ | x) = 1/Z(x) exp(…).
  • C.6. Proof of Prop. 6: Proposition 6 bounds the bias introduced by omitting terminal gradient contributions and shows that the resulting terminal-gradient estimator has lower variance under the stated condition.The variance result is stated as Var(gall) > Var(gterm).
  • C.7. Proof of Prop. 7: Proposition 7 derives a sample-size condition for estimating the weighted objective by applying Hoeffding’s inequality under bounded returns R(τ) ∈ [Rmin, Rmax].The proof bounds the range of the weighted-return variables and rearranges the concentration requirement to obtain equation (23).
  • D.1. Baseline: The experiments compare DRIFT with SFT-5TURN, STAR-2TURN, SCORE-2TURN, and UFO-5TURN under correction-based multi-turn protocols.SFT-5TURN is the unweighted ablation, while the other baselines use self-training, KL-regularized two-turn correction, or repeated trial-and-error prompting.
  • D.2. Setup: The offline trajectory-generation setup uses K = 16 rollouts at temperature 1.0, with maximum trajectory length T = 5 and 512 new tokens.These settings define the rollout budget and trajectory limits for the experiments.
  • D.2. Setup: Optimization uses global batch size 128 for 200 steps, with β = 0.1, γ = 0.9, and repetition penalty coefficient λ = 0.5.These are the stated optimization-phase and regularization hyperparameters.

D.3. Benchmark

The benchmark suite evaluates multi-turn correction on mathematical reasoning and general-domain knowledge and science tasks, distinguishing same-family performance from out-of-domain transfer beyond MATH-subset training.

  • Benchmark scope: Evaluation spans mathematical reasoning and general-domain knowledge and science reasoning to test correction behavior within and beyond the training task.Training uses only the MATH subset of MetaMathQA; mathematical benchmarks are same-family evaluations, while general reasoning benchmarks assess out-of-domain transfer.
  • Mathematical reasoning: MATH is the primary mathematical benchmark, featuring competition-style problems that require multi-step derivations.MATH is supplemented by MATH500, a 500-problem evaluation subset commonly used for faster iteration.
  • Mathematical reasoning: TheoremQA evaluates theorem-driven STEM problem solving, emphasizing selection and application of appropriate theorems.It complements MATH and MATH500 within the mathematical reasoning evaluation family.
  • General reasoning: MMLU-Redux evaluates cross-domain transfer using a cleaned and re-annotated MMLU designed to reduce ambiguity and labeling errors.The benchmark is part of the general reasoning suite.
  • General reasoning: MMLU-Pro increases difficulty by expanding answer options and filtering trivial items, while GPQA-diamond targets graduate-level science questions.Both benchmarks extend evaluation beyond the mathematical training task.

D.4. Additional Model Results … Evaluation for GPQA

Additional experiments show that DRIFT scales to Qwen2.5-7B-Instruct, benefits modestly from rollout refresh, and provides a useful warm start for online RL. The section also specifies the prompts used across rollout and benchmark evaluation.

  • D.4. Additional Model Results: DRIFT raises Qwen2.5-7B-Instruct’s all-benchmark average from 64.8% to 68.3% and slightly exceeds online multi-turn RL, with largest gains on MATH and MATH500.The stronger-backbone experiment uses the same protocol as the main evaluation.
  • D.5. A Simple Rollout-Refresh Variant: A two-stage rollout-refresh variant first trains DRIFT for 100 steps, regenerates correction trajectories, and then continues training to mitigate offline coverage limitations.The refresh is evaluated on Qwen2.5-3B-Instruct without fully reverting to online RL.
  • D.5. A Simple Rollout-Refresh Variant: 60.5% to 61.2%: rollout refresh modestly improves the all-benchmark average, providing an initial diagnostic rather than a complete replacement for online exploration.The authors leave more systematic rollout-refresh schedules for future work.
  • D.6. DRIFT as an Initialization for Online RL: The hybrid initialization experiment fixes the total budget at 200 steps, allocating 100 steps to SFT or DRIFT and 100 steps to UFO.Pure UFO and pure DRIFT each use all 200 steps.
  • D.6. DRIFT as an Initialization for Online RL: DRIFT followed by UFO achieves the best all-benchmark average, improving over pure UFO by 2.3 points, whereas SFT followed by UFO does not improve over pure UFO.The results provide preliminary evidence that DRIFT is a useful warm start for verifier-guided online RL, while broader schedule studies remain future work.
  • D.7. Prompt, Rollout for MetaMathQA: The section presents the prompts used for rollout and evaluation, including a helpful-assistant system instruction and required <think> and <answer> response formatting.The rollout prompt additionally specifies five actions left and a 400-word maximum response length.
  • Evaluation for MATH / MATH500, Evaluation for TheoremQA: Math and theorem evaluations require final answers enclosed between <answer> tags, with theorem problems supplied using an image placeholder.The MATH and MATH500 prompts include explicit think-and-answer formatting, while TheoremQA uses <image> for the problem.
  • Evaluation for MMLU-Redux / MMLU-Pro, Evaluation for GPQA: Multiple-choice prompts require only the final answer letter between <answer> tags, using A-J for MMLU-Redux/MMLU-Pro and A-D for GPQA.Both prompts also require the specified think-and-answer format with no extra text.

D.8. Ablation on return shaping … E.2. Case 2

DRIFT’s return shaping improves correction behavior and accuracy, while its training remains stable across hyperparameter settings and depends on sufficiently diverse rollout groups. Case studies show that DRIFT avoids behavioral collapse and explores alternative reasoning paths to recover from persistent errors.

  • D.8. Ablation on return shaping: Removing the trajectory penalty decreases accuracy, correction rate, and average unique answers, indicating that the penalty encourages changing incorrect responses.The ablation is reported in Tables 6 and 7.
  • D.8. Ablation on return shaping: Correction rate is defined relative to turn 1 as Corr = Acc@5−Acc@1.The reward-shaping ablation reports first-turn and cumulative 5-turn accuracy alongside correction rate.
  • D.9. Learning Curves Across Hyperparameter Settings: Across broad ranges of γ and β, multi-turn accuracy improves steadily and remains stable without abrupt collapse or divergence.Different settings produce modest differences in final accuracy.
  • D.9. Learning Curves Across Hyperparameter Settings: Small rollout numbers K produce more degenerate all-correct or all-wrong groups, yielding nearly uniform weights and performance close to SFT-5Tu.The analysis fixes batch size and training steps and summarizes group-level distributions in Table 8.
  • E.1. Case 1: In Case 1, the base model repeats the inconsistent claim that “c can be 0 or 1” across four turns, while generic feedback fails to break the loop.The MATH-500 problem requires finding integer x satisfying x^3|10! with the constraint 3c ≤2 for prime factor 5.
  • E. Case Study: Case 1 evaluates DRIFT’s error-recovery behavior on MATH-500 through a four-turn comparison of repeated failure and correction.The case is presented in Figure 10, with effective groups defined as neither all-correct nor all-wrong.
  • E.1. Case 1: DRIFT uses negative feedback to prune the search space, re-verifying in Turn 3 that 5^3 ∤10! and correcting the exponent count to 1.This demonstrates exploration of alternative reasoning paths rather than resampling high-probability errors.
  • E.2. Case 2: In Case 2, the base model hallucinates verification of a wrong modular-arithmetic answer and loops despite external feedback, whereas DRIFT explores different candidate solutions across turns.The example is drawn from MATH-500 and is shown in Figure 11.

E.3. Case 3: Limitations in Knowledge-Intensive Domains

DRIFT learns strategic multi-turn correction behavior, but its reasoning accuracy remains bounded by the model’s underlying domain knowledge. In knowledge-intensive tasks, it can avoid blind correction loops without reliably supplying the expertise needed for valid solutions.

  • E.3. Case 3: Limitations in Knowledge-Intensive Domains: In organic chemistry, DRIFT’s carbocation reasoning is heuristically plausible but factually flawed for the concerted mCPBA mechanism.The Base Model also hallucinates an impossible electrophilic aromatic substitution mechanism, while UFO repeats reasoning and cycles through options.
  • E.3. Case 3: Limitations in Knowledge-Intensive Domains: DRIFT instills multi-turn correction strategy, yet does not spontaneously inject missing subject-matter expertise.The limitation reflects insufficient underlying capability rather than a defect in the DRIFT objective; domain-specific continual pre-training is proposed as a future direction.
  • E.3. Case 3: Limitations in Knowledge-Intensive Domains: DRIFT initially errs on the factorial problem but uses feedback to correct its analysis, whereas the Base Model repeatedly miscalculates despite feedback.DRIFT identifies that 5^3 does not divide 10! and derives the correct answer in Turn 3.
Loading 2605.31455v1…