PaperlayerView exact revision on arXiv

Source-linked AI summary

One-Token Rollout: Guiding Supervised Fine-Tuning of LLMs with Policy Gradient

Rui Ming, Haoyuan Wu, Shoubo Hu, Zhuolun He, Bei Yu

arXiv:2509.26313v2cs.CL

TL;DR

SFT often generalizes less well than RL, and the paper examines whether their different data regimes help explain that gap. It introduces OTR, which applies policy gradient to token-level rollouts using supervised ground truth as a reward. OTR consistently outperforms SFT across in-domain and out-of-domain benchmarks, while its scalability beyond the evaluated models remains unvalidated.

  • Problem

    SFT often struggles with generalization compared with RL, whose training data is sampled on-policy rather than from SFT’s fixed, pre-collected dataset.

    Source: §1 · p. 1

  • Method

    OTR treats each token generation as a single-step policy-gradient trajectory, sampling candidate tokens from the current policy and using the ground-truth token as a reward signal.

    Source: §1 · p. 2, §3.1 · p. 3

  • Results

    OTR consistently outperforms SFT across a wide array of in-domain and out-of-domain benchmarks.

    Source: §4.3 · pp. 6–7, §6 · p. 9

  • Takeaways & Limitations

    The authors’ results provide evidence that on-policy data is a critical factor in fine-tuned models’ generalization performance.

    Source: §1 · p. 2

    The study evaluates models up to 8 billion parameters and text-only tasks, so OTR’s scalability to larger models remains unvalidated.

    Source: LIMITATIONS · p. 9

Abstract

from arXiv · show

Supervised fine-tuning (SFT) is the predominant method for adapting large language models (LLMs), yet it often struggles with generalization compared to reinforcement learning (RL). In this work, we posit that this performance disparity stems not just from the loss function, but from a more fundamental difference: SFT learns from a fixed, pre-collected dataset, whereas RL utilizes on-policy data sampled from the current policy. Building on this hypothesis, we introduce one-token rollout (OTR), a novel fine-tuning algorithm that guides SFT with the policy gradient method. OTR reframes the autoregressive learning process by treating each token generation as a single-step reinforcement learning trajectory. At each step, it performs a Monte Carlo ``rollout'' by sampling multiple candidate tokens from the current policy's distribution. The ground-truth token from the supervised data is then used to provide a reward signal to these samples. Guided by policy gradient, our algorithm repurposes static, off-policy supervised data into a dynamic, on-policy signal at the token level, capturing the generalization benefits of on-policy learning while bypassing the costly overhead of full sentence generation. Through extensive experiments on a diverse suite of challenging benchmarks spanning mathematical reasoning, code generation, and general domain reasoning, we demonstrate that OTR consistently outperforms standard SFT. Our findings establish OTR as a powerful and practical alternative for fine-tuning LLMs and provide compelling evidence that the on-policy nature of data is a critical driver of generalization, offering a promising new direction for fine-tuning LLMs.

1 INTRODUCTION

The introduction frames SFT’s generalization gap with RL as potentially rooted in static versus on-policy training data, not just the loss function. It introduces OTR as a token-level policy-gradient approach and reports consistent gains over SFT across challenging benchmarks.

1 INTRODUCTION

  • SFT often struggles to generalize compared with RL, while also risking degradation of pre-trained general capabilities.The introduction describes this degradation as catastrophic forgetting.Source: §1 · p. 1
  • SFT uses static, pre-collected off-policy demonstrations, whereas RL updates from on-policy data sampled iteratively from the current policy.The introduction therefore treats data regime, alongside the loss function, as a possible source of the performance disparity.Source: §1 · p. 1
  • OTR applies policy gradient at each token position, sampling candidate tokens from the current policy and using the ground-truth token as a reward signal.This turns static supervised data into a token-level on-policy signal.Source: §1 · p. 2
  • OTR consistently outperforms SFT across challenging mathematical, coding, and general reasoning benchmarks, while avoiding full sentence-level on-policy generation.The authors present this as evidence for the importance of on-policy data in generalization.Source: §1 · p. 2
  • The authors posit that on-policy training data is a critical factor in the SFT–RL generalization gap, and cite OTR’s success as evidence for that view.They motivate this perspective by noting that on-policy sampling updates the policy in regions it can already reach.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries describe SFT as autoregressive next-token training on prompt-response data, then introduce policy gradient as direct policy optimization using rewards accumulated over trajectories.

2.1 SUPERVISED FINE-TUNING

  • SFT maximizes the conditional probability of a ground-truth prompt-response sequence by minimizing autoregressive next-token negative log-likelihood.At each step, the model predicts the next token given the prompt and preceding ground-truth tokens.Source: §2.1 · p. 2

2.2 POLICY GRADIENT

  • Policy gradient optimizes a parameterized policy by ascending the expected-reward objective’s gradient, weighting action log-probability updates by trajectory rewards.The state is the prompt plus prior tokens, and the action is the next token; Figure 1 illustrates computational divergence between SFT and OTR.Source: §2.2 · pp. 2–3, Figure 1 · p. 3

3 METHODOLOGY

OTR applies policy gradient to supervised fine-tuning by treating each token generation as a single-step trajectory and using supervised data to provide token-level rewards. This design turns fixed supervised data into an on-policy learning signal.

3.1 FROM POLICY GRADIENT TO ONE TOKEN ROLLOUT

  • OTR treats generation of one token from a state to an action as a complete length-one trajectory, reducing full-trajectory sampling to sampling one action.The policy-gradient sums over timesteps consequently reduce to individual-token terms.Source: §3.1 · p. 3

3.2 TOKEN-LEVEL ROLLOUT AND ON-POLICY REWARD

  • At each sequence position, Monte Carlo rollout samples candidate actions from a temperature-flattened policy to increase exploration.A temperature κ > 1 increases the likelihood of sampling less probable tokens.Source: §3.1 · p. 3, §3.2 · p. 4
  • The ground-truth token rewards sampled matches with 1 and other tokens with β, converting off-policy examples into an on-policy token-level signal.The candidates come from the current policy, while the supervised token supplies their reward; this avoids sentence-level importance-sampling corrections.Source: §3.2 · p. 4

3.3 THE OTR OBJECTIVE FUNCTION

  • The per-token loss averages the Monte Carlo policy-gradient objective over K samples, dynamically weights the ground-truth-token loss by its sampling frequency, and penalizes sampled incorrect tokens.If the ground-truth token is never sampled, its loss contribution is zero; sequence loss averages the per-timestep losses.Source: §3.3 · p. 4
  • OTR focuses optimization on correct predictions already within the model’s reach while gently suppressing plausible alternatives.The paper presents Figure 1 to illustrate how the OTR update differs computationally from standard SFT.Source: §3.3 · p. 5

4 EXPERIMENTS

The experiments compare OTR with SFT across mathematical, code, and general-reasoning benchmarks. OTR outperforms SFT across these domains, and the ablations examine reward settings and training dynamics.

4.1 EXPERIMENT SETTINGS

  • Training uses 5,000 sampled instances from OpenR1-Math-220k and five open-source base models, with SFT as the primary baseline.OTR and SFT have comparable training duration and peak GPU memory usage.Source: §4.1 · p. 5

4.2 EVALUATION

  • Evaluation covers seven math benchmarks, two code benchmarks, and three general-reasoning benchmarks, with distinct decoding setups for base and fine-tuned models.Base models use 5-shot prompts on MATH-500 and GSM8K and greedy decoding; fine-tuned models use 0-shot chat templates and stochastic decoding.Source: §4.2 · p. 6

4.3 RESULTS

  • OTR achieves higher average in-domain math scores than SFT across four model families and higher average OOD code and reasoning performance across all models.The OOD comparison is reported in Table 2.Source: §4.3 · pp. 6–7, Table 2 · p. 6
  • OTR falls below base-model performance in 4 in-domain and 5 OOD instances, versus SFT’s 10 and 7, respectively.On AMC23 with Qwen3-8B, the score drop is 2.97 points for OTR and 7.34 for SFT.Source: §4.3 · p. 7

4.4 ABLATION STUDY

  • Positive rewards for incorrect samples may destabilize training, while OTR converges to more sampled ground-truth tokens than SFT across tested models.With β = 0.1, the ground-truth-token count rises and then collapses; the authors hypothesize that positive rewards can mislead optimization.Source: §4.4 · pp. 7–8, Figure 2 · p. 8
  • Across nearly all domains and models, OTR variants outperform SFT, with β = −0.1 giving the most consistent performance and serving as the default.Tables 3 and 4 report the in-domain and OOD β ablations, respectively.Source: Table 3 · p. 7, Table 4 · p. 7, §4.4 · p. 8
  • Compared with β = 0, β = −0.1 almost universally performs better, supporting an explicit penalty for negatively sampled tokens.The paper views β = 0 as a Monte Carlo approximation of DFT.Source: §4.4 · p. 8

5 RELATED WORK

Prior work simplifies RL objectives and modifies SFT losses to address fine-tuning challenges. OTR is distinguished by changing the training signal through current-policy sampling rather than only reweighting static expert data.

RL for LLMs

  • Some RL fine-tuning methods simplify preference objectives, while GPG expresses RL as weighted maximum likelihood and can match more complex algorithms.The discussion also describes PPO as computationally intensive, unstable, and complex to implement.Source: §5 · pp. 8–9

Improving SFT

  • SFT improvement methods include proximal regularization, RL-based reformulations, and DFT’s reweighting of the ground-truth-token loss.The cited approaches address divergence from the initial policy, connect SFT and RL objectives, or reweight the static expert data.Source: §5 · p. 9
  • Unlike approaches that modify the loss on static expert data, OTR samples from the current policy to turn the training data into a dynamic on-policy signal.Source: §5 · p. 9

6 CONCLUSION

The conclusion presents OTR as a data-centric response to the difference between SFT’s static off-policy data and RL’s dynamic on-policy data. It reports consistent OTR gains on in-domain and OOD benchmarks while retaining computational efficiency.

6 CONCLUSION

  • OTR treats each token generation as a single-step RL trajectory, turning static supervised data into a dynamic on-policy learning signal.Source: §6 · p. 9
  • The token-level approach incorporates an on-policy signal into SFT while maintaining computational efficiency.Source: §6 · p. 9
  • Experiments show OTR consistently outperforms SFT across a wide array of in-domain and OOD benchmarks.Source: §6 · p. 9

LIMITATIONS

The study’s evidence is bounded by model scale, training data, and modality, and its reward mechanism is simple.

LIMITATIONS

  • Experiments use models up to 8 billion parameters and a subset of a mathematics-focused dataset.Scalability to models such as 70B+ remains to be validated.Source: LIMITATIONS · p. 9
  • The investigation is confined to text-only training.The authors identify vision-language tasks as a direction for future work.Source: LIMITATIONS · p. 9
  • The reward mechanism is relatively simple, and the study leaves more sophisticated reward functions for future work.Source: LIMITATIONS · p. 9

A COMPUTATIONAL EFFICIENCY ANALYSIS

OTR has only a slight training-time increase and a marginal peak-memory increase over SFT, with single-token rollouts avoiding full-sentence generation costs.

A COMPUTATIONAL EFFICIENCY ANALYSIS

B ADDITIONAL EXPERIMENT

Under training settings adapted from DFT, OTR outperforms SFT on most benchmarks, testing its generalization across a different data and hyperparameter setup.

B ADDITIONAL EXPERIMENT

  • The experiment tests Qwen2.5-3B and Qwen3-4B using the NuminaMath CoT dataset, sampling 50,000 instances from roughly 860,000 problems.Source: B ADDITIONAL EXPERIMENT · p. 12
  • Both OTR and SFT use identical settings: batch size 256, maximum input length 4096 tokens, and one training epoch.The setup also uses AdamW and a cosine decay schedule with a warm-up ratio of 0.1.Source: B ADDITIONAL EXPERIMENT · p. 12
  • OTR outperforms SFT across the majority of benchmarks under training settings adapted from DFT.The authors say this indicates its generalization benefits are not confined to one data and hyperparameter setting.Source: B ADDITIONAL EXPERIMENT · p. 12

C COMPARISON WITH DYNAMIC FINE-TUNING

OTR matches DFT in a special β = 0 case, while nonzero β adds a penalty on incorrect rollout samples.

C COMPARISON WITH DYNAMIC FINE-TUNING

Related papers

Loading 2509.26313v2…