Resumen Técnico: Máscaras de Divergencia Predictiva para RL de LLM

1. Planteamiento del Problema

El Aprendizaje por Refuerzo (RL) para Modelos de Lenguaje Grandes (LLMs) es inherentemente fuera de la política (off-policy) debido a los desajustes entre los motores de inferencia utilizados para los rollouts y las pilas de entrenamiento, así como al estancamiento de la política (policy staleness). Para estabilizar estas actualizaciones fuera de la política, los métodos estándar como PPO emplean máscaras de región de confianza. Estas máscaras se basan en dos criterios para decidir si se enmascara (se pone a cero) el gradiente de un token:

1. Criterio de Proximidad: Determina si la política se ha alejado demasiado de la política de comportamiento.
2. Criterio de Dirección: Determina si la actualización empuja la política más lejos de la política de comportamiento (hacia afuera) o la atrae de vuelta (correctiva).

Un trabajo reciente, específicamente DPPO, mejoró el criterio de proximidad reemplazando la relación de importancia de una sola muestra por una divergencia distributiva (por ejemplo, la divergencia KL) estimada sobre el top-K del vocabulario. Sin embargo, DPPO conserva el criterio de dirección original de PPO, que depende únicamente de la relación de importancia del token muestreado ($r_t - 1$).

Los autores identifican una inconsistencia fundamental en este enfoque híbrido: mientras que el criterio de proximidad ahora se define mediante una divergencia distributiva global, el criterio de dirección sigue siendo un sustituto de una sola muestra. Los autores observan que el signo del cambio en la relación muestreada puede discrepar con el signo del cambio real en la divergencia. En consecuencia, un token podría mantenerse (porque la relación sugiere una actualización correctiva) incluso cuando la actualización en realidad aumenta la divergencia global, o viceversa.

2. Metodología

El artículo propone la Máscara de Divergencia Predictiva, un método que alinea el criterio de dirección con la naturaleza distributiva del criterio de proximidad. En lugar de preguntar "¿Se alejará la relación muestreada de una?", el método pregunta "¿Aumentará el siguiente paso de gradiente la divergencia?".

Derivación Teórica

Para las políticas softmax utilizadas en LLMs autorregresivos, los autores derivan una expresión de forma cerrada para la derivada direccional de la divergencia de Kullback-Leibler (KL), denotada como $\dot{D}$, a lo largo de un paso de gradiente de la política.
$$\dot{D} = \underbrace{(\pi_k - \mu_k)}_{\text{Término Local}} + \underbrace{\sum_{i} \pi_i (\mu_i - \pi_i)}_{\text{Término Global}}$$
Donde:

• $\pi_k$ y $\mu_k$ son las probabilidades del token muestreado $k$ bajo las políticas de entrenamiento y de comportamiento, respectivamente.
• El Término Local corresponde exactamente a la cantidad utilizada en el estándar criterio de dirección basado en la relación ($r - 1$).
• El Término Global captura el acoplamiento de normalización a través de todo el vocabulario. Este término es invisible para una sola relación muestreada, pero es crucial para determinar el cambio real en la divergencia.

La máscara predictiva mantiene una actualización solo si el signo de $\hat{A}_t \cdot \dot{D}_t$ es no positivo (indicando un movimiento correctivo o neutral), siempre que el token ya esté fuera de la región de confianza ($D_t > \delta$).

Estimación Práctica

Los motores de rollout de producción suelen exponer solo las probabilidades de los tokens top-$K$, no la distribución completa del vocabulario. Para calcular $\dot{D}$ bajo esta vista truncada, los autores proponen dos estimadores ligeros para la cola no vista de la distribución:

1. Estimador de Cola Agregada: Colapsa toda la masa de probabilidad residual en un único token agregado.
2. Estimador de Cola Uniforme: Distribuye la masa residual uniformemente entre los tokens del vocabulario no vistos.

Ambos estimadores requieren solo reducciones simples sobre las probabilidades retenidas del top-$K$ y las masas de la cola, lo que añade una sobrecarga computacional insignificante.

3. Contribuciones Clave

• Identificación de la Inconsistencia: El artículo destaca que mezclar un criterio de proximidad distributivo (divergencia) con un criterio de dirección de una sola muestra (relación) crea un desalineamiento donde la máscara de dirección puede fallar al rastrear el cambio real en la divergencia.
• Máscara de Divergencia Predictiva: Una nueva regla de enmascaramiento que utiliza el cambio de primer orden de la divergencia misma para determinar la dirección de la actualización.
• Derivación de Forma Cerrada: La derivación de la derivada direccional para políticas softmax, separando explícitamente los componentes local (muestreado) y global (normalización).
• Estimadores de Vocabulario Truncado: Dos estimadores de baja sobrecarga (cola agregada y uniforme) que permiten que el método funcione con las restricciones de vista top-$K$ de los entornos de producción estándar.

4. Resultados Experimentales

Los autores evaluaron el método en el conjunto de datos DAPO-Math-17k utilizando modelos Qwen3 (4B, 8B y 30B) con diversas configuraciones de precisión (FP8).

• Estabilidad de Entrenamiento: En todas las escalas de modelos y configuraciones de precisión (incluyendo pruebas de estrés con FP8), la Máscara de Divergencia Predictiva mejoró la estabilidad del entrenamiento en comparación con la línea base DPPO-TopK-KL. En contraste, el GRPO estándar con trucos de clip-higher mostró inestabilidad y colapso en configuraciones de modelos más grandes.
• Alineación del Criterio de Dirección: El análisis a nivel de token reveló que el criterio de dirección basado en la divergencia está mejor alineado con el cambio realizado en la divergencia que el criterio basado en la relación.
  • Entre los tokens donde los dos criterios discrepaban, el criterio basado en la divergencia tuvo una menor tasa de "mantención insegura" (unsafe-keep rate) (34.2% vs 36.9%), lo que significa que era menos propenso a mantener actualizaciones que realmente aumentaban la divergencia.
  • Por el contrario, mantuvo una mayor fracción de actualizaciones que lograron contraer la divergencia (65.1% vs 62.9%).
• Robustez: El método demostró una estabilidad mejorada respecto al hiperparámetro del umbral de la región de confianza ($\delta$), funcionando mejor que la línea base incluso bajo restricciones más estrictas ($\delta = 0.05$).
• Equivalencia de Estimadores: Los estimadores de cola agregada y de cola uniforme produjeron resultados casi idénticos, lo que sugiere que el método no es sensible al supuesto de modelado de la cola.

5. Significado y Reivindicaciones

El artículo afirma que la Máscara de Divergencia Predictiva resuelve una inconsistencia estructural en los métodos actuales de RL basados en divergencia para LLMs. Al asegurar que tanto el criterio de proximidad como el de dirección se deriven de la misma divergencia distributiva, el método proporciona un control de región de confianza más fiel.

Los autores enfatizan que, si bien la mejora en las tasas de "mantención insegura" es modesta (debido al pequeño conjunto de tokens donde los criterios discrepan), la consistencia es teóricamente sólida y empíricamente beneficiosa a través de las escalas. Concluyen que el criterio de dirección debe recibir el mismo tratamiento distributivo que el criterio de proximidad para aprovechar plenamente los beneficios de las regiones de confianza basadas en la divergencia. El método se presenta como un reemplazo directo para las implementaciones existentes de DPPO que no requiere nuevos hiperparámetros más allá del umbral de divergencia existente.