ワントークン・ロールアウト：ポリシー勾配を用いたLLMの教師ありファインチューニングのガイダンス
One-Token Rollout: Guiding Supervised Fine-Tuning of LLMs with Policy Gradient
September 30, 2025
著者: Rui Ming, Haoyuan Wu, Shoubo Hu, Zhuolun He, Bei Yu
cs.AI
要旨
教師ありファインチューニング（SFT）は、大規模言語モデル（LLM）を適応させるための主要な手法であるが、強化学習（RL）と比較して一般化に苦戦することが多い。本研究では、この性能差が単に損失関数の違いだけでなく、より根本的な差異に起因すると仮定する。すなわち、SFTは事前に収集された固定データセットから学習するのに対し、RLは現在のポリシーからサンプリングされたオンラインデータを利用する。この仮説に基づき、我々はポリシー勾配法を用いてSFTを導く新しいファインチューニングアルゴリズムであるワントークンロールアウト（OTR）を提案する。OTRは、各トークン生成を単一ステップの強化学習軌道として扱うことで、自己回帰的学習プロセスを再構築する。各ステップにおいて、現在のポリシーの分布から複数の候補トークンをサンプリングし、モンテカルロ「ロールアウト」を実行する。教師データからの正解トークンは、これらのサンプルに対する報酬信号として使用される。ポリシー勾配に導かれた我々のアルゴリズムは、静的でオフラインデータである教師データを、トークンレベルでの動的でオンラインデータに変換し、オンラインデータ学習の一般化の利点を捉えつつ、全文生成の高コストなオーバーヘッドを回避する。数学的推論、コード生成、一般領域推論にわたる多様な挑戦的なベンチマークでの広範な実験を通じて、OTRが標準的なSFTを一貫して上回ることを実証する。我々の知見は、OTRをLLMのファインチューニングにおける強力で実用的な代替手法として確立し、データのオンラインデータ性が一般化の重要な要因であることを示唆する。これにより、LLMのファインチューニングに向けた新たな有望な方向性を提供する。
English
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.