Resumen Técnico: Replanteando la Regularización de la Divergencia en el RL para LLMs

1. Planteamiento del Problema

El aprendizaje por refuerzo (RL) es una piedra angular del post-entrenamiento de los modelos de lenguaje extensos (LLMs) para alinearlos con las preferencias humanas y mejorar sus capacidades de razonamiento. Sin embargo, el entrenamiento de RL para LLMs modernos es inherentemente fuera de la política (off-policy) debido a las discrepancias numéricas entre los motores de entrenamiento e inferencia y a la práctica de dividir los rollouts en múltiples mini-lotes (mini-batches). Esto requiere la Optimización de Política de Región de Confianza (TRPO, por sus siglas en inglés) para asegurar actualizaciones estables mediante la restricción de la divergencia entre la política actual ($\pi$) y la política de comportamiento ($\mu$).

Los métodos convencionales actuales enfrentan dos limitaciones principales en este contexto:

1. Restricciones Basadas en Razones (PPO, GRPO, SPO): Estos métodos utilizan la razón de importancia $r_t = \pi(y_t|s_t)/\mu(y_t|s_t)$ para definir la región de confianza. En los LLMs con vocabularios grandes y de cola larga, la razón es un mal indicador del cambio distributivo. Un pequeño aumento de probabilidad en un token de baja probabilidad puede generar una razón enorme con un cambio de masa insignificante, mientras que un cambio moderado en la razón de un token de alta probabilidad puede inducir cambios de política significativos. Esto conduce a una sobre-restricción de los tokens raros y a una sub-restricción de los comunes.
2. Enmascaramiento Duro (DPPO): Trabajos recientes como la Optimización de Política de Divergencia Proximal (DPPO) abordan el problema de la razón utilizando una máscara de divergencia (específicamente Binary-TV, que mide el desplazamiento de probabilidad absoluta $|\pi - \mu|$). Sin embargo, DPPO impone la región de confianza mediante una máscara binaria dura. Una vez que un token cruza el límite en una dirección perjudicial, su gradiente se descarta (se establece en cero). Esto evita una mayor divergencia, pero no proporciona una señal correctiva para atraer la política de vuelta hacia la región de confianza, lo que genera una optimización frágil cerca del umbral.

2. Metodología: Optimización de Política Regularizada por Divergencia (DRPO)

Los autores proponen DRPO, un método que reemplaza la máscara dura de DPPO con un regularizador cuadrático suave y ponderado por la ventaja sobre el desplazamiento de la política, preservando al mismo tiempo la misma geometría de la región de confianza de Binary-TV.

Formulación Central

DRPO deriva su objetivo reescribiendo la restricción Binary-TV $|\pi(y_t|s_t) - \mu(y_t|s_t)| \leq \delta$ como un límite de razón adaptativo al token. En lugar de aplicar un recorte (clipping), aplica un regularizador suave al objetivo subsidiario:

$$\mathcal{L}_{DRPO}(x, \pi) = \mathbb{E}_{y \sim \mu} \left[ \sum_{t=1}^{|y|} \left( r_t \hat{A}_t - \frac{|\hat{A}_t|}{2\delta} \mu(y_t|s_t) (r_t - 1)^2 \right) \right]$$

Donde:

• $r_t$ es la razón de importancia.
• $\hat{A}_t$ es la estimación de la ventaja.
• $\delta$ es el umbral de la región de confianza.
• El término $\mu(y_t|s_t)$ escala el regularizador, convirtiendo la penalización de una forma $\chi^2$ basada en la razón (como en SPO) a una forma de desplazamiento de probabilidad absoluta ($\ell_2^2$).

Comportamiento del Gradiente

El gradiente del objetivo de DRPO introduce un peso continuo $w_t$ para el gradiente de política de cada token:
$$w_t = 1 - \text{sign}(\hat{A}_t(r_t - 1)) \frac{|\pi(y_t|s_t) - \mu(y_t|s_t)|}{\delta}$$

Este peso exhibe tres comportamientos clave:

1. Dentro de la Región de Confianza: El peso es positivo pero se atenúa suavemente a medida que el desplazamiento se acerca a $\delta$, permitiendo que las actualizaciones continúen pero con una magnitud reducida.
2. Fuera de la Región de Confianza (Divergiendo): Cuando la actualización se aleja más de la política de comportamiento ($|\pi - \mu| > \delta$), el peso se vuelve negativo. Esto revierte el gradiente, proporcionando una señal correctiva que atrae la política de vuelta hacia el límite, a diferencia del gradiente cero de las máscaras duras.
3. Actualizaciones de Convergencia: Cuando la actualización regresa hacia la política de comportamiento, el peso se amplifica ($>1$), fomentando un retorno suave a la región de confianza.

Distinciones Teóricas

• vs. SPO: SPO utiliza un regularizador de tipo $\chi^2$ (dividiendo por $\mu$), lo que causa que los pesos del gradiente crezcan sin control en la cola larga de los tokens de baja probabilidad. DRPO utiliza un regularizador de tipo $\ell_2^2$ (multiplicando por $\mu$), asegurando que los pesos del gradiente permanezcan acotados independientemente de la probabilidad del token.
• vs. DPPO: DRPO preserva la geometría de Binary-TV de DPPO, pero reemplaza la máscara discontinua con un regularizador suave y diferenciable.

3. Contribuciones Clave

1. Algoritmo DRPO: Un nuevo método de optimización de política que reemplaza las máscaras de divergencia duras con un regularizador cuadrático suave basado en el desplazamiento de probabilidad absoluta (Binary-TV).
2. Pesos de Gradiente Acotados: Demostración teórica y empírica de que los pesos del gradiente de DRPO permanecen acotados en vocabularios de cola larga, evitando la inestabilidad causada por los métodos basados en razones (PPO, SPO) donde los pesos pueden explotar para tokens de baja probabilidad.
3. Mecanismo Correctivo: Introducción de una señal correctiva continua para las actualizaciones que exceden la región de confianza, evitando la "muerte del gradiente" observada en los enfoques de máscara dura como DPPO.
4. Principios de Diseño: Identificación de tres criterios prácticos para regularizadores efectivos en el RL de LLMs:
   • Inducir un límite estable alineado con el cambio distributivo (no solo con las razones).
   • Mantener los pesos del gradiente por token acotados en vocabularios de cola larga.
   • Proporcionar señales correctivas suaves cuando la política diverge.

4. Resultados Experimentales

Los autores evaluaron DRPO en diversos entornos, incluyendo:

• Modelos: Qwen3-4B, Qwen3-30B-A3B, Qwen3.5-35B-A3B, y DeepSeek-R1-Distill-Qwen-1.5B.
• Arquitecturas: Densa y Mezcla de Expertos (MoE).
• Precisión: BF16, FP8 (solo para rollout), y FP8 (extremo a extremo).
• Benchmarks: AIME 2024 y AIME 2025 (tareas de razonamiento matemático).

Hallazgos Clave:

• Estabilidad: DRPO permite consistentemente un entrenamiento estable en todos los entornos. En contraste, los métodos basados en razones (GRPO, SPO) sufrieron frecuentemente de inestabilidad o colapso, particularmente en configuraciones de baja precisión (FP8).
• Rendimiento: DRPO igualó o superó la mejor precisión de evaluación de todas las líneas base. Superó a los métodos de máscara dura (DPPO) en velocidad de convergencia y precisión final.
• Ablaciones:
  • Eliminar la ponderación de la ventaja absoluta $|\hat{A}_t|$ del regularizador provocó caídas de rendimiento e inestabilidad, confirmando que el límite de la región de confianza debe ser independiente de la escala de la recompensa.
  • Otras medidas de divergencia (penalizaciones KL, TV) tuvieron un rendimiento inferior al de DRPO, lo que sugiere que la geometría específica del regularizador Binary-TV es crucial.
  • Aplicar el regularizador solo fuera de la región de confianza (Mask-DRPO) produjo un rendimiento comparable al de DRPO completo, indicando que la principal ganancia proviene de la regularización correctiva fuera del límite.

5. Significado y Reivindicaciones

El artículo afirma que la geometría del gradiente inducida por un regularizador es más crítica que la medida de divergencia nominal en la función objetivo. Mientras que trabajos previos suelen centrarse en divergencias estándar (KL, JS), este trabajo demuestra que:

• Las restricciones basadas en razones son fundamentalmente defectuosas para los LLMs debido a sus vocabularios de cola larga.
• Las máscaras duras, aunque mejores que las razones, introducen fragilidad al descartar las señales correctivas.
• DRPO ofrece una solución unificada que combina la corrección geométrica de las restricciones basadas en la divergencia (Binary-TV) con la estabilidad de optimización de los regularizadores suaves.

Los autores posicionan a DRPO no solo como un algoritmo específico, sino como una lente empírica para diseñar objetivos de optimización de política estables, enfatizando que los regularizadores efectivos deben proporcionar señales de gradiente acotadas, continuas y correctivas para manejar los desafíos únicos del RL para LLMs.