Resumen Técnico: POISE (Optimización de Políticas con Estimación de Valor de Estado Interno)

Declaración del Problema

El aprendizaje por refuerzo con recompensas verificables (RLVR) es un motor central para mejorar los Modelos de Razonamiento Avanzado (LRM) en tareas complejas como el razonamiento matemático. Sin embargo, un RLVR estable y eficiente depende en gran medida de una estimación precisa de la línea base para reducir la varianza de los gradientes de la política. Los enfoques existentes incurren en costos computacionales significativos:

• PPO (Optimización de Política Proximal): Requiere entrenar un modelo crítico separado, a escala de LLM, que procesa la secuencia generada completa en cada actualización, duplicando aproximadamente el consumo de memoria y la complejidad de la optimización.
• GRPO (Optimización de Política Relativa por Grupo): Evita un crítico separado estimando las líneas base a partir de la recompensa media de un grupo de trayectorias muestreadas por prompt. Esto "intercambia parámetros por muestras", requiriendo múltiples trayectorias por prompt para mantener una línea base estable. Bajo un presupuesto de cómputo fijo, esto reduce la diversidad de prompts en un lote, inflando la varianza del gradiente. Además, se desperdicia un cómputo sustancial en prompts "degenerados" donde todas las trayectorias reciben recompensas idénticas (ventaja cero).

La pregunta central abordada por este trabajo es: ¿Se puede extraer una línea base efectiva de los cálculos ya realizados durante el entrenamiento de la política, eliminando la necesidad de recursos adicionales?

Metodología: POISE

Los autores proponen POISE (Optimización de Políticas con Estimación de Valor de Estado Interno), un método que aprovecha las representaciones internas del propio modelo de política para estimar la función de valor $V^\pi(x)$ a un costo negligible.

1. Estimación de Valor de Estado Interno

En lugar de entrenar un crítico separado, POISE entrena una sonda ligera para predecir la recompensa verificable esperada directamente a partir de los estados ocultos del modelo de política y las estadísticas de tokens recopiladas durante el paso hacia adelante.

• Características de Entrada: La sonda utiliza tres señales complementarias:
  1. Características a nivel de prompt: Estados ocultos de los últimos tokens del prompt, que capturan la representación del modelo del prompt y la dificultad anticipada.
  2. Características a nivel de trayectoria: Estados ocultos del final de la trayectoria de razonamiento, exponiendo información relevante para el valor no presente en el prompt por sí solo.
  3. Estadísticas de entropía a nivel de token: Características de incertidumbre ligeras.
• Objetivo de Entrenamiento: La sonda se entrena para predecir el valor a nivel de prompt $V^\pi(x)$. Para evitar que el objetivo filtre información sobre la trayectoria específica que se está evaluando, los autores utilizan un objetivo Monte Carlo leave-one-out. Para un prompt $x$ con $K$ trayectorias muestreadas, el objetivo para la trayectoria $i$ es la recompensa promedio de las otras $K-1$ trayectorias.

2. Construcción Cruzada de Trayectorias para Imparcialidad

Un desafío teórico crítico es garantizar que la línea base permanezca condicionalmente independiente de la acción (la respuesta muestreada) para preservar la imparcialidad del gradiente de la política. Si la línea base depende de la misma trayectoria que evalúa, introduce sesgo.

• Solución: POISE emplea una construcción cruzada de trayectorias. Para cada prompt, se muestrean dos trayectorias independientes, $y^{(1)}$ y $y^{(2)}$.
  • La línea base para $y^{(1)}$ se predice utilizando los estados internos de $y^{(2)}$.
  • La línea base para $y^{(2)}$ se predice utilizando los estados internos de $y^{(1)}$.
• Esto asegura que la línea base utilizada para actualizar una trayectoria específica se derive de una muestra independiente, satisfaciendo la condición $E[b(x) \nabla \log \pi(y|x)] = 0$.

3. Integración en la Optimización de Políticas

POISE integra este estimador en un bucle de optimización estilo PPO:

• Entrenamiento en Línea: El estimador de valor se entrena conjuntamente con la política sobre un búfer deslizante de trayectorias recientes. Dado que las características se extraen del paso hacia adelante ya requerido para las actualizaciones de la política, la sobrecarga computacional es negligible.
• Cálculo de la Ventaja: Las ventajas se calculan como $A^{(i)}(x) = R(x, y^{(i)}) - b^{(i)}(x)$, donde $b^{(i)}(x)$ es la línea base cruzada de trayectorias.
• Eficiencia: A diferencia de GRPO, que requiere un grupo grande (por ejemplo, 8 trayectorias) para estimar una línea base, POISE solo requiere un par de trayectorias. Esto permite asignar el presupuesto de cómputo fijo a un mayor número de prompts distintos por lote, reduciendo teóricamente la varianza del gradiente.

Contribuciones Clave

1. Estimación de Línea Base Novel: Introduce un método para derivar líneas base a partir de los estados internos del modelo de política, eliminando la necesidad de un crítico separado a escala de LLM o de grandes grupos de trayectorias.
2. Imparcialidad Teórica: Propone una arquitectura cruzada de trayectorias que garantiza la independencia condicional requerida para gradientes de política imparciales mientras utiliza características condicionadas a la trayectoria.
3. Eficiencia Computacional: Demuestra que una sonda ligera (regresión lineal) puede lograr un rendimiento de estimación de valor comparable o superior al de un crítico separado, requiriendo significativamente menos cómputo.
4. Generalización: Muestra que las señales de estado interno se generalizan a través de diferentes dominios (matemáticas, codificación, uso de herramientas, seguimiento de instrucciones) y escalas de modelos.

Resultados Experimentales

El método se evaluó en Qwen3-4B y DeepSeek-R1-Distill-Qwen-1.5B en benchmarks de razonamiento matemático de nivel olímpico (AMC, AIME, HMMT, BRUMO).

• Rendimiento: POISE logra un rendimiento comparable al de DAPO (un algoritmo basado en GRPO de última generación).
  • En Qwen3-4B, POISE obtuvo una puntuación promedio Avg@32 de 0.500, coincidiendo estrechamente con el 0.508 de DAPO, mientras que lo superó en benchmarks específicos como AMC23 y HMMT25.
  • En DeepSeek-R1-Distill-Qwen-1.5B, POISE mejoró la puntuación promedio a 0.303 en comparación con el 0.296 de DAPO.
• Eficiencia de Entrenamiento: POISE requiere sustancialmente menos tiempo de reloj. En DeepSeek-R1-Distill-Qwen-1.5B, alcanzar un rendimiento comparable tomó 18 horas con POISE frente a 24 horas con DAPO en una sola GPU B200. En Qwen3-4B, el tiempo fue de 36 horas frente a 49 horas.
• Estabilidad: Las normas de los gradientes en POISE fueron consistentemente más bajas que en DAPO, indicando una varianza de gradiente reducida y una optimización más estable.
• Calidad de la Estimación de Valor: La sonda de estado interno logró una correlación de Pearson de 0.870 con las recompensas empíricas, superando a un crítico separado entrenado a escala de LLM (0.676) en la misma tarea.
• Generalización: El estimador mantuvo un alto rendimiento en tareas de codificación, uso de herramientas y seguimiento de instrucciones, superando a menudo a los críticos separados en estos dominios.

Significado y Afirmaciones

El artículo afirma que POISE proporciona una ruta eficiente en cómputo hacia un RLVR estable y escalable para modelos de razonamiento avanzados. Al aprovechar las representaciones internas del propio modelo, el método:

• Elimina la sobrecarga de memoria y cómputo de entrenar un modelo crítico separado.
• Elimina la ineficiencia de muestreo de las líneas base relativas al grupo (GRPO), permitiendo una mayor diversidad de prompts dentro de un presupuesto fijo.
• Evita el desperdicio de cómputo en prompts degenerados que producen una ventaja cero.
• Demuestra que los estados internos, previamente utilizados principalmente como herramientas de diagnóstico, pueden servir como señales de optimización prácticas y de alta fidelidad para el aprendizaje por refuerzo.

Los autores concluyen que POISE permite una optimización de políticas más estable y eficiente sin necesidad de líneas base relativas al grupo o modelos de valor auxiliares, convirtiéndolo en una alternativa viable para entrenar modelos de razonamiento avanzados bajo restricciones de recursos.