Beyond Token-Level Policy Gradients for Complex Reasoning with Large Language Models

1. El Problema: Desajuste de Granularidad en el Razonamiento

Los métodos actuales de optimización de políticas (como PPO, GRPO y DAPO) para el ajuste fino de Modelos de Lenguaje Grandes (LLMs) operan a nivel de token individual. Tratan la generación de cada token como una acción aislada en el proceso de aprendizaje por refuerzo.

Sin embargo, en tareas de razonamiento complejo (como matemáticas o programación), las decisiones semánticas no ocurren token a token, sino a través de bloques semánticos coherentes (por ejemplo, definir una variable completa, escribir una ecuación o formular una función).

• Limitación actual: La optimización token a nivel fragmenta estas acciones semánticas, rompiendo la coherencia interna de los pasos de razonamiento.
• Consecuencia: Existe un desajuste fundamental entre la granularidad de la optimización (token) y la naturaleza de la tarea (bloques semánticos), lo que puede llevar a una asignación de crédito ineficiente y a una inestabilidad en el entrenamiento debido a la alta volatilidad de las razones de muestreo de importancia.

2. Metodología: Optimización de Gradiente de Política Multi-token (MPO)

Los autores proponen MPO (Multi-token Policy Gradient Optimization), un marco que extiende el concepto de predicción multi-token (MTP) al ámbito de la optimización de políticas.

Conceptos Clave:

• Acciones de Bloque: En lugar de optimizar la probabilidad de un solo token siguiente, MPO agrupa secuencias de K tokens consecutivos como una única "acción semántica" unificada.
• Razón de Muestreo de Importancia Agregada:
  • En lugar de usar la razón estándar $r_t = \frac{\pi_\theta(a_t|s_t)}{\pi_{old}(a_t|s_t)}$, MPO calcula una razón agregada sobre un bloque de K tokens.
  • Para evitar que el producto de múltiples razones (que aumenta drásticamente la varianza) inestabilice el entrenamiento, los autores proponen una formulación alternativa basada en una suma logarítmica ponderada:
    $$\tilde{R}^{(K)}_{i,t}(\theta) = \exp \left( \sum_{n=1}^{K} \beta_n \log r_{i,t+n}(\theta) \right)$$
    Donde $\beta_n$ son pesos que decaen con la distancia (controlados por un hiperparámetro $\lambda$), dando más peso a los tokens inmediatos y menos a los más lejanos dentro del bloque.
• Arquitectura y Entrenamiento:
  • Utiliza módulos de Predicción Multi-Token (MTP) (inspirados en DeepSeek-V3) que predicen los siguientes K tokens en paralelo.
  • Estrategia de Calentamiento (Warm-up): Antes del entrenamiento por refuerzo, los módulos MTP se "calientan" mediante una etapa de ajuste supervisado para asegurar que las predicciones de múltiples tokens sean precisas.
  • Compatibilidad: MPO solo modifica la razón de muestreo de importancia dentro de los marcos existentes de gradiente de política (como PPO), manteniendo la compatibilidad con las tuberías de entrenamiento actuales.

3. Contribuciones Principales

1. Propuesta de MPO: Un nuevo marco que permite optimizar la política sobre unidades de razonamiento estructuradas (bloques de tokens) en lugar de tokens aislados.
2. Integración en Post-Entrenamiento: Extiende el concepto de predicción multi-token (usualmente usado en pre-entrenamiento o inferencia) al entrenamiento por refuerzo (RLHF), revelando su potencial para mejorar la coherencia del razonamiento.
3. Superación de Límites: Demuestra empíricamente que superar la granularidad del token es crucial para tareas intensivas en razonamiento, logrando mejoras consistentes sobre los métodos basados en tokens.

4. Resultados Experimentales

Los autores evaluaron MPO en tres modelos base (Llama3.2-1B, DeepSeek-Distilled-Qwen2.5-1.5B y 7B) y en tres benchmarks:

• GSM8K y MATH (Razonamiento Matemático): MPO superó consistentemente a los baselines (PPO, GRPO, DAPO) en todos los tamaños de modelo. Las mejoras fueron particularmente notables en el conjunto de datos MATH (nivel competición), donde la estructura del razonamiento es más compleja.
• HumanEval (Generación de Código): MPO también obtuvo mejoras en la síntesis de programas, demostrando que la dependencia semántica de múltiples pasos es crítica también en la codificación.
• Estabilidad del Entrenamiento:
  • MPO redujo significativamente la varianza de la razón de muestreo de importancia y la fracción de recorte (clip fraction) en comparación con PPO.
  • Se identificó un equilibrio óptimo: incorporar aproximadamente un 10-20% de señal de información de múltiples tokens (pesos acumulados de los módulos MTP) ofrece el mejor rendimiento. Un exceso de peso introduce ruido y sesgo.
• Eficiencia: Aunque MPO añade módulos adicionales, el costo computacional es manejable. El tiempo de entrenamiento es aproximadamente un 30% más rápido que GRPO (debido a la eficiencia en la asignación de crédito y menor necesidad de recortes), y el uso de memoria es comparable.

5. Significado e Impacto

Este trabajo marca un cambio metodológico en el post-entrenamiento de LLMs. Sugiere que para tareas de razonamiento complejo, la granularidad de la unidad de decisión debe alinearse con la estructura semántica del lenguaje (bloques, ecuaciones, funciones) en lugar de la unidad sintáctica mínima (token).

• Implicación Teórica: Cuestiona la suposición de que las actualizaciones token a nivel son suficientes para capturar dependencias de largo alcance en tareas estructuradas.
• Dirección Futura: Abre la puerta a investigar estrategias de aprendizaje de políticas que operen a niveles de granularidad más gruesos y semánticamente significativos, mejorando la coherencia global y la capacidad de planificación de los modelos de lenguaje.

En resumen, MPO demuestra que tratar secuencias de tokens como acciones semánticas unificadas mejora tanto la estabilidad del entrenamiento como la capacidad de razonamiento de los LLMs, superando las limitaciones inherentes de los métodos de gradiente de política tradicionales.