¿Diversidad o Precisión? Un Análisis Profundo de la Predicción del Siguiente Token

Resumen

Los avances recientes han demostrado que el aprendizaje por refuerzo (RL) puede mejorar sustancialmente las capacidades de razonamiento de los modelos de lenguaje grandes (LLMs). Sin embargo, la eficacia de dicho entrenamiento con RL depende críticamente del espacio de exploración definido por la distribución de salida de tokens del modelo preentrenado. En este artículo, revisitamos la función de pérdida de entropía cruzada estándar, interpretándola como un caso específico de optimización de gradiente de política aplicada dentro de un episodio de un solo paso. Para estudiar sistemáticamente cómo la distribución preentrenada moldea el potencial de exploración para el RL posterior, proponemos un objetivo de preentrenamiento generalizado que adapta los principios de RL *on-policy* al aprendizaje supervisado. Al enmarcar la predicción del siguiente token como un proceso de decisión estocástico, introducimos una estrategia de configuración de recompensas que equilibra explícitamente la diversidad y la precisión. Nuestro método emplea un factor de escala de recompensa positivo para controlar la concentración de probabilidad en los tokens de verdad fundamental y un mecanismo consciente del rango que trata de manera asimétrica los tokens negativos de rango alto y bajo. Esto nos permite remodelar la distribución de salida de tokens preentrenada e investigar cómo proporcionar un espacio de exploración más favorable para el RL, mejorando en última instancia el rendimiento de razonamiento de extremo a extremo. Contrario a la intuición de que una entropía de distribución más alta facilita una exploración efectiva, encontramos que imponer un prior orientado a la precisión produce un espacio de exploración superior para el RL.

English

Recent advancements have shown that reinforcement learning (RL) can substantially improve the reasoning abilities of large language models (LLMs). The effectiveness of such RL training, however, depends critically on the exploration space defined by the pre-trained model's token-output distribution. In this paper, we revisit the standard cross-entropy loss, interpreting it as a specific instance of policy gradient optimization applied within a single-step episode. To systematically study how the pre-trained distribution shapes the exploration potential for subsequent RL, we propose a generalized pre-training objective that adapts on-policy RL principles to supervised learning. By framing next-token prediction as a stochastic decision process, we introduce a reward-shaping strategy that explicitly balances diversity and precision. Our method employs a positive reward scaling factor to control probability concentration on ground-truth tokens and a rank-aware mechanism that treats high-ranking and low-ranking negative tokens asymmetrically. This allows us to reshape the pre-trained token-output distribution and investigate how to provide a more favorable exploration space for RL, ultimately enhancing end-to-end reasoning performance. Contrary to the intuition that higher distribution entropy facilitates effective exploration, we find that imposing a precision-oriented prior yields a superior exploration space for RL.

¿Diversidad o Precisión? Un Análisis Profundo de la Predicción del Siguiente Token

Diversity or Precision? A Deep Dive into Next Token Prediction

Resumen

Support