Future-KL Regularized GRPO: Process-Level Credit Assignment from $f$-Divergence Regularization

Resumen Técnico: GRPO Regularizado con KL del Futuro

Enunciado del Problema
La Optimización de Política Relativa por Grupos (GRPO) ha surgido como un paradigma dominante para el aprendizaje por refuerzo (RL) sin crítico en el entrenamiento posterior de Modelos de Lenguaje Grande (LLM), particularmente para el razonamiento matemático. GRPO elimina la necesidad de una función de valor aprendida normalizando las recompensas de resultado dentro de un grupo de respuestas muestreadas. Sin embargo, las implementaciones estándar de la regularización KL en GRPO suelen tratar la penalización como una penalización local de tokens en la pérdida (por ejemplo, utilizando el estimador $k_3$).

Los autores identifican un defecto fundamental en este enfoque: ignora la asignación de crédito a nivel de proceso inducida por la generación autoregresiva. En una trayectoria autoregresiva, un token temprano no solo incurre en su propio costo de regularización local, sino que también altera la distribución de los prefijos futuros, afectando así los costos KL de todos los tokens subsiguientes. Los términos de pérdida locales estándar diferencian solo la penalización del token muestreado, omitiendo el "retorno al futuro de regularización causal".

Además, el artículo destaca una incompatibilidad estructural entre la normalización por grupos de GRPO y las formulaciones estándar de KL en la recompensa. A diferencia de los objetivos de recompensa lineales (por ejemplo, RL de Máxima Entropía) donde la recompensa y el regularizador pueden fusionarse, la normalización por grupos de GRPO induce una utilidad no lineal a nivel de prompt (específicamente, una transformación arcsin para recompensas binarias). Esta no linealidad significa que añadir penalizaciones KL a las recompensas antes de la normalización distorsiona el objetivo implícito, colapsando la información KL a nivel de token en un único escalar a nivel de trayectoria y fallando en preservar la señal de gradiente correcta.

Metodología
El artículo deriva el gradiente de política en política para objetivos al estilo de GRPO bajo regularización general de divergencia $f$ token a token. La idea teórica clave es que el gradiente se descompone en dos términos distintos con propiedades de escalado diferentes:

1. El Término de Recompensa: Heredado de la utilidad no lineal $\ell(p)$, este término lleva una ganancia dependiente del prompt $\ell'(p)$. Para recompensas binarias, esto recupera la ventaja estandarizada de GRPO.
2. El Término del Regularizador: Este término es independiente de $\ell'$ y se descompone en una penalización local de token más un retorno al futuro de regularización causal.

Basándose en esta derivación, los autores proponen la Optimización de Política Regularizada con KL del Futuro (FRPO). El algoritmo modifica la actualización estándar de GRPO de manera desacoplada:

1. Calcular la ventaja de resultado normalizada por grupos ($\hat{A}^{GRPO}$) estrictamente a partir de las recompensas originales del verificador, preservando la utilidad no lineal implícita.
2. Construir un término de corrección de KL futuro a nivel de token calculando la suma acumulada inversa de las razones de logaritmo por token ($\delta_{i,j} = \log \frac{\pi_\theta}{\pi_{ref}}$) desde el token actual hasta el final de la trayectoria.
3. Añadir este término de KL futuro a la ventaja después de la construcción de la ventaja, en lugar de dar forma a la recompensa antes de la normalización.

La regla de actualización para un token $t$ en la respuesta $i$ es:
$$\hat{A}^{FRPO}_{i,t} = \text{sg}\left( \hat{A}^{GRPO}_i - \beta \sum_{j=t}^{T_i} \delta_{i,j} \right)$$
donde $\text{sg}$ denota gradiente detenido. Este enfoque no requiere un crítico, ni pasadas adicionales del modelo, y mantiene la misma sobrecarga computacional que GRPO estándar.

Contribuciones Clave

1. Identificación de Utilidad No Lineal: El artículo demuestra que la normalización por grupos de GRPO induce una utilidad no lineal ($\ell(p) = 2 \arcsin \sqrt{p}$ para recompensas binarias), haciendo que la recompensa y el KL no sean fusionables. Añadir KL a las recompensas antes de la normalización distorsiona estructuralmente el objetivo.
2. Derivación del Gradiente en Política: Los autores derivan el gradiente en política correcto para la regularización de divergencia $f$ token a token, mostrando que esto requiere un término de retorno al futuro de regularización que está ausente en las implementaciones locales en la pérdida.
3. Algoritmo FRPO: Proponen FRPO, una actualización simple y desacoplada que preserva la ventaja de resultado de GRPO mientras inyecta la asignación de crédito de KL futuro necesaria. Este método separa la clasificación a nivel de trayectoria de la regularización KL a nivel de proceso.
4. Interpretación de Optimalidad: Se muestra que la estructura de KL futuro se alinea con las condiciones de optimalidad del RL regularizado por entropía, donde la política óptima satisface una forma de recompensa de proceso que involucra costos KL futuros.

Resultados Experimentales
Los autores evalúan FRPO en tareas de razonamiento matemático utilizando modelos base Qwen3 (1.7B, 4B y 30B-A3B) entrenados en el conjunto de datos DAPO-Math-17k.

• Rendimiento: En el modelo 30B-A3B, FRPO mejora el pass@16 en casi un 5% de ganancia absoluta sobre la línea base más fuerte (GRPO estándar con KL en la pérdida) en la prueba MATH500. También supera a las líneas base en AIME24, AIME25 y otras pruebas de nivel olímpico.
• Estabilidad y Deriva: FRPO mantiene una entropía final más alta y una deriva de política menor (medida por PPO-KL) en comparación con las líneas base de KL en la pérdida. Esto sugiere que el método elicita patrones de razonamiento de Cadena de Pensamiento (CoT) largos sin desviarse excesivamente de la política de referencia.
• Ablación: Los experimentos confirman que el KL inverso (búsqueda de modo) funciona mejor que el KL directo (cobertura de modo) para esta tarea. El artículo también señala que usar el estimador $k_1$ (razón de logaritmo) es preferible a $k_3$ para el término de KL futuro, ya que la no negatividad de $k_3$ crea una señal de penalización unilateral que dificulta la asignación de crédito.
• Modos de Fallo: El artículo observa que los enfoques de "KL en recompensa" (añadir KL a las recompensas antes de la normalización) conducen al colapso del modelo o a la degradación del rendimiento, validando la afirmación teórica de que dicha fusión es estructuralmente incorrecta bajo la utilidad no lineal de GRPO.

Significado
El artículo argumenta que la importancia de este trabajo radica en corregir un error de implementación sutil pero crítico en un marco de RL ampliamente utilizado. Al cambiar la perspectiva de "penalización local de pérdida" a "asignación de crédito a nivel de proceso", los autores muestran que GRPO puede mejorarse sin introducir funciones de valor complejas ni costos computacionales adicionales. El método propuesto FRPO proporciona una modificación ligera y fundamentada teóricamente que mejora las capacidades de razonamiento mientras preserva las propiedades de estabilidad y generalización del modelo base. El trabajo subraya que, en entornos autoregresivos, la regularización es inherentemente una señal a nivel de proceso, e ignorar el costo futuro al futuro conduce a una asignación de crédito subóptima.