OR Else: A Differentiable Trust Region for Policy Optimization

Resumen Técnico: OR ELSE: UNA REGIÓN DE CONFIANZA DIFERENCIABLE PARA LA OPTIMIZACIÓN DE POLÍTICAS

Planteamiento del Problema
El Aprendizaje por Refuerzo a partir de la Retroalimentación Humana (RLHF) para Modelos de Lenguaje Extensos (LLMs) se basa comúnmente en la Optimización de Política Próxima (PPO) y su variante, la Optimización de Política Relativa de Grupo (GRPO). Ambos métodos utilizan un objetivo de sustitución recortado (clipped surrogate objective) para restringir las actualizaciones de la política. Si bien el mecanismo de recorte crea una región de gradiente cero de un solo lado para evitar un desplazamiento de política excesivo en la dirección "favorable", la transición hacia esta región es geométricamente abrupta: la derivada de la función objetivo con respecto a la razón de probabilidad cambia discontinuamente en el límite de recorte. Los autores investigan si el Output Reset (OR), una regla de saturación suave de un solo lado diseñada originalmente para el entrenamiento de clasificadores, ofrece una alternativa superior para el post-entrenamiento de LLMs. Específicamente, el artículo plantea si reemplazar el sustituto recortado por una pérdida basada en OR que sea diferenciable puede mejorar la estabilidad de la optimización y las puntuaciones finales de recompensa sin introducir nuevas patologías.

Metodología
El artículo introduce PPO-OR y GRPO-OR, dos variantes que reemplazan el término de política recortado con una pérdida de margen cuadrático derivada de OR en el espacio de los log-ratios de los tokens.

• Mecanismo Central: En lugar de utilizar la razón de probabilidad $r_t(\theta)$, los métodos operan sobre el log-ratio de probabilidad por token $\rho_t = \log \pi_\theta(a_t|s_t) - \log \pi_{\theta_{old}}(a_t|s_t)$.
• La Pérdida OR: Para un token con una ventaja estimada no nula $\hat{A}_t$, se establece un margen objetivo $\sigma_t \alpha_{OR}$ (donde $\sigma_t = \text{sign}(\hat{A}_t)$). Este es un margen cuadrático de un solo lado: $\ell_{OR,t} = [\max(0, \alpha_{OR} - \sigma_t \rho_t)]^2$.
• Comportamiento de Output Reset: Una vez que $\rho_t$ cruza el margen en la dirección favorecida por el signo de la ventaja, el objetivo se "reinicia" al $\rho_t$ actual. Esto resulta en un residuo cero y una contribución de gradiente directo nula para ese token. La derivada se acerca a cero continuamente en el límite, lo que hace que la pérdida sea $C^1$ continua (a diferencia del recorte de PPO, que es $C^0$ continua pero no diferenciable).
• Estimación de la Ventaja:
  • PPO-OR: Conserva la Estimación de Ventaja Generalizada (GAE) con una cabeza de valor aprendida.
  • GRPO-OR: Conserva las ventajas relativas al grupo (calculadas a partir de $G=2$ respuestas por prompt) sin un crítico.
• Configuración Experimental: El estudio utiliza Llama-3.2-1B-Instruct en el conjunto de datos Anthropic hh-rlhf. Todos los métodos comparten un único modelo de recompensa congelado y son evaluados sobre tres semillas. Los experimentos están diseñados como comparaciones emparejadas dentro de la misma familia: PPO-clip vs. PPO-OR (bajo GAE) y GRPO vs. GRPO-OR (bajo ventajas relativas al grupo). Las comparaciones entre familias se señalan como no controladas debido a las diferencias en la estimación de la ventaja, el uso del crítico y las longitudes de respuesta.

Contribuciones Clave

1. Propuesta Algorítmica: La introducción de PPO-OR y GRPO-OR, que sustituyen el sustituto recortado por una pérdida de margen cuadrático de un solo lado y continuamente diferenciable en el espacio del log-ratio.
2. Caracterización Teórica: Prueba formal de que el objetivo OR es $C^1$ continuo en el límite de saturación y que la dirección de actualización a nivel de muestra está estrictamente determinada por el signo de la ventaja, desvaneciéndose continuamente al cruzar el margen.
3. Evaluación Empírica: Una comparación rigurosa entre el recorte y OR dentro de dos familias distintas de estimación de la ventaja, destacando que los beneficios de OR no son uniformes sino que dependen del pipeline de optimización circundante.
4. Perspectiva Diagnóstica: La identificación de que reemplazar el recorte por OR no resuelve inherentemente el problema del gran desplazamiento acumulativo de la política observado en configuraciones relativas al grupo, distinguiendo la saturación del gradiente local del control global de la política.

Resultados
El artículo reporta las puntuaciones del modelo de recompensa durante el entrenamiento y los diagnósticos de optimización a través de tres semillas:

• Familia GAE (PPO-clip vs. Pppo-OR):

  • Puntuación de Recompensa: PPO-OR logró una puntuación media final más alta (+0.500) en comparación con PPO-clip (+0.195), una diferencia de +0.305.
  • Estabilidad: PPO-OR exhibió una desviación estándar entre semillas mayor (0.158) en comparación con PPO-clip (0.110).
  • Desplazamiento (Drift): Ambos métodos mantuvieron un desplazamiento de log-ratio (rollout-to-current) similar y bajo (aprox. 0.3–0.5).

• Familia Relativa al Grupo (GRPO vs. GRPO-OR):

  • Puntuación de Recompensa: GRPO-OR no mejoró la puntuación media; fue ligeramente inferior (+0.457) a GRPO (+0.488).
  • Estabilidad: GRPO-OR mostró una desviación estándar entre semillas significativamente menor (0.027) en comparación con GRPO (0.080).
  • Diagnósticos: GRPO-OR logró un residuo OR terminal cercano a cero y una fracción de sobreimpulso (overshoot) decreciente, mientras que la traza de GRPO emparejada se mantuvo grande y altamente variable.
  • Desplazamiento (Drft): Crucialmente, GRPO-OR no redujo el gran desplazamiento de la política observado en GRPO. Ambos métodos relativos al grupo exhibieron un desplazamiento sustancialmente mayor (aprox. 5–13) en comparación con los métodos GAE.

Significado y Reivindicaciones
El artículo concluye de manera modesta, enfatizando que los efectos observados de OR dependen del contexto:

• No es una Región de Confianza Universal: Los autores afirman explícitamente que la saturación suave del objetivo OR no constituye una región de confianza formal. Suprime la contribución directa del gradiente de los tokens individuales después de que cruzan un margen, pero no restringe el desplazamiento acumulativo de la política causado por parámetros compartidos, pérdidas auxiliares o cambios en los signos de la ventaja a través de las iteraciones.
• Dependencia del Pipeline: La eficacia de OR está ligada al método de estimación de la ventaja. En el entorno GAE, OR se correlacionó con mayores puntuaciones de recompensa pero aumentó la varianza. En el entorno relativo al grupo, OR redujo la varianza y estabilizó la traza de la pérdida, pero no mejoró la puntuación media de recompensa ni redujo el desplazamiento de la política.
• Limitaciones: El estudio reconoce que los resultados son específicos para un modelo ligero de 1B, un único modelo de recompensa congelado y un tamaño de grupo restringido ($G=2$). Los hallazgos caracterizan este entorno experimental específico en lugar de establecer una superioridad general sobre PPO o GRPO. El artículo postula que el gran desplazamiento de la política en los métodos relativos al grupo sigue siendo un problema no resuelto que OR por sí solo no puede solucionar, sugiriendo que el trabajo futuro debe abordar el desplazamiento acumulativo mediante una regularización de referencia más fuerte o penalizaciones de desplazamiento explícitas.

En resumen, el artículo demuestra que el Output Reset ofrece una alternativa diferenciable al recorte que altera el comportamiento de la optimización y puede mejorar las puntuaciones de recompensa en contextos específicos, pero no resuelve inherentemente los desafíos de la estabilidad de la política o el desplazamiento en el RLHF relativo al grupo.