One-Token Rollout: Guiding Supervised Fine-Tuning of LLMs with Policy Gradient
Ming, Rui, Wu, Haoyuan, Hu, Shoubo, He, Zhuolun, Yu, Bei
–arXiv.org Artificial Intelligence
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. 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. 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. Supervised fine-tuning (SFT) has become a cornerstone for adapting large language models (LLMs) to downstream tasks (Ouyang et al., 2022; Chung et al., 2022; Zhang et al., 2025).
arXiv.org Artificial Intelligence
Oct-1-2025