One-Token Rollout: Begeleiding van Supervised Fine-Tuning van LLM's met Policy Gradient

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

September 30, 2025

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

cs.AI

Samenvatting

Supervised fine-tuning (SFT) is de dominante methode voor het aanpassen van grote taalmodel(len) (LLMs), maar het heeft vaak moeite met generalisatie in vergelijking met reinforcement learning (RL). In dit werk stellen we dat dit prestatieverschil niet alleen voortkomt uit de verliesfunctie, maar uit een fundamenteler verschil: SFT leert van een vast, vooraf verzameld dataset, terwijl RL gebruikmaakt van on-policy data die wordt gesamplet uit het huidige beleid. Op basis van deze hypothese introduceren we one-token rollout (OTR), een nieuw fine-tuning algoritme dat SFT begeleidt met de policy gradient-methode. OTR herformuleert het autoregressieve leerproces door elke token-generatie te behandelen als een single-step reinforcement learning traject. Bij elke stap voert het een Monte Carlo "rollout" uit door meerdere kandidaat-tokens te samplet uit de distributie van het huidige beleid. De ground-truth token uit de supervised data wordt vervolgens gebruikt om een beloningssignaal te geven aan deze samples. Geleid door policy gradient, hergebruikt ons algoritme statische, off-policy supervised data als een dynamisch, on-policy signaal op tokenniveau, waardoor het de generalisatievoordelen van on-policy leren behoudt terwijl het de kostbare overhead van volledige zinsgeneratie omzeilt. Door uitgebreide experimenten op een diverse reeks uitdagende benchmarks, variërend van wiskundig redeneren, codegeneratie en algemeen domeinredeneren, tonen we aan dat OTR consistent beter presteert dan standaard SFT. Onze bevindingen vestigen OTR als een krachtig en praktisch alternatief voor het fine-tunen van LLMs en leveren overtuigend bewijs dat de on-policy aard van data een cruciale drijver is van generalisatie, wat een veelbelovende nieuwe richting biedt voor het fine-tunen van 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.