Rollout de Um Token: Orientando o Ajuste Fino Supervisionado de LLMs com Gradiente de Política

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

September 30, 2025

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

cs.AI

Resumo

O ajuste fino supervisionado (SFT) é o método predominante para adaptar grandes modelos de linguagem (LLMs), mas frequentemente enfrenta dificuldades em termos de generalização em comparação com o aprendizado por reforço (RL). Neste trabalho, propomos que essa disparidade de desempenho não se deve apenas à função de perda, mas a uma diferença mais fundamental: o SFT aprende a partir de um conjunto de dados fixo e pré-coletado, enquanto o RL utiliza dados on-policy amostrados da política atual. Com base nessa hipótese, introduzimos o one-token rollout (OTR), um novo algoritmo de ajuste fino que guia o SFT com o método de gradiente de política. O OTR reformula o processo de aprendizado autoregressivo tratando cada geração de token como uma trajetória de aprendizado por reforço de passo único. A cada passo, ele realiza um "rollout" de Monte Carlo ao amostrar múltiplos tokens candidatos da distribuição da política atual. O token verdadeiro dos dados supervisionados é então usado para fornecer um sinal de recompensa a essas amostras. Guiado pelo gradiente de política, nosso algoritmo transforma dados supervisionados estáticos e off-policy em um sinal dinâmico e on-policy no nível do token, capturando os benefícios de generalização do aprendizado on-policy enquanto evita o custoso overhead da geração de frases completas. Por meio de extensos experimentos em um conjunto diversificado de benchmarks desafiadores que abrangem raciocínio matemático, geração de código e raciocínio em domínio geral, demonstramos que o OTR supera consistentemente o SFT padrão. Nossos resultados estabelecem o OTR como uma alternativa poderosa e prática para o ajuste fino de LLMs e fornecem evidências convincentes de que a natureza on-policy dos dados é um fator crítico para a generalização, oferecendo uma nova direção promissora para o ajuste fino de 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.