Despliegue de Un Solo Token: Guiando el Ajuste Fino Supervisado de Modelos de Lenguaje con 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

Resumen

El ajuste fino supervisado (SFT, por sus siglas en inglés) es el método predominante para adaptar modelos de lenguaje de gran escala (LLMs), aunque a menudo enfrenta dificultades en la generalización en comparación con el aprendizaje por refuerzo (RL). En este trabajo, proponemos que esta disparidad en el rendimiento no se debe únicamente a la función de pérdida, sino a una diferencia más fundamental: el SFT aprende de un conjunto de datos fijo y previamente recolectado, mientras que el RL utiliza datos generados bajo la política actual. Basándonos en esta hipótesis, introducimos el algoritmo de ajuste fino llamado "one-token rollout" (OTR), que guía el SFT mediante el método de gradiente de política. OTR reformula el proceso de aprendizaje autoregresivo al tratar la generación de cada token como una trayectoria de aprendizaje por refuerzo de un solo paso. En cada paso, realiza un "rollout" de Monte Carlo muestreando múltiples tokens candidatos de la distribución de la política actual. Luego, el token verdadero de los datos supervisados se utiliza para proporcionar una señal de recompensa a estas muestras. Guiado por el gradiente de política, nuestro algoritmo transforma datos supervisados estáticos y fuera de política en una señal dinámica y bajo política a nivel de token, capturando los beneficios de generalización del aprendizaje bajo política mientras evita el costoso proceso de generación de oraciones completas. A través de experimentos exhaustivos en un conjunto diverso de benchmarks desafiantes que abarcan razonamiento matemático, generación de código y razonamiento en dominios generales, demostramos que OTR supera consistentemente al SFT estándar. Nuestros hallazgos establecen a OTR como una alternativa poderosa y práctica para el ajuste fino de LLMs y proporcionan evidencia convincente de que la naturaleza bajo política de los datos es un factor crítico para la generalización, ofreciendo una nueva dirección prometedora para el 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.