Resumen Técnico: Optimización de Rollout Adaptativa entre Épocas para el Post-entrenamiento de RL (CERO)

1. Formulación del Problema

El post-entrenamiento de Grandes Modelos de Lenguaje (LLM) mediante Aprendizaje por Refuerzo (RL), particularmente utilizando algoritmos como GRPO, depende típicamente de muestrear un número fijo de rollouts por cada prompt. Esta estrategia de asignación uniforme ignora la heterogeneidad en las señales de entrenamiento proporcionadas por diferentes prompts. Algunos prompts son triviales (generando éxito o fracaso consistentes con baja varianza), mientras que otros son demasiado difíciles o ya han sido dominados, ofreciendo rendimientos marginales decrecientes. En consecuencia, una parte significativa del presupuesto global de rollouts se desperdicia en prompts con un valor de aprendizaje limitado.

El artículo formula esto como un problema de asignación de recursos online entre épocas (cross-epoch) bajo un presupuesto global fijo $B$. A diferencia de métodos previos que realizan una reasignación dentro del lote (within-batch) o filtrado post-generación, este trabajo aborda la asignación de rollouts tanto a través de prompts como de épocas de entrenamiento. El objetivo es maximizar la utilidad acumulada de las muestras recolectadas, donde la utilidad está sujeta a rendimientos decrecientes a nivel de prompt. El problema se define matemáticamente como la maximización de la suma de las funciones de utilidad específicas de cada prompt $U_i(\sum_k N_{i,k})$ sujeta a la restricción de presupuesto global $\sum_{i,k} N_{i,k} \le B$, donde $N_{i,k}$ es el número de rollouts asignados al prompt $i$ en la época $k$.

2. Metodología: CERO

Los autores proponen CERO (Optimización de Rollout Adaptativa entre Épocas), un algoritmo prímal-dual bayesiano diseñado para asignar adaptativamente los rollouts basándose en la estimación de la informatividad del prompt.

2.1 Diseño de Informatividad y Utilidad

• Modelado de la Posterior Beta: Para cada prompt $x_i$, CERO mantiene una distribución posterior Beta sobre su probabilidad de éxito latente $p_i$, actualizada basada en el historial de rollouts $H_i^k$.
• Puntuación de Informatividad ($q_{k,i}$): En lugar de usar una estimación puntual de la probabilidad de éxito, CERO estima la varianza de Bernoulli esperada: $q_{k,i} = E[p_i(1-p_i) | H_i^k]$. Esta puntuación se maximiza cuando $p_i \approx 0.5$ (dificultad intermedia) y desaparece para prompts fáciles ($p_i \approx 1$) o difíciles ($p_i \approx 0$), identificando eficazmente los prompts que ofrecen las señales de aprendizaje más inciertas y, por lo tanto, más informativas.
• Función de Utilidad: Se define una función de utilidad cóncava y saturante $U_i(n) = 1 - \exp(-\eta q_i n)$, donde $n$ es el número acumulado de rollouts. Esto captura los rendimientos decrecientes: los primeros rollouts en un prompt informativo son altamente valiosos, pero el valor marginal disminuye a medida que se recolectan más muestras.

2.2 Optimización mediante Reformulación Fenchel-Dual

El acoplamiento temporal de la función de utilidad (que depende de las asignaciones acumuladas) dificulta la optimización online directa. CERO aborda esto mediante una reformulación Fenchel-dual:

• El problema se transforma en un objetivo dual que involucra variables duales a nivel de prompt $\theta_i$ y un multiplicador de presupuesto global $\mu$.
• Esta reformulación descompone el problema temporalmente acoplado en subproblemas de asignación por época coordinados por las variables duales.

2.3 Algoritmo Primal-Dual Online

CERO opera como una capa de inserción (drop-in) sobre los pipelines de RL existentes (por ejemplo, GRPO):

1. Decisión de Asignación: En cada época, para cada prompt, el algoritmo resuelve un subproblema lineal. Un prompt recibe el máximo de rollouts permitidos ($N_{max}$) si su valor marginal estimado (variable dual $\theta_i$) supera el precio del presupuesto global actual ($\mu$); de lo contrario, recibe cero.
2. Actualización de Retroalimentación: Tras generar las respuestas y observar las recompensas, la posterior Beta para el prompt se actualiza, y la puntuación de informatividad $q_{k,i}$ se refresca.
3. Actualizaciones de Variables Duales:
   • A nivel de prompt ($\theta$): Se actualiza mediante descenso de gradiente online proyectado para equilibrar los rollouts asignados frente a la utilidad implicada por la estimación de informatividad actual.
   • A nivel de presupuesto ($\mu$): Se actualiza basándose en la discrepancia entre el consumo real de rollouts y la tasa de presupuesto objetivo. Si el consumo excede la tasa, $\mu$ aumenta, haciendo que la asignación futura sea más selectiva; si es inferior, $\mu$ disminuye.

3. Principales Contribuciones

1. Formulación del Problema: El artículo es el primero en formular la asignación de rollouts en el post-entrenamiento de RL de LLM como un problema de optimización online de presupuesto fijo entre épocas con rendimientos decrecientes específicos de cada prompt, acoplando explícitamente las decisiones a través del horizonte de entrenamiento mediante un presupuesto global compartido.
2. Diseño del Algoritmo: Introduce CERO, un algoritmo primal-dual bayesiano que utiliza estimaciones posteriores del valor de los rollouts a nivel de prompt para guiar el presupuesto. Utiliza la varianza de Bernoulli posterior esperada como un proxy de la capacidad de aprendizaje y emplea una reformulación Fenchel-dual para manejar la no-separabilidad temporal.
3. Validación Teórica y Empírica:
   • Teoría: Bajo utilidades de prompt fijas, los autores demuestran un límite de arrepentimiento (regret bound) de $O(\sqrt{K})$ contra un benchmark de asignación offline, donde $K$ es el número de épocas.
   • Experimentos: Experimentos extensos en tareas de razonamiento matemático (AIME, AMC) en cuatro modelos de lenguaje de pesos abiertos (que van desde 1.5B hasta 7B de parámetros) demuestran que CERO supera consistentemente al GRPO vanilla.

4. Resultados Experimentales

Los autores evaluaron CERO en el conjunto de datos DAPO-Math-17K utilizando DeepSeek-R1-Distill-1.5B, Qwen3-4B (Base e Instruct) y Qwen2.5-Math-7B.

• Rendimiento: CERO logró mejoras consistentes sobre el GRPO vanilla en todos los modelos y benchmarks. Las ganancias promedio oscilaron entre +1.53 y +6.04 puntos dependiendo del modelo.
• Análisis de Eficiencia: La "Relación de Prompts Efectivos" (la fracción de prompts en un lote que contiene tanto respuestas exitosas como fallidas, las cuales generan gradientes no nulos) fue significativamente mayor para CERO que para GRPO.
  • En el baseline GRPO, la relación de prompts efectivos cayó bruscamente a medida que avanzaba el entrenamiento, indicando que el modelo había dominado muchos prompts, lo que llevó a grupos de rollouts degenerados.
  • CERO mantuvo una relación de prompts efectivos más alta al desplazar adaptativamente el presupuesto de los prompts dominados hacia aquellos con mayor incertidumbre.
• Sobrecarga (Overhead): El mecanismo de actualización de la asignación introduce una sobrecarga computacional insignificante (menos de un segundo por actualización).

5. Significado y Reivindicaciones

El artículo sostiene que la asignación adaptativa de rollouts es una dirección fundamentada y efectiva para mejorar la eficiencia de muestreo en el post-entrenamiento de RL. Al alejarse de la asignación uniforme, CERO asegura que el presupuesto de generación fijo se gaste en prompts que proporcionan las señales de aprendizaje más informativas.

Los autores posicionan a CERO como una capa de recolección de datos adaptativa de "inserción directa" (drop-in) que modifica qué prompts reciben rollouts y cuántos, sin alterar la regla de gradiente de política subyacente. El trabajo destaca que la ineficiencia de los actuales procesos de post-entrenamiento de RL proviene no solo de la regla de actualización de la política, sino de la estrategia de recolección de datos subóptima.

Limitaciones señaladas por los autores:

• La garantía teórica de arrepentimiento asume utilidades de prompt fijas, mientras que la implementación utiliza estimaciones posteriores dinámicas por época. Por lo tanto, el análisis no captura plenamente la no-estacionariedad inducida por las actualizaciones de la política y los refrescos de la posterior.
• La validación empírica está actualmente limitada a tareas de razonamiento matemático; se sugiere como trabajo futuro una validación más amplia en otros dominios y modelos de recompensa.