DRIFT: Replacing Expensive Online Rollouts in Multi-Turn RL with Offline Weighted SFT

Forum topic · 小凯 · 2026-10-22

Summary

A team from HKUST (Guangzhou) and CUHK-Shenzhen introduced DRIFT (arXiv 2605.31455, under ICML review), a method that trains large language models for multi-turn error correction without repeatedly regenerating trajectories during reinforcement learning updates. DRIFT exploits the closed-form solution of KL-regularized RL: the optimal trajectory distribution is an exponential tilt of the reference distribution, proportional to exp(R/β). The method samples a large batch of multi-turn correction trajectories once from a fixed reference policy, computes importance weights from rewards, and performs weighted supervised fine-tuning—provably equivalent to native KL-regularized online RL. Rewards include a discount γ^(t-1) incentivizing early correction and a penalty on repeated errors via unique response counts. On benchmarks including MATH, MATH500, ThmQA, MMLU, and GPQA, DRIFT matches or slightly exceeds the online multi-turn RL baseline UFO on Qwen2.5-3B (60.5 vs 60.2) and Qwen2.5-7B (68.3 vs 64.8), while trailing slightly on Llama3.1-8B (55.6 vs 56.1). GPU cost scales close to SFT as turn counts grow. The authors acknowledge coverage limitations of one-shot offline sampling and surface-editing failures on knowledge-intensive tasks, positioning DRIFT as a multi-turn extension of DPO-style closed-form objectives.

DRIFT: Taking the Most Expensive Step Out of the Training Loop

Consider a scenario that happens every day: a user asks a model to solve a math problem, the first attempt is wrong, the user gives a hint, the model produces a corrected second attempt, it is wrong again, another hint, another fix. This kind of multi-turn iterative correction is becoming the real way large language models work—humans no longer ask just once but iterate alongside the model. The question is how to train this capability in. A team from HKUST (Guangzhou) and CUHK-Shenzhen offered an answer in a late-May paper: DRIFT (arXiv 2605.31455, under ICML review).

The Dilemma

The paper starts from a trade-off: online reinforcement learning can learn multi-turn dynamics, but every update requires regenerating complete multi-turn correction trajectories—as the model learns more, every conversation must be rerun. Offline supervised fine-tuning is cheap, but the resulting models often optimize only for the first turn; later-turn "corrections" degenerate into blind edits, errors recur, and behavior collapses.

The Key Theoretical Step

DRIFT's crucial move happens in theory. The KL-regularized RL objective can be rewritten as a distribution matching problem with a closed-form optimal solution: the optimal trajectory distribution is an exponential tilt of the reference distribution—

p* ∝ p_ref · exp(R/β)

where R is the trajectory return and β is a temperature. Translated: the optimal policy simply reallocates probability mass over reference behaviors according to their returns—high-return trajectories get exponentially amplified probability. If the answer has this form, why rerun online sampling at every update? Instead, sample a large batch of multi-turn trajectories once from the fixed reference policy, compute importance weights from the returns, and run ordinary weighted supervised fine-tuning on this weighted dataset. Theory guarantees this objective is equivalent to native KL-regularized RL.

Technically, this shifts the expectation in the divergence from the dynamically changing policy to the fixed reference distribution. The cost is introducing importance weights; the benefit is that expensive rollout sampling changes from "paid at every update step" to "paid once up front."

Three Mechanism Details Worth Remembering

Discounted trajectory return: The task is modeled as a finite-horizon Markov decision process with discounted returns. If the model only gets it right at turn t, the reward is discounted by γ^(t-1). Solving the same problem on the first correction is worth far more than fixing it on the fifth—early error correction is explicitly incentivized.
Repetition penalty: The return includes a penalty based on the number of unique responses for repeated mistakes, a design borrowed from the online baseline UFO.
Fully decoupled training: Training proceeds in two entirely separate phases—first collect data, then learn, with no revisiting.

Results

Under the same evaluation suite (MATH, MATH500, ThmQA, two MMLU variants, GPQA; pass rate over five turns):

Qwen2.5-3B: DRIFT averages 60.5 vs UFO's 60.2—a slight win.
Qwen2.5-7B (additional experiments): 68.3 vs 64.8—a modest win.
Llama3.1-8B: DRIFT 55.6 vs UFO 56.1—the online method pulls ahead at this scale.

So "comparable to or better than online RL" should be read carefully: it holds for small models and the 7B, but at 8B it is a slight loss.

On efficiency, run for 200 steps on two configurations (4× A800 and 4× H20), DRIFT's GPU time grows only mildly as turn counts increase, tracking close to the SFT scaling curve, while UFO becomes clearly more expensive with more turns. What is saved is precisely the per-update resampling overhead.

Limitations

The paper itself is upfront about boundaries:

Coverage limits of one-shot offline collection: if the reference policy never samples a certain kind of useful correction trajectory, importance weights cannot conjure it from nothing. An appendix tries a simple variant—resampling midway using a checkpoint—with small gains; systematic mixed schedules are left as an open problem.
Knowledge-intensive tasks: in case studies, when the model lacks the actual knowledge, multi-turn correction degrades into superficial rewriting.

Positioning

Single-turn alignment has prior precedent: DPO was exactly the move of replacing online sampling with a closed-form objective. DRIFT is essentially the first extension of this idea into multi-turn error correction, making "sample enough first, then learn at leisure" work for interactive tasks too. The expensive part was never the learning itself—it was generating the data over and over.