One-Token Rollout: Steuerung des überwachten Feinabstimmens von LLMs mittels Policy-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

Supervised Fine-Tuning (SFT) ist die vorherrschende Methode zur Anpassung großer Sprachmodelle (LLMs), zeigt jedoch im Vergleich zum Reinforcement Learning (RL) oft Schwächen in der Generalisierung. In dieser Arbeit stellen wir die These auf, dass diese Leistungsunterschiede nicht nur auf die Verlustfunktion zurückzuführen sind, sondern auf einen grundlegenderen Unterschied: SFT lernt aus einem festen, vorab gesammelten Datensatz, während RL On-Policy-Daten nutzt, die aus der aktuellen Policy stammen. Basierend auf dieser Hypothese führen wir One-Token Rollout (OTR) ein, einen neuartigen Fine-Tuning-Algorithmus, der SFT mit der Policy-Gradient-Methode kombiniert. OTR reformuliert den autoregressiven Lernprozess, indem es die Generierung jedes Tokens als einen einstufigen Reinforcement-Learning-Pfad betrachtet. In jedem Schritt führt es einen Monte-Carlo-„Rollout“ durch, indem es mehrere Kandidatentokens aus der Verteilung der aktuellen Policy sampelt. Das Ground-Truth-Token aus den überwachten Daten wird dann verwendet, um ein Belohnungssignal für diese Samples zu liefern. Durch die Anleitung des Policy-Gradient-Algorithmus wandelt unsere Methode statische, Off-Policy-Daten in ein dynamisches, On-Policy-Signal auf Token-Ebene um und nutzt so die Vorteile der Generalisierung durch On-Policy-Lernen, ohne den aufwändigen Overhead der vollständigen Satzgenerierung zu benötigen. Durch umfangreiche Experimente auf einer vielfältigen Reihe anspruchsvoller Benchmarks, die mathematisches Denken, Code-Generierung und allgemeines Domänenverständnis abdecken, zeigen wir, dass OTR durchweg besser abschneidet als Standard-SFT. Unsere Ergebnisse etablieren OTR als eine leistungsstarke und praktische Alternative für das Fine-Tuning von LLMs und liefern überzeugende Beweise dafür, dass die On-Policy-Natur der Daten ein entscheidender Faktor für die Generalisierung ist. Dies eröffnet eine vielversprechende neue Richtung für das Fine-Tuning von 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.