One-Token Rollout: Guidare il Fine-Tuning Supervisionato di LLM con il Gradiente della Politica

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

September 30, 2025

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

cs.AI

Abstract

Il fine-tuning supervisionato (SFT) è il metodo predominante per adattare i grandi modelli linguistici (LLM), ma spesso incontra difficoltà nella generalizzazione rispetto all'apprendimento per rinforzo (RL). In questo lavoro, ipotizziamo che questa disparità di prestazioni non derivi solo dalla funzione di perdita, ma da una differenza più fondamentale: l'SFT apprende da un dataset fisso e pre-raccolto, mentre l'RL utilizza dati on-policy campionati dalla politica corrente. Basandoci su questa ipotesi, introduciamo il one-token rollout (OTR), un nuovo algoritmo di fine-tuning che guida l'SFT con il metodo del gradiente della politica. L'OTR riformula il processo di apprendimento autoregressivo trattando ogni generazione di token come una traiettoria di apprendimento per rinforzo a singolo passo. Ad ogni passo, esegue un "rollout" Monte Carlo campionando più token candidati dalla distribuzione della politica corrente. Il token ground-truth proveniente dai dati supervisionati viene quindi utilizzato per fornire un segnale di ricompensa a questi campioni. Guidato dal gradiente della politica, il nostro algoritmo trasforma dati supervisionati statici e off-policy in un segnale dinamico e on-policy a livello di token, catturando i benefici della generalizzazione dell'apprendimento on-policy evitando il costoso sovraccarico della generazione di frasi complete. Attraverso esperimenti estesi su una suite diversificata di benchmark impegnativi che spaziano dal ragionamento matematico, alla generazione di codice e al ragionamento in domini generali, dimostriamo che l'OTR supera costantemente l'SFT standard. I nostri risultati stabiliscono l'OTR come un'alternativa potente e pratica per il fine-tuning degli LLM e forniscono prove convincenti che la natura on-policy dei dati è un fattore critico per la generalizzazione, offrendo una nuova direzione promettente per il fine-tuning degli 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.