Déploiement à un jeton : Guider l'affinage supervisé des modèles de langage avec une politique de gradient

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

September 30, 2025

papers.authors: Rui Ming, Haoyuan Wu, Shoubo Hu, Zhuolun He, Bei Yu

cs.AI

papers.abstract

Le réglage fin supervisé (SFT) est la méthode prédominante pour adapter les grands modèles de langage (LLMs), mais il peine souvent à généraliser par rapport à l'apprentissage par renforcement (RL). Dans ce travail, nous postulons que cette disparité de performance ne découle pas uniquement de la fonction de perte, mais d'une différence plus fondamentale : le SFT apprend à partir d'un ensemble de données fixe et pré-collecté, tandis que le RL utilise des données on-policy échantillonnées à partir de la politique actuelle. Sur la base de cette hypothèse, nous introduisons le one-token rollout (OTR), un nouvel algorithme de réglage fin qui guide le SFT avec la méthode du gradient de politique. L'OTR reformule le processus d'apprentissage autorégressif en traitant chaque génération de token comme une trajectoire d'apprentissage par renforcement en une seule étape. À chaque étape, il effectue un « rollout » de Monte Carlo en échantillonnant plusieurs tokens candidats à partir de la distribution de la politique actuelle. Le token de vérité terrain des données supervisées est ensuite utilisé pour fournir un signal de récompense à ces échantillons. Guidé par le gradient de politique, notre algorithme réutilise des données supervisées statiques et off-policy en un signal dynamique et on-policy au niveau du token, capturant ainsi les avantages de généralisation de l'apprentissage on-policy tout en évitant le coût élevé de la génération de phrases complètes. À travers des expériences approfondies sur une suite diversifiée de benchmarks complexes couvrant le raisonnement mathématique, la génération de code et le raisonnement dans des domaines généraux, nous démontrons que l'OTR surpasse systématiquement le SFT standard. Nos résultats établissent l'OTR comme une alternative puissante et pratique pour le réglage fin des LLMs et fournissent des preuves convaincantes que la nature on-policy des données est un facteur critique de généralisation, ouvrant ainsi une nouvelle voie prometteuse pour le réglage fin des LLMs.

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.