Blockwise Advantage Estimation for Multi-Objective RL with Verifiable Rewards

Resumen Técnico: Estimación de Ventaja por Bloques (BAE) para RL Multiobjetivo con Recompensas Verificables

1. El Problema: Interferencia de Objetivos y Asignación de Crédito

En el entrenamiento de Grandes Modelos de Lenguaje (LLMs) mediante Aprendizaje por Refuerzo (RL), métodos como GRPO (Group Relative Policy Optimization) son populares por su eficiencia de memoria, ya que no requieren un modelo crítico (critic model) separado. Sin embargo, GRPO estándar trata cada generación como un bloque monolítico: calcula una única ventaja escalar y la aplica uniformemente a todos los tokens de la respuesta.

Esto presenta dos problemas críticos en tareas estructuradas (donde la respuesta tiene segmentos con objetivos distintos, como una solución matemática seguida de un análisis de confianza):

• Interferencia de objetivos: Los gradientes de un objetivo (ej. precisión de la respuesta) se mezclan con los de otro (ej. calibración de la confianza), degradando el rendimiento de ambos.
• Asignación de crédito deficiente: Es difícil determinar qué tokens específicos fueron responsables de una recompensa particular si la recompensa se aplica a toda la secuencia.
• Sesgo en la línea base (Baseline): En generaciones secuenciales, los bloques posteriores dependen de los prefijos generados anteriormente. Usar una media de grupo global para normalizar recompensas en bloques tardíos es estadísticamente incorrecto porque ignora que el "estado" inicial de ese bloque varía entre las muestras del grupo.

2. Metodología Propuesta: Blockwise Advantage Estimation (BAE)

Los autores proponen una familia de métodos compatibles con GRPO que descomponen la optimización en segmentos de texto.

A. Descomposición por Bloques:
En lugar de una ventaja única, el método asigna una ventaja de proximidad (proxy-advantage) $\hat{A}_k^{(i)}$ a cada bloque $k$. La función de pérdida se optimiza sumando los objetivos de cada bloque, aplicando la ventaja correspondiente solo a los tokens que pertenecen a dicho bloque.

B. Línea Base Condicionada por Resultado (Outcome-Conditioned Baseline - OCB):
Para resolver el problema de la estimación de la línea base en bloques tardíos sin recurrir a costosos rollouts de Monte Carlo (que requerirían generar múltiples continuaciones desde estados intermedios), proponen el OCB.

• Funcionamiento: El grupo de muestras se estratifica en subgrupos basados en un resultado intermedio discreto $o$ (por ejemplo, si la solución previa fue "correcta" o "incorrecta").
• Cálculo: La línea base para un bloque $k$ se calcula como la media de las recompensas de las muestras que comparten el mismo resultado intermedio. Esto aproxima el valor esperado del estado intermedio de manera eficiente y sin costo computacional adicional.

3. Contribuciones Clave

1. Marco BAE: Un marco de trabajo critic-free y compatible con GRPO que permite optimizar múltiples objetivos de forma modular en generaciones estructuradas.
2. Estimador OCB: Un método computacionalmente eficiente para estimar líneas base condicionales utilizando únicamente las estadísticas del grupo actual, reduciendo la varianza sin necesidad de inferencia extra.
3. Mitigación de Reward Hacking: Demuestran que su método permite usar funciones de recompensa más complejas (como la entropía o la pérdida de entropía cruzada - BCE) sin que el modelo caiga en comportamientos degenerados (como dar respuestas vacías para evitar penalizaciones).

4. Resultados Experimentales

El estudio se centró en la tarea de razonamiento matemático conjunto con estimación de incertidumbre (donde el modelo debe resolver un problema y luego reportar su confianza).

• Rendimiento en Calibración: En el modelo Qwen2.5-3B-Instruct, el método OCB superó significativamente a los métodos estándar en métricas de calibración, logrando un ECE (Expected Calibration Error) de 0.030 frente al 0.059 de los métodos basados en diseño de recompensas tradicionales (RLCR).
• Precisión y Robustez: El método es competitivo con el estado del arte (SoTA) en precisión matemática y mantiene la capacidad de mejora mediante Test-Time Scaling (TTS). Es decir, las estimaciones de confianza producidas por el modelo son lo suficientemente fiables como para usarse para seleccionar la mejor respuesta entre múltiples muestras.
• Generalización: Se demostró eficacia en escenarios de "refinamiento de dos intentos", donde el modelo genera una solución, la reflexiona y genera una segunda, mejorando consistentemente el rendimiento.

5. Significancia y Conclusión

Este trabajo es significativo porque ofrece una "receta modular" para el entrenamiento de agentes y modelos de razonamiento de largo horizonte. Al desacoplar los objetivos, permite que los desarrolladores añadan nuevas capacidades (como reflexión, verificación o uso de herramientas) de forma independiente, sin tener que rediseñar cuidadosamente una función de recompensa escalar única que equilibre todos los aspectos, lo cual es uno de los mayores retos actuales en el RL para LLMs.