One-Token Rollout: Guiding Supervised Fine-Tuning of LLMs with Policy
  Gradient
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.
68dfc7cf73e20ab5778419b3
Rui Ming
claimed_verified
2025-10-03T12:56:45.158Z
66d6ce5e0b8483353663a6ce
/avatars/f008af97b87b6288f8f3279e3d0e9107.svg
Rui
user
Yalimu
68dfc7cf73e20ab5778419b4
Haoyuan Wu
claimed_verified
2025-10-05T12:46:36.641Z
6445e7b1b272430bdbf64e80
/avatars/d3e59a3b488f8539966c944bb16f7b90.svg
Haoyuan WU
user
hywu
68dfc7cf73e20ab5778419b5
Shoubo Hu
68dfc7cf73e20ab5778419b6
Zhuolun He
68dfc7cf73e20ab5778419b7
Bei Yu
One-Token Rollout: Guiding Supervised Fine-Tuning of LLMs with Policy
  Gradient
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.
One-token rollout (OTR) enhances supervised fine-tuning of large language models by incorporating policy gradient methods to improve generalization using on-policy data.
supervised fine-tuning large language models reinforcement learning one-token rollout policy gradient method autoregressive learning Monte Carlo rollout on-policy data off-policy data generalization mathematical reasoning code generation domain reasoning