Resumen Técnico: REFT – Exploración con Despliegue y Diversificación del Primer Token

Declaración del Problema

El Aprendizaje por Refuerzo con Recompensas Verificables (RLVR) ha surgido como un método principal para mejorar las capacidades de razonamiento de los modelos de lenguaje grandes (LLM) sin depender de trayectorias de razonamiento etiquetadas. En los métodos de RLVR agrupados como GRPO y DAPO, una política genera un conjunto finito de despliegues (trazas de razonamiento) para un único prompt, los cuales son puntuados por un verificador automático. La señal de aprendizaje se deriva de las recompensas relativas dentro de este grupo.

Un cuello de botella crítico en este paradigma es la diversidad de despliegues. Si todos los despliegues en un grupo siguen rutas de razonamiento similares o reciben recompensas idénticas, el grupo proporciona poca señal contrastiva para las actualizaciones de la política. Los enfoques existentes intentan ampliar la exploración aumentando la temperatura de muestreo, ramificando en pivotes de razonamiento de alta entropía o seleccionando respuestas finales diversas. Sin embargo, estos métodos a menudo asumen que los tokens de prefijo tempranos son andamiajes fijos o que la exploración valiosa reside únicamente en pasos de razonamiento sustantivos de alta incertidumbre.

Los autores identifican una ineficiencia estructural: el primer token generado después del marcador de razonamiento (<thought>) es típicamente un iniciador de discurso de bajo contenido (por ejemplo, "Primero," "Vamos," "Para"). Las políticas de muestreo estándar exhiben una distribución agudamente picada sobre estos tokens, sesgando fuertemente el grupo de despliegues hacia un único iniciador. Crucialmente, los autores observan que, aunque la probabilidad de estos primeros tokens está altamente concentrada, la correctitud de los despliegues resultantes es casi plana entre los 20 candidatos principales. Esto crea una oportunidad de "baja carga, alta palanca": el primer token tiene una carga semántica mínima pero ejerce una alta palanca distribucional, enrutando toda la continuación subsiguiente hacia diferentes regiones semánticas.

Metodología: REFT

El artículo introduce REFT (Exploración con Despliegue y Diversificación del Primer Token), una intervención mínima en la tubería de RLVR que diversifica el primer token sin alterar la arquitectura del modelo, el verificador, la función de recompensa ni el objetivo de RL.

Mecanismo Central

REFT reemplaza el muestreo estocástico estándar del primer token con una estrategia de muestreo estratificado:

1. Selección de Candidatos: Para un prompt dado $x$, se evalúa la distribución de la política en la posición del primer token. Los $N$ tokens semánticos válidos superiores (excluyendo espacios en blanco y tokens especiales) se identifican como el conjunto candidato $F_N(x)$.
2. Muestreo Uniforme: De $F_N(x)$, se muestrean $K$ primeros tokens distintos uniformemente sin reemplazo para formar un conjunto $S_K(x)$.
3. Asignación de Presupuesto: El tamaño total del grupo de despliegues $G$ se divide equitativamente entre los $K$ tokens seleccionados. Para cada primer token seleccionado $f \in S_K(x)$, el modelo genera $G/K$ continuaciones condicionadas a $f$, utilizando el decodificador base sin cambios para el resto de la secuencia.
4. Entrenamiento: El grupo diversificado resultante es procesado por el algoritmo estándar de RLVR (por ejemplo, GRPO o DAPO) sin correcciones de ponderación de importancia para la discrepancia en el muestreo del primer token.

Decisiones Clave de Diseño

• Baja Carga, Alta Palanca: El método apunta a una posición donde el modelo está "artificialmente seguro" (alta probabilidad para un iniciador específico) pero donde la correctitud es estable entre las alternativas.
• Sin Ruido Global: A diferencia de aumentar la temperatura, que aplanaría las distribuciones en toda la trayectoria (potencialmente dañando pasos de razonamiento críticos posteriores), REFT inyecta diversidad solo en el primer token, preservando la integridad de la cadena de razonamiento subsiguiente.
• Reemplazo Directo: REFT no requiere cambios en el objetivo de RL subyacente, la estimación de ventaja ni la lógica del verificador.

Contribuciones Clave

1. Insight Diagnóstico: Los autores demuestran que la distribución del primer token en modelos de razonamiento está agudamente concentrada pero desacoplada de la correctitud. Muestran que los primeros tokens de baja probabilidad dentro de los 20 principales pueden conducir a despliegues correctos tan a menudo como el token más probable, siempre que la continuación se muestree normalmente.
2. Efecto de Enrutamiento: Los experimentos confirman que diferentes primeros tokens enrutan al modelo hacia regiones de continuación semántica distintas. El muestreo uniforme de los 20 primeros tokens principales produce una diversidad semántica significativamente mayor en la restante continuación (tras eliminar el primer token) en comparación con el muestreo estándar.
3. Algoritmo REFT: Un método simple y computacionalmente eficiente para diversificar grupos de despliegues reasignando el presupuesto fijo de despliegues entre iniciadores de primer token plausibles.
4. Mitigación del Sobre-Crédito: El artículo muestra que las actualizaciones estándar de GRPO tienden a agudizar la preferencia por el primer token número 1 a lo largo de los pasos de entrenamiento, incluso cuando el verificador no apoya fuertemente esa preferencia. REFT mitiga esto asegurando la cobertura de múltiples iniciadores plausibles, evitando que la política otorgue sobre-crédito a una preferencia de discurso específica.

Resultados Experimentales

Los autores evaluaron REFT en cuatro modelos base (Qwen2.5-0.5B, 3B, 7B y Llama3.2-3B) y dos objetivos de RLVR (GRPO y DAPO) en tres conjuntos de datos (GSM8K, BigMath-Easy, BigMath).

• Mejoras de Rendimiento: REFT mejora consistentemente el rendimiento agregado en las métricas Pass@1, Pass@8 y Pass@64.
  • En GSM8K, GRPO y DAPO aumentados con REFT superan a sus líneas base.
  • En el "Promedio Matemático" (promediando MATH-500, AIME 2024/2025, AMC 2023), las ganancias son particularmente pronunciadas, lo que sugiere que REFT es efectivo cuando la exploración más amplia es crítica.
  • Se observan mejoras incluso en la escala más pequeña (0.5B de parámetros), indicando que el método no depende de una gran capacidad del modelo.
• Análisis de Diversidad:
  • Diversidad de Entrenamiento: REFT produce los grupos de despliegues semánticamente más diversos durante el entrenamiento.
  • Concentración en Inferencia: Paradójicamente, los modelos entrenados con REFT producen salidas más concentradas en tiempo de inferencia, lo que sugiere que el método ayuda a la política a aprender a converger en patrones exitosos de manera más efectiva después de explorar rutas diversas.
  • Descomposición de Varianza Cero: REFT reduce la fracción de grupos "todos incorrectos" (grupos de varianza cero donde no se encuentra ninguna trayectoria correcta) sin aumentar la saturación de "todos correctos", indicando una exploración más útil.
• Eficiencia: REFT introduce una sobrecarga de tiempo de ejecución negligible. En algunos casos, reduce ligeramente las horas de GPU debido a longitudes promedio de finalización más cortas resultantes de una mayor precisión.

Significado y Afirmaciones

El artículo afirma que REFT ofrece una "inversión simple de la intuición habitual de exploración". En lugar de buscar diversidad en bifurcaciones de razonamiento de alta entropía e inciertas, una diversidad útil puede recuperarse de elecciones de prefijo de baja carga que están fuertemente sesgadas pero débilmente vinculadas a la correctitud.

Los autores enfatizan que REFT no es un nuevo algoritmo de RL, sino un reemplazo dirigido para la primera decisión del muestreador de despliegues. Su significado radica en:

• Eficiencia: Amplía la señal de entrenamiento sin aumentar el presupuesto de despliegues ni el costo computacional.
• Robustez: Mejora tanto la precisión de una sola muestra (Pass@1) como la probabilidad de encontrar una solución correcta dentro de un presupuesto finito (Pass@8/64), abordando la compensación a menudo vista en RLVR donde mejorar Pass@1 estrecha la distribución accesible.
• Generalizabilidad: El método funciona en diferentes escalas de modelos y objetivos de RLVR, lo que sugiere que el fenómeno de "enrutamiento del primer token" es una característica fundamental de los modelos de razonamiento actuales que puede ser explotada para un mejor entrenamiento.

El artículo concluye que el primer token después del marcador de razonamiento es un punto de intervención inesperadamente barato y efectivo para la exploración de RLVR, permitiendo a los modelos explorar regiones de continuación más amplias mientras mantienen la estabilidad de la traza de razonamiento.