1 Introduction↩︎

Large language models (LLMs) [1]–[3] have evolved from static query-response engines into interactive agents capable of advanced reasoning. While standard training pipelines focus predominantly on single-turn accuracy [4]–[6], real-world deployment necessitates multi-turn capabilities where users iteratively provide feedback to guide the model [7], [8]. However, models trained strictly on single-turn data often exhibit fragility when confronted with negative feedback [9]–[12], frequently repeating errors or degrading in performance after a revision attempt. As shown in Figure 1, effectively leveraging lightweight feedback signals (e.g., “Incorrect, please try again”) to improve robust multi-turn reasoning remains a critical open challenge.

Current approaches to multi-turn optimization face a sharp dilemma between effectiveness and efficiency. On one hand, Supervised Fine-Tuning (SFT) on offline correction trajectories is sample-efficient but often fails to learn genuine correction policies. As noted by [10], naive SFT suffers from distribution shift and behavioral collapse, where models over-optimize for first-turn accuracy while failing to produce meaningful edits in subsequent turns. On the other hand, Online Reinforcement Learning (RL) [9], [10], [12] approaches like PPO [13], [14] or GRPO [1], [15] address these distribution issues but incur a prohibitive computational cost. Unlike single-turn settings, multi-turn optimization requires generating full interaction trajectories for every policy update. As illustrated in Figure 7, this rollout cost scales poorly with interaction length, making standard online RL hard for training on multi-turn reasoning tasks.

We address this bottleneck by establishing a fundamental connection between the KL-regularized RL objective and weighted supervised learning. Our core insight is that the gradient of the online RL objective can be approximated using offline trajectories sampled from a fixed reference policy, provided they are re-weighted by their exponentiated rewards. This theoretical equivalence implies that the expensive rollout generation can be completely decoupled from the policy optimization process. By shifting the computational burden to an offline, parallelizable generation phase, we can project the benefits of online RL into a high-throughput supervised training framework without the complexity of online interactions.

Building on this insight, we propose DRIFT ( ecoupled Rollouts and Importance-Weighted Fine- Tuning), a practical framework for verifiable multi-turn optimization. DRIFT operates in two distinct stages: (1) an offline rollout stage where interaction trajectories are collected under a reference policy, and (2) a trajectory-weighted SFT stage where the model is optimized using our derived importance weights. This design enables DRIFT to match the asymptotic performance of online RL baselines while retaining the computational efficiency of standard SFT.

Our contributions are summarized as follows:

• We propose DRIFT, a method that decouples rollout from optimization to realize an RL objective via weighted SFT, thereby enabling efficient training under multi-turn interaction protocols.

• We provide theoretical results demonstrating that our weighted objective is equivalent to the underlying KL-regularized RL objective, ensuring the effectiveness and stability of the optimization.

• Extensive experiments show that DRIFT achieves performance comparable to or better than online RL baselines across mathematical and general reasoning benchmarks while offering substantially higher training efficiency.

2 Problem Setup↩︎

We formulate the multi-turn answer correction task as a finite-horizon Markov Decision Process (MDP), defined by the tuple \(\mathcal{M} = (\mathcal{S}, \mathcal{A}, \mathcal{P}, \mathcal{R}, \gamma, T)\), where \(T\) represents the maximum turn budget. We detail the components and the underlying assumptions below.

States and Actions. Let \(\mathcal{V}\) denote the vocabulary. The action space \(\mathcal{A} = \mathcal{V}^*\) consists of sequences of tokens generated by the model. At turn \(t\), the action \(y_t \in \mathcal{A}\) represents the model’s response. The state space \(\mathcal{S}\) represents the interaction history. Given an initial problem prompt \(x_1\), the state at turn \(t\) is defined as the full interaction history up to that point: \(x_t = (x_1, y_1, f_1, \dots, y_{t-1}, f_{t-1})\), where \(f_i\) denotes the feedback received at turn \(i\). Since Large Language Models are autoregressive and condition on the entire input context, treating the complete interaction history \(x_t\) as the state ensures that the Markov property holds strictly, as \(x_t\) contains all sufficient statistics for the next transition.

Transition Dynamics. The transition function \(\mathcal{P}: \mathcal{S} \times \mathcal{A} \to \mathcal{S}\) is deterministic in our setting. A verifier function \(\mathcal{V}(y_t)\) evaluates the correctness of the response \(y_t\). If \(y_t\) is correct, the episode terminates successfully. If \(y_t\) is incorrect and \(t < T\), the environment transitions to \(x_{t+1}\) by appending a lightweight feedback message \(f\) (e.g., “Incorrect, please try again”): \[x_{t+1} = \text{concat}(x_t, y_t, f).\] We assume a deterministic transition dynamics with fixed feedback \(f\). This assumption aligns with standard evaluation protocols in reasoning benchmarks [10], [12], where the user provides consistent, rigorous feedback to elicit self-correction without introducing stochastic user noise.

Objective. Let \(\pi_\theta(y|x)\) be the parameterized policy initialized from a reference model \(\pi_{\text{ref}}\). A trajectory is a sequence \(\tau = (x_1, y_1, \dots, x_{L}, y_{L})\), where \(L \le T\) is the effective episode length. We define a discounted trajectory return: \[R(\tau) \triangleq \sum_{t=1}^{L} \gamma^{t-1} r(x_t, y_t), \quad \gamma \in (0, 1),\] where \(r(x_t, y_t) = 1\) if \(y_t\) is correct and \(0\) otherwise. The discount factor \(\gamma\) strictly penalizes delayed success, incentivizing the model to correct errors as early as possible. Our goal is to optimize the standard KL-regularized RL objective [4], [14], which balances maximizing the expected return against maintaining fidelity to the reference policy: \[\begin{align} \max_{\theta}\; J(\theta) &\triangleq \mathbb{E}_{\tau \sim p_\theta}\!\left[R(\tau)\right] \\ &\quad - \beta\, \mathrm{KL}\!\left( p_\theta(\cdot \mid x) \,\|\, p_{\mathrm{ref}}(\cdot \mid x) \right). \end{align} \label{eq:rl95objective}\tag{1}\] where \(\beta > 0\) controls the strength of the regularization, and \(p_\theta(\tau\mid x)\) and \(p_{\text{ref}}(\tau\mid x)\) are the trajectory distributions defined by \(\pi_{\theta}\) and \(\pi_{\text{ref}}\), respectively. This formulation grounds our multi-turn correction problem as a specific instance of return maximization under distribution constraints. However, optimizing Eq. 1 with standard online RL requires generating fresh on-policy multi-turn trajectories from \(p_\theta\) at every update, so the rollout cost scales with the interaction horizon.

3 From KL-Regularized RL to Weighted SFT↩︎

In this section, we establish a fundamental equivalence between the online KL-regularized reinforcement learning objective and importance-weighted supervised fine-tuning.

3.1 The Optimal Trajectory Distribution↩︎

To manipulate this objective effectively, it is instructive to first abstract away the parametric constraints and identify the theoretically optimal distribution \(p^\star(\cdot|x)\) that maximizes Eq. 1 over the space of all valid probability distributions.

This result is pivotal because it reveals that the ideal correction behavior is not arbitrary; it essentially re-weights the reference behavior such that trajectories with higher returns \(R(\tau)\) are assigned exponentially higher probability mass. The parameter \(\beta\) serves as the temperature of this distribution, controlling the sharpness of the tilt towards high-return regions.

3.2 Connecting RL to Distribution Matching↩︎

Having identified the theoretically optimal target \(p^\star\), we now establish the fundamental connection between the standard reinforcement learning objective and distribution matching. While \(p^\star\) was derived as the maximizer of the variational problem, it is not immediately obvious how the parametric RL objective \(J(\theta)\) relates to this distribution. The following theorem bridges this gap by explicitly reframing the reward maximization problem as a divergence minimization task.

Theorem 2 is significant because it explicitly characterizes the ideal optimization path: the policy \(p_\theta\) should strive to cover the mode of \(p^\star\). Standard online RL algorithms, such as PPO, effectively minimize this Reverse-KL divergence. However, a critical bottleneck arises from the direction of the divergence. The term \(\text{\normalfont KL}(p_\theta \| p^\star)\) involves an expectation under the current policy \(p_\theta\) (i.e., \(\mathbb{E}_{\tau \sim p_\theta}[\dots]\)). Estimating its gradient requires generating fresh rollouts from \(p_\theta\) at every optimization step, incurring the high computational costs we aim to avoid.

To overcome this, we propose utilizing the Forward-KL divergence, \(\text{\normalfont KL}(p^\star \,\|\, p_\theta)\), as a surrogate objective. This substitution allows us to shift the expectation from the changing policy \(p_\theta\) to the fixed optimal distribution \(p^\star\). We formalize this substitution in two complementary ways: an exact global statement under realizability, and a local surrogate guarantee that does not require realizability.

Lemma 1 gives an exact global guarantee, but its role is limited to the realizable case: it uses \(p^\star\in\Pi_\theta\) only to identify the global minimizers of the two KL directions. In finite-capacity models, \(p^\star\) may not be exactly realizable, and the global projections induced by Forward-KL and Reverse-KL need not coincide. We therefore complement Lemma 1 with a local comparison result showing that the two KL objectives still have the same second-order geometry around \(p^\star\).

Together, Lemmas 1 and 2 justify the use of Forward-KL as a training surrogate at two levels. In the realizable case, the substitution is exact at the level of global optima; in the non-realizable case, it remains locally faithful whenever \(p_\theta\) is close to \(p^\star\). We now exploit the computational advantage of the Forward-KL objective: its only \(\theta\)-dependent term is the cross-entropy under the fixed target distribution \(p^\star\). Hence,

\[\min_{\theta} \text{\normalfont KL}(p^\star \,\|\, p_\theta) \iff \max_{\theta} \mathbb{E}_{\tau \sim p^\star(\cdot|x)} [\log p_\theta(\tau|x)]. \label{eq:forward95likelihood}\tag{2}\]

3.3 Deriving the Offline Weighted Objective↩︎

The challenge remains that we cannot sample directly from the optimal distribution \(p^\star\) to compute this expectation. However, since \(p^\star\) is absolutely continuous with respect to the reference policy \(p_{\text{ref}}\), we can apply importance sampling to change the measure from \(p^\star\) to \(p_{\text{ref}}\). This transformation is the core engine of DRIFT.

This formulation fully decouples the rollout phase from the optimization phase. We can effectively approximate \(\mathbb{E}_{\tau \sim p_{\text{ref}}}\) using a static dataset of trajectories sampled offline from the reference policy. The learning process then reduces to a standard supervised fine-tuning loop where each trajectory loss is scaled by its normalized exponential return. For autoregressive language models, the term \(\log p_\theta(\tau|x)\) naturally decomposes into a sum of token-level log-probabilities, allowing Eq. [thm:weighted95sft] to be implemented by simply applying the scalar weight \(w(\tau|x)\) to the cross-entropy loss of every token in the response.

4 The DRIFT Algorithm↩︎

4.1 Overview of DRIFT↩︎

The theoretical analysis in Section 3 establishes that the complex online RL objective can be equivalently transformed into a tractable weighted supervised learning problem. Guided by this theoretical equivalence, we propose DRIFT, a framework that operationalizes this insight to fundamentally decouple interaction rollouts from policy optimization. As illustrated in Figure 2, DRIFT effectively bridges the gap between the rigorous RL objective and practical supervised fine-tuning through two distinct stages:

Stage 1: Offline Trajectory Generation. Instead of expensive online interactions during training, DRIFT samples trajectories offline from a fixed reference policy \(\pi_{\text{\normalfont ref}}\). Crucially, to align with the optimal RL objective derived in Eq. ?? , we concurrently compute a scalar importance weight \(w(\tau)\) for each trajectory. This process, detailed in Algorithm 3, efficiently transforms return signals into a high-quality, importance-weighted dataset.

Stage 2: Weighted Supervised Optimization. The policy \(\pi_\theta\) is then updated via weighted supervised fine-tuning on this pre-collected dataset. DRIFT avoids the extra overhead of repeatedly performing rollouts during the updates of online RL methods. This is particularly critical in multi-turn settings, where the cost of a single rollout grows with the interaction horizon and can be substantially higher than that of a single-turn rollout. Meanwhile, the theoretical guarantees of the KL-regularized objective enable DRIFT to achieve performance close to strong multi-turn RL baselines.

4.2 Stage 1: Offline Trajectory Generation↩︎

The primary goal of this stage is to construct a high-quality dataset \(\mathcal{D}\) that allows us to approximate expectations under the optimal distribution \(p^\star\) using samples from the reference policy \(\pi_{\text{\normalfont ref}}\). As detailed in Lines 1–17 of Algorithm 3, this process involves three key components:

Interaction Protocol. For each problem prompt \(x_1\), we sample \(K\) trajectories \(\{\tau^{(k)}\}_{k=1}^K\) from the reference policy \(\pi_{\text{\normalfont ref}}\) under a deterministic multi-turn protocol. At any turn \(t < T\), if the response \(y_t\) is incorrect, a fixed lightweight feedback \(f\) is deterministically appended to construct the next state: \[\label{eq:transition} x_{t+1} = \text{\normalfont concat}(x_t, y_t, f).\tag{3}\] The interaction terminates immediately upon generating a correct answer or reaching the maximum turn budget \(T\). This protocol ensures that the collected trajectories explicitly capture the model ability to recover from errors under fixed feedback.

Return Computation. To strictly align the model behavior with the desired multi-turn outcome, we design a shaped return \(R(\tau)\) that incorporates both efficiency and diversity. Let \(L\) denote the effective length of trajectory \(\tau\). The return is defined as: \[\label{eq:reward95shaped} R(\tau) \triangleq \mathbb{I}(y_L \text{ is correct}) \cdot \gamma^{L-1} - \lambda \left( 1 - \frac{E(\tau)}{L} \right),\tag{4}\] where \(\gamma \in (0,1)\) serves as a discount factor to prioritize solving problems in fewer turns, and the second term imposes a penalty based on the unique response count \(E(\tau)\) to discourage repetitive errors, a design inspired by UFO [12]. We analyze its effect in Section 10.8.

Importance Weight Calculation. To operationalize the theoretical equivalence established in Theorem 3, we assign a scalar importance weight \(w^{(k)}\) to each trajectory. This weight acts as the Radon-Nikodym derivative that adjusts the sampling distribution: \[\label{eq:weight95calculation} \begin{align} w^{(k)} & \leftarrow \frac{\exp\left(R(\tau^{(k)})/\beta\right)}{\widehat{Z}(x_1)}, \\ \widehat{Z}(x_1) & \triangleq \frac{1}{K}\sum_{j=1}^K \exp\left(\frac{R(\tau^{(j)})}{\beta}\right). \end{align}\tag{5}\] Here, the partition function estimate \(\widehat{Z}(x_1)\) provides prompt-level normalization. This ensures that the weights reflect the relative quality of trajectories within the specific prompt’s solution space, stabilizing the variance of the importance sampling estimator.

Terminal-Step Retention. Finally, to construct the dataset \(\mathcal{D}\), we retain only the terminal turn \((x_{L}, y_{L})\) of each trajectory paired with its weight \(w^{(k)}\). This is a protocol-specific approximation to the full-trajectory weighted objective in Theorem 3, rather than an exact implementation of that objective. Under the stop-on-success protocol, intermediate turns in a successful trajectory are verifier-rejected attempts generated under fixed negative feedback. Applying the same large trajectory weight to all such intermediate turns can therefore assign positive credit to responses that did not realize success. Terminal-only retention still conditions on the full interaction history through \(x_L\), but concentrates the trajectory weight on the final response that determines the verifier outcome. We analyze and empirically validate this approximation in Section 4.4.

4.3 Stage 2: Weighted Supervised Optimization↩︎

Once the importance-weighted dataset \(\mathcal{D}\) is constructed, the second stage focuses on distilling the optimal behavior into the parameterized policy \(\pi_\theta\). This phase effectively converts the reinforcement learning problem into a standard supervised learning task, characterized by the following three aspects:

Optimization Objective. Intuitively, we aim to optimize the policy as if it were learning from the theoretical optimal distribution \(p^\star\). Since we only have access to trajectories from the reference policy, the importance weights act as a statistical correction mechanism, effectively reshaping the offline data to serve as a proxy for the ideal behavior. Formally, we optimize \(\pi_\theta\) by minimizing the weighted negative log-likelihood: \[\label{eq:weighted95objective} \mathcal{L}(\theta) = \mathbb{E}_{(x, y, w) \sim \mathcal{D}} \left[ -w \cdot \log \pi_\theta(y \mid x) \right].\tag{6}\] This objective is not heuristically chosen; rather, it serves as a Monte Carlo approximation of the cross-entropy between the optimal distribution \(p^\star\) and the policy \(\pi_\theta\). As derived in Eq. [thm:weighted95sft], minimizing this weighted loss is theoretically equivalent to minimizing the forward KL divergence \(\text{\normalfont KL}(p^\star \| \pi_\theta)\), thereby ensuring the policy converges towards the optimal solution implied by the KL-regularized RL target.

Token-Level Realization. Since the trajectory return is a holistic outcome of the entire generation sequence, we must propagate the trajectory-level evaluation to individual generation steps. In the context of autoregressive models, we implement this by applying the scalar weight \(w\) uniformly to the loss of every token in the terminal response \(y\): \[\label{eq:token95loss} -\log \pi_\theta(y \mid x) = \sum_{j=1}^{M} -\log \pi_\theta(y_j \mid x, y_{<j}),\tag{7}\] where \(M\) is the sequence length. This formulation renders our optimization mechanically equivalent to standard weighted Supervised Fine-Tuning (SFT). This mechanism explicitly amplifies the gradients for desirable reasoning paths while suppressing suboptimal ones, effectively incentivizing the model to assign higher probability mass to tokens that belong to efficient and correct solutions.

A critical advantage of this design is the complete decoupling of rollout generation from parameter updates. Unlike online RL algorithms (e.g., PPO) that require computationally expensive trajectory sampling at every training step, DRIFT treats the weighted dataset \(\mathcal{D}\) as a fixed offline corpus. This allows the optimization phase to enjoy the high throughput and stability of standard Supervised Fine-Tuning (SFT), making it significantly more scalable for large language models.

4.4 Protocol-Specific Motivation for Terminal-Step Retention↩︎

Theorem 3 gives a full-trajectory weighted objective. In the practical implementation of DRIFT, however, we apply the trajectory weight only to the terminal response \((x_L,y_L)\). This terminal-only loss is not an exact implementation of the full-trajectory objective; it is a protocol-specific approximation for the stop-on-success correction setting.

The motivation is credit assignment. Since the episode terminates immediately after the first correct response, every non-terminal response has been rejected by the verifier under the fixed feedback protocol. Applying the same large trajectory weight to all turns may therefore assign positive imitation signal to rejected attempts. Terminal-step retention avoids this issue by supervising only the final response, while still keeping the full interaction history in the conditioning state \(x_L\). We formalize this as a bias-variance motivation below.

4.4.0.1 Gradient Decomposition.

Let \(\ell_t(\theta) \triangleq -\log \pi_\theta(y_t \mid x_t)\). The full-trajectory and terminal-only gradient estimators are \[g_{\mathrm{all}} \triangleq w(\tau) \sum_{t=1}^{L} \nabla_\theta \ell_t, \qquad g_{\mathrm{term}} \triangleq w(\tau) \nabla_\theta \ell_L .\] Their difference is \[\Delta g \triangleq g_{\mathrm{all}} - g_{\mathrm{term}} = w(\tau) \sum_{t=1}^{L-1} \nabla_\theta \ell_t,\] which corresponds to the omitted gradients from intermediate verifier-rejected responses.

2.1pt 2.3em

Remark. Proposition 4 should be interpreted as a bias-variance motivation rather than an optimizer-equivalence result. Terminal-only retention introduces bias relative to the full-trajectory objective by omitting intermediate turns as imitation targets. Under our stop-on-success protocol, however, those turns are verifier-rejected attempts. When their gradients are noisy or misaligned with the terminal correction signal, terminal-only retention can reduce variance and improve credit assignment. This approximation is therefore justified by the protocol structure.

Empirical results. On MATH500 with Qwen2.5-3B-Instruct, we test this interpretation by comparing DRIFT with an all-turn variant that uses the same offline trajectories and trajectory weights, but applies each weight to every response in the trajectory. As shown in Figure 4, terminal-only supervision reaches higher accuracy and yields a smoother optimization curve under the same training schedule. This supports the prediction of Proposition 4: in stop-on-success trajectories, retroactively assigning large weights to verifier-rejected intermediate responses can amplify misaligned gradients, whereas concentrating supervision on the terminal response provides a better bias-variance tradeoff in practice.

5 Experiments↩︎

5.1 Experimental Setup↩︎

Training setup. We employ Qwen2.5-3B-Instruct [2] and Llama3.1-8B-Instruct [16] as the base model and train it on the MATH subset of the MetaMathQA [17] dataset. We introduce more training details in Section 10.2.

Benchmark. All evaluations are performed using greedy decoding. Since all methods are trained on the MATH subset of MetaMathQA, we organize the evaluation benchmarks into two groups: math reasoning benchmarks from the same broad task family, and out-of-domain general reasoning benchmarks. This split tests whether learned multi-turn correction transfers beyond the training domain rather than only improving MATH-style problem solving. We detail the benchmarks and experimental settings in Section 10.3.

Metrics. We report the cumulative accuracy with a maximum budget of 5 turns, denoted as \(\text{multi}@5\). Generation will terminate immediately if the answer is correct, otherwise it continues up to \(k\) turns. This metric is formally defined as: \[\text{multi}@k \;=\; \frac{1}{N}\sum_{i=1}^N \max_{t\in\{1,\dots,k\}} \mathcal{V}(y_{i,t}).\] where \(\mathcal{V}(\cdot)\) denotes the verifier function and \(N\) is the number of test samples.

5.2 Main Results↩︎

Table [tab:main95cross95benchmark] summarizes the main results. Across most benchmarks, the RL-based method UFO outperforms SFT-based baselines, and DRIFT further improves upon UFO under matched settings. All methods are trained on math data. Single-turn training attains multi@k performance close to multi-turn training on MATH-style benchmarks mainly by improving the first-turn accuracy; however, it does not learn to condition on negative feedback and thus yields little gain on non-math benchmarks. In contrast, multi-turn training explicitly learns correction under negative feedback, leading to substantial improvements and enabling effective multi-turn behavior beyond math tasks. Among multi-turn methods, RL-based approaches remain stronger than SFT-based ones, with the gap most pronounced on non-math benchmarks. By realizing the RL objective with an SFT-style optimization, DRIFT matches or surpasses RL baselines on most benchmarks. In small discrete action spaces (e.g., multiple-choice), UFO significantly suppresses repetition, enabling models to guess via elimination when intrinsic capability is lacking, notably on Llama-3.1-8B-Instruct. We illustrate examples of blind guessing and briefly discuss them in Section 11.3.

5.3 Training efficiency comparison↩︎

We compare training efficiency for Qwen2.5-3B-Instruct and Llama3.1-8B-Instruct on two hardware configurations: 4\(\times\) NVIDIA A800 (80GB) GPUs and 4\(\times\) NVIDIA H20 (96GB) GPUs, and report GPU time. GPU time measures the end-to-end wall-clock time, including rollout latency for methods that perform rollouts (both UFO and DRIFT). We run 200 training steps with a global batch size of 128 distributed over 4 GPUs. As the number of turns increases, SFT-5Turn and DRIFT incur only a small increase in GPU time, while UFO-5Turn becomes substantially slower due to the growing cost of multi-turn rollouts. Overall, DRIFT achieves higher training-time efficiency than UFO, and its scaling behavior remains close to SFT-style training as the turn count grows.

5.4 Turn-by-turn performance comparison↩︎

In Fig. 6, we report the cumulative accuracy and correction rate (# correct this turn / # wrong in the previous turn) at each turn after training for different methods. Among all methods, DRIFT achieves a significantly higher correction rate in early turns. And the results show that single-turn training primarily improves the turn 1 accuracy, with only marginal gains in later turns. In particular, the model trained with single-turn PPO is almost unable to continue correcting its answers under multi-turn interactions. Among multi-turn training methods, RL-based approaches exhibit stronger turn-by-turn performance than SFT-based ones. However, as shown in Fig. 7, RL-based methods are substantially less training-efficient. In contrast, DRIFT achieves multi-turn improvements that match or surpass RL-based methods, while retaining a training efficiency comparable to SFT-based methods.

5.5 Impact of hyperparameters↩︎

We analyze the effect of different hyperparameters \(\gamma\) and \(\beta\), as shown in Fig. 5. To compare convergence behaviors across different \(\gamma\) values, we define the solved rate as the ratio of the cumulative success rate at a given turn to the final success rate achieved by the end of the episode. This metric indicates the fraction of ultimately solvable problems that are solved by the current turn. The results demonstrate that a smaller \(\gamma\) encourages the model to solve problems in fewer turns, aligning with the discount factor’s role in incentivizing early success.

For different \(\beta\) values, we report the cumulative accuracy at each turn in Fig. 5 (middle). The results show that a moderate \(\beta\) (we use \(\beta=0.1\)) performs best: larger \(\beta\) leads to weaker exponential tilting and smaller multi-turn gains, while smaller \(\beta\) makes the weights overly concentrated and degrades stability. This matches the role of \(\beta\) in Theorem 1 (controlling the sharpness of the tilted distribution) and the stability trade-off discussed in Proposition 5.

Although we train with negative feedback string \(f=\) “Incorrect. Please think again.”, the learned correction behavior is not sensitive to the exact wording. As shown in Fig. 5 (right), different feedback strings \(f\) have only a minor impact on performance and simpler feedback such as “Incorrect” can better elicit the model’s multi-turn capability, suggesting stable performance across different feedback variants.

6 Conclusion & Limitations↩︎

Conclusion. In this work, we introduced DRIFT, a practical framework for multi-turn optimization under lightweight negative feedback. DRIFT decouples trajectory generation from optimization by sampling multi-turn correction trajectories once from a frozen reference policy, assigning each trajectory an exponential weight derived from a KL-regularized multi-turn objective, and training the target model using an importance-weighted supervised fine-tuning objective. This formulation provides a simple and stable alternative to online multi-turn reinforcement learning, avoiding repeated rollouts during training while still optimizing for both accuracy and efficiency across turns. Empirically, under matched settings, DRIFT achieves performance comparable to or better than strong multi-turn RL baselines across math and general-domain benchmarks, while approaching the efficiency profile of standard supervised fine-tuning as the interaction horizon increases.

Limitations. DRIFT is intended for short-horizon, verifier-guided correction with lightweight deterministic feedback. This boundary matches our problem formulation and evaluation protocol: (1) The verifier provides an unambiguous correctness signal, and each episode contains only a small number of correction attempts. Settings with stochastic or preference-based human feedback, open-ended dialogue objectives, or substantially longer-horizon interactive planning require additional modeling of feedback uncertainty, credit assignment, and exploration, and are left for future work. (2) A second limitation comes from offline rollout coverage. Because DRIFT samples correction trajectories from a fixed reference policy and does not alternate rollout collection with policy optimization, it can only upweight useful behaviors that appear in the collected trajectories. It may therefore miss strategies that online RL could discover through repeated exploration. As a simple initial attempt to reduce this limitation, Appendix 10.5 studies a two-stage rollout-refresh variant that regenerates rollouts after an intermediate DRIFT checkpoint and obtains a modest gain over single-stage DRIFT. More systematic rollout-refresh schedules and hybrid offline-online training remain open directions.

Impact Statement↩︎

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

Acknowledgements↩︎

This work was supported in part by the National Natural Science Foundation of China (Grant Nos. 62506319), the Guangdong Basic and Applied Basic Research Foundation (Grant No. 2026A1515030032), the Shenzhen Science and Technology Program (Grant No. JCYJ20250604141031003), and the Pearl River Talent Program of Guangdong Province (Grant No. 2024QN11X069).

7 Related Work↩︎

Multi-turn self-correction under lightweight negative feedback. Large language models (LLMs) exhibit a non-trivial but fragile ability to revise their answers when presented with minimal negative feedback such as “Incorrect, please try again” [10], [12]. However, this behavior is unreliable. Extended multi-turn contexts can cause models to drift or “get lost” as dialogue history grows [7], [8], and intrinsic self-correction may degrade answer quality, produce superficial edits, or amplify errors [11], [18]. Beyond prompting-based elicitation, a line of work studies iterative improvement via self-critique and refinement, for example self-refinement and reflection-style agents [19], [20]. These approaches can help, but they typically do not optimize a principled trajectory-level objective and they remain sensitive to prompting, verification, and context accumulation. Empirically, multi-step reasoning can sometimes be elicited by generic “try again” messages [12], but the gains vary by task and model family, motivating learning objectives that directly target multi-turn success rather than only first-turn correctness. These observations align with our setting, a deterministic multi-turn correction protocol with lightweight negative feedback, and motivate optimizing the trajectory-level outcome that DRIFT is designed to improve.

Learning from correction trajectories and multi turn fine tuning. A common approach to improve revision behavior is to collect correction traces and apply supervised fine-tuning (SFT) on multi-turn demonstrations [21]–[23]. Self-training and bootstrapping methods such as STaR generate candidate solutions or rationales and then distill them into the model via SFT, improving problem solving with limited additional supervision [21], [23]. For interactive correction, offline multi-turn SFT can teach models to condition on feedback, but it can also suffer from distribution shift. The model trains on traces whose early turns are produced by a different policy than the one deployed at test time, so correction quality may drop when the first-turn response is generated by the updated model [10]. Moreover, SFT can exhibit behavioral collapse in which the model over-optimizes first-attempt accuracy while performing shallow or non-committal revisions in later turns [10]. Recent analyses further suggest that intermediate turns may contain noisy or spurious reasoning paths, complicating stable learning from full trajectories [11]. Complementary to multi-turn trace SFT, recent offline RL formulations for conversational policies recast short-horizon interaction learning as return-weighted fine-tuning, enabling direct use of scalar rewards on logged dialogue trajectories [24]. In contrast to unweighted multi-turn SFT on fixed traces, DRIFT retains SFT-like optimization but reweights offline trajectories to approximate the intended multi-turn objective, thereby addressing these failure modes while staying compatible with efficient fine-tuning pipelines.

KL regularized policy optimization and return weighted fine tuning. To optimize multi-turn outcomes, reinforcement learning (RL) and RLHF update the policy from rollouts and reward signals, often with KL regularization to stay close to a reference model [9], [13], [14], [25]. In LLM alignment, PPO-style RLHF and sequence-level variants such as GRPO [6], [15] can be viewed as KL-regularized distribution matching, while single-turn alternatives such as DPO replace online RL with closed-form preference objectives [4]. More broadly, RL-as-inference and reward-weighted or advantage-weighted regression show that KL-regularized objectives induce exponential reweighting of actions or trajectories [26]–[29], enabling offline optimization via importance weighting. Related connections also appear in structured prediction, where reward augmented maximum likelihood (RAML) integrates task rewards into maximum likelihood training through an exponentiated reward distribution [30]. Building on these views, recent alignment methods derive reward-driven reweighted SFT objectives through variational or bound-tightening analyses, including VAR [31] and importance-weighted SFT [32], and successive policy reweighting schemes that target RL objectives with SFT-like compute [33]. Complementary analyses study when RLVR with answer-only rewards can still promote correct reasoning and what normalization or supervision improves stability [34]. Although online multi-turn RL can be effective [9], [10], its rollout cost grows quickly with horizon because each update requires full multi-turn trajectories. DRIFT applies these KL-regularized, weighted-likelihood principles at the trajectory level using offline rollouts under a reference policy, prompt-level weight normalization, and terminal-only supervision, decoupling rollout from optimization and achieving an RL-motivated target with SFT-like efficiency.

8 Additional theoretical analyses↩︎

8.1 Estimation Stability and Sample Complexity↩︎

The stability of the gradient estimator relies on the accuracy of the normalized weights \(\widehat{w}(\tau) = \exp(R(\tau)/\beta)/\widehat{Z}(x)\). Since the Monte Carlo estimate \(\widehat{Z}(x)\) appears in the denominator, small estimation errors can be amplified, particularly when the partition function is small. Here, we analyze the sample complexity required to bound this error.

Stability of Normalized Weights. Bounding the error of \(\widehat{Z}\) is a prerequisite for stable training. Since \(Z(x) \ge m_\beta > 0\), the function \(f(z) = 1/z\) is Lipschitz continuous on \([m_\beta, \infty)\) with constant \(1/m_\beta^2\). Consequently, an estimation error \(|\widehat{Z}-Z| \le \epsilon\) propagates to the importance weights as: \[\left| \widehat{w}(\tau) - w(\tau) \right| \le M_\beta \left| \frac{1}{\widehat{Z}} - \frac{1}{Z} \right| \le \frac{M_\beta}{m_\beta^2} \epsilon. \label{eq:estimation95error}\tag{8}\] This confirms that sufficiently large \(K\) ensures the convergence of the empirical weights to their theoretical values.

The Regularization-Complexity Trade-off. Eq. ?? reveals that the sample complexity is governed by the range \(\Delta_\beta\). In the weak regularization regime (\(\beta \to 0\)), \(\Delta_\beta\) grows exponentially, indicating that in the worst case, the required \(K\) to maintain a fixed precision \(\epsilon\) scales with \(\exp(2R_{\max}/\beta)\). This theoretical bound justifies our use of a moderate \(\beta\): it balances the sharpness of the distribution matching against the need for sample efficiency, preventing the gradient estimator from being dominated by high-variance sampling noise.

9 Proofs.↩︎

9.1 Proof of Thm 1↩︎

9.2 Proof of Thm 2↩︎

9.3 Proof of Lemma 1↩︎

9.4 Proof of Lemma 2↩︎

9.5 Proof of Thm 3↩︎

9.6 Proof of Prop. 4↩︎

9.7 Proof of Prop. 5↩︎

10 Experiments↩︎

10.1 Baseline↩︎

SFT-5Turn. We fine-tune the model with SFT on the same offline correction trajectories as DRIFT, with up to five turns per example. Since in our setting only the final-turn response can be correct, we follow DRIFT and train only on the final-turn response conditioned on the full preceding dialogue context. This baseline is equivalent to DRIFT with all trajectory weights set to 1, and thus serves as the unweighted ablation.

STaR-2Turn. We implement STaR-2Turn [21] as a two-turn self-training baseline adapted to our correction protocol. For each training prompt, we first sample a turn-1 response from the current model and evaluate it using the same verifier as in our main setup. If the turn-1 response is incorrect, we append the fixed lightweight negative feedback and prompt the model to produce a turn-2 revision. We keep only those trajectories whose turn-2 response is verified as correct, and then fine-tune the model with SFT on the turn-2 responses conditioned on the full two-turn context. Under our setting, only the final-turn response can be correct, so we do not fit the turn-1 responses. This yields an implementation of STaR that bootstraps training data from verified two-turn corrections.

SCoRe-2Turn. SCoRe [10] studies learning from corrective feedback with verifiable rewards in a two-attempt setting. The model produces an initial answer, receives lightweight negative feedback if it is incorrect, and then generates a second-turn revision. Training uses KL regularization to control deviation from a reference policy and encourages successful correction from an incorrect first attempt to a correct second attempt. In our comparisons, we treat SCoRe as a two-turn training approach under this fixed correction protocol.

UFO-5Turn. UFO [12] considers multi-turn trial-and-error prompting and training under generic negative feedback. The model repeatedly attempts the same problem, and after each incorrect attempt it receives a fixed message such as “Incorrect, please try again” before producing the next response. The process stops once a correct answer is obtained or a maximum turn budget is reached. In our setting, we refer to UFO with a maximum of five turns as UFO-5Turn.

10.2 Setup↩︎

During the offline trajectory generation phase, we set the number of rollouts to \(K=16\) with a sampling temperature of \(1.0\). The maximum trajectory length is restricted to \(T=5\), and the maximum number of new tokens is set to 512. In the optimization phase, we utilize a global batch size of 128 and train the model for 200 steps. The hyperparameters are configured as \(\beta=0.1\) and \(\gamma=0.9\), and we set the repetition penalty coefficient to \(\lambda=0.5\).

10.3 Benchmark↩︎

We evaluate multi-turn correction on benchmarks spanning mathematical reasoning and general-domain knowledge and science reasoning. Because training uses only the MATH subset of MetaMathQA, we treat the mathematical reasoning benchmarks as same-family evaluations and the general reasoning benchmarks as out-of-domain evaluations that test whether correction behavior transfers beyond the training task. Table 1 summarizes the benchmark scale.

Mathematical reasoning. We use MATH [35] as the primary math benchmark with competition-style problems that require multi-step derivations. We also report MATH500 [35], a 500-problem evaluation subset commonly used for faster iteration. In addition, we include TheoremQA (ThmQA) [36], which tests theorem-driven problem solving on STEM questions and emphasizes selecting and applying appropriate theorems.

General reasoning. To evaluate cross-domain transfer, we include MMLU-Redux (MMLU-R) [37], a cleaned and re-annotated version of MMLU designed to reduce ambiguity and labeling errors. We further evaluate on MMLU-Pro (MMLU-P) [38], which increases difficulty by expanding answer options and filtering trivial items. Finally, we use GPQA-diamond [39], a high-quality and challenging subset of GPQA that focuses on graduate-level science questions.

10.4 Additional Model Results↩︎

To further evaluate whether DRIFT scales to a stronger backbone, we additionally train and evaluate Qwen2.5-7B-Instruct under the same protocol. Table [tab:qwen25957b95results] shows that DRIFT improves the all-benchmark average from 64.8% to 68.3%, and slightly outperforms the online multi-turn RL baseline while retaining the largest gains on MATH and MATH500.

2.1pt 2.1em

10.5 A Simple Rollout-Refresh Variant↩︎

A natural limitation of a single offline rollout collection is coverage: if the reference policy does not produce a useful correction trajectory, importance weighting cannot recover it during optimization. To test whether this limitation can be partially mitigated without fully reverting to online RL, we run a simple two-stage rollout-refresh variant on Qwen2.5-3B-Instruct. We first train DRIFT for 100 steps, use the resulting checkpoint to regenerate correction trajectories, and then continue DRIFT training for another 100 steps. The total optimization budget is kept the same as the single-stage 200-step DRIFT run.

Table [tab:iterative95rollout95refresh] shows that this simple refresh variant improves the all-benchmark average from 60.5% to 61.2%. The gain is modest, but it suggests that periodically refreshing the rollout distribution can partially address the coverage limitation of one-shot offline rollouts. This result is best interpreted as an initial diagnostic rather than a complete replacement for online exploration; designing more systematic rollout-refresh schedules is left for future work.

2.1pt 2.1em

10.6 DRIFT as an Initialization for Online RL↩︎

We further examine whether DRIFT can be used as an initialization before online multi-turn RL. This experiment uses Qwen2.5-3B-Instruct and keeps the total optimization budget fixed at 200 steps. Pure UFO and pure DRIFT are trained for 200 steps, while the hybrid schedules use the first 100 steps for either SFT or DRIFT and the remaining 100 steps for UFO.

Table [tab:drift95rl95initialization] reports all improvements relative to the base model. DRIFT followed by UFO achieves the best all-benchmark average among these schedules and improves over pure UFO by 2.3 points. In contrast, SFT followed by UFO does not improve over pure UFO, suggesting that the benefit is not simply due to adding an offline warm-up stage. These results provide preliminary evidence that DRIFT can serve as a useful warm start for online RL in this verifier-guided setting, although a broader study of hybrid training schedules is left for future work.

1.8pt 2.0em

10.7 Prompt↩︎

In this section, we present the prompts used for rollout and evaluation.

10.8 Ablation on return shaping↩︎

We add a penalty term of the form \(\lambda \left( 1 - \frac{E(\tau)}{L} \right)\) to the trajectory return to encourage generating diverse answers. We report the ablation results in Table 2 and Table [tab:reward95penalty95ablation95cross95benchmark]. After removing the trajectory penalty term, the model exhibits decreases in accuracy, correction rate, and the average number of unique answers. This indicates that the trajectory penalty encourages the model to change its responses when it is wrong, which in turn leads to improved accuracy.

5pt

2.4pt 2.1em

10.9 Learning Curves Across Hyperparameter Settings↩︎

We examine DRIFT’s training dynamics under different choices of the discount factor \(\gamma\) in the shaped trajectory return and the parameter \(\beta\) in the exponential trajectory reweighting with prompt-level normalization. Figures 8 and 9 report multi-turn accuracy over training steps. Across a broad range of \(\gamma\) and \(\beta\), the learning curves improve steadily and remain stable throughout training, without abrupt collapse or divergence. While different settings lead to modest differences in the final accuracy and the amount of fluctuation, the overall trend is consistent, suggesting that DRIFT is not brittle to these hyperparameters in practice.

We plot the training curves under different rollout numbers \(K\) in Fig. 10 while keeping the batch size and training steps fixed. Table 3 summarizes the group-level data distribution for different \(K\). The results show that when \(K\) is small, groups are more likely to become degenerate (all-correct or all-wrong), increasing the proportion of such groups. In this regime, the method gradually approaches an SFT-like update with nearly uniform (all-one) weights, and the performance becomes close to SFT-5Turn, leading to worse-than-expected results. In contrast, larger rollouts typically yield better performance.

11 Case Study↩︎

11.1 Case 1↩︎

We qualitatively analyze a failure case from MATH-500 in Figure 11 to demonstrate DRIFT’s error-recovery capability. The problem requires finding integer values of \(x\) where \(x^3 | 10!\) (implying the exponent constraint \(3c \le 2\) for prime factor 5).

As shown in Figure 11, the base model suffers from behavioral collapse, repeating the logically inconsistent claim that “\(c\) can be 0 or 1” (despite \(3(1) > 2\)) across four turns. Generic feedback fails to break this local optimum. In contrast, DRIFT effectively leverages negative feedback to prune the search space. In Turn 3, it explicitly re-verifies that \(5^3 \nmid 10!\), correcting the exponent count to 1. This illustrates that DRIFT learns to explore alternative reasoning paths rather than merely resampling high-probability errors.

11.2 Case 2↩︎

In Figure 12, we examine a modular arithmetic problem from MATH-500. This case illustrates a different failure mode: hallucinated verification. The base model incorrectly verifies a wrong answer and becomes stuck in a loop because it cannot correct its own arithmetic error despite external feedback. DRIFT, conversely, successfully explores different candidate solutions across turns.

11.3 Case 3: Limitations in Knowledge-Intensive Domains↩︎

he primary contribution of DRIFT is to enable verifiable multi-turn optimization with the training efficiency of SFT. While DRIFT demonstrates strong performance on reasoning-intensive tasks (e.g., MATH) by effectively learning correction policies, we analyze an out-of-distribution case in Figure 13 to distinguish between learning how to correct (strategy) and possessing the domain knowledge required to correct (capability).

In this organic chemistry problem from GPQA:

• Base Model: It selects the correct answer (C) in Turn 1, but relies on a hallucinated mechanism ("electrophilic aromatic substitution") that is chemically impossible for the given aliphatic structure.

• UFO: Lacking the knowledge to verify the answer, it falls into an exhaustive guessing loop, repeating the exact same reasoning text while cycling through options (B \(\to\) A \(\to\) B \(\to\) D \(\to\) C).

• : Its reasoning ("forming a carbocation") remains heuristically plausible but factually flawed for the concerted mCPBA mechanism.

This comparison highlights that DRIFT successfully instills the strategic behavior of multi-turn correction, allowing the model to avoid the blind trial-and-error loops seen in baselines like UFO. However, the accuracy of the reasoning is bounded by the model’s intrinsic knowledge. Since DRIFT focuses on the optimization framework for efficiency rather than domain adaptation, it does not spontaneously inject missing subject matter expertise. This is not a limitation of the DRIFT objective itself, but rather indicates that for knowledge-intensive domains, the learned correction policy must be supported by sufficient underlying capabilities. A promising direction for future work is to couple the efficient correction training of DRIFT with domain-specific continual pre-training, ensuring the model possesses both the policy to correct and the knowledge to correct accurately.

References↩︎