Resumen Técnico: Optimización de Políticas Libres de Valor mediante la Partición de Recompensas

Declaración del Problema

Alinear los Grandes Modelos de Lenguaje (LLMs) con las preferencias humanas depende típicamente del Aprendizaje por Refuerzo con Retroalimentación Humana (RLHF) o de la Optimización de Preferencias Directas (DPO). Sin embargo, estos paradigmas enfrentan desafíos significativos de escalabilidad y estabilidad:

• RLHF requiere la anotación humana costosa de preferencias por pares, el entrenamiento de un Modelo de Recompensa (RM) separado y bucles de aprendizaje por refuerzo inestables (por ejemplo, PPO) que son ineficientes en muestras y sensibles a los hiperparámetros.
• Los métodos de tipo DPO (Optimización de Preferencias Directas) eliminan el RM explícito y el bucle de RL, pero requieren datos de preferencia por pares curados. A medida que los modelos mejoran, distinguir entre respuestas de alta calidad se vuelve cada vez más difícil, lo que aumenta los costos de anotación.
• El Aprendizaje de Trayectoria Única ofrece una alternativa práctica utilizando tuplas $(prompt, respuesta, recompensa\_escalar)$ (que son más fáciles de anotar, por ejemplo, con pulgares arriba/abajo). Sin embargo, los métodos existentes para este escenario, como la Optimización de Recompensa Directa (DRO), dependen del aprendizaje conjunto de una política y una función de valor. Esta optimización conjunta introduce varianza, inestabilidad y sensibilidad a los datos fuera de la política (off-policy), ya que los errores en la estimación del valor pueden sesgar el entrenamiento de la política.

El problema central abordado es la necesidad de un método estable, escalable y simple para optimizar políticas directamente a partir de datos de recompensa escalar sin la complejidad e inestabilidad de aprender una función de valor auxiliar.

Metodología: Optimización de Partición de Recompensa (RPO)

Los autores proponen la Optimización de Partición de Recompensa (RPO), un objetivo libre de valor que aprovecha la forma analítica de la política óptima bajo regularización KL para eliminar la necesidad de una función de valor aprendida.

Fundamento Teórico

Partiendo de la formulación de la política óptima en la maximización de recompensa regularizada por KL:
$$\pi^*(y|x) = \pi_{ref}(y|x) \exp\left(\frac{1}{\tau} (r(x, y) - V^*(x))\right)$$
donde $V^*(x)$ es la función de valor suave que actúa como una función de partición para normalizar las recompensas. Los métodos tradicionales como DRO aproximan $V^*(x)$ utilizando una red neuronal separada, lo que conduce a una optimización acoplada de la política ($\pi_\theta$) y el valor ($V_\phi$).

RPO evita esto mediante la estimación empírica de la función de partición directamente desde los datos. Para un prompt dado $x$, el método agrega todas las respuestas disponibles $y_j$ en el conjunto de datos para calcular una función de partición empírica $\hat{Z}(x)$:
$$\hat{Z}(x) = \sum_{j \in I_x} \pi_{ref}(y_j|x) \exp\left(\frac{1}{\tau} r(x, y_j)\right)$$
donde $I_x$ es el conjunto de índices para las muestras con el prompt $x$. El valor empírico es entonces $\hat{V}(x) = \tau \log \hat{Z}(x)$.

Objetivo de Optimización

RPO formula la tarea de aprendizaje como una regresión directa sobre el logaritmo de la razón entre la política aprendida y la política de referencia. La función de pérdida minimiza el error cuadrático entre el logaritmo de la razón predicha y la recompensa centrada y escalada:
$$\mathcal{L}_{RPO}(\theta) = \frac{1}{2} \mathbb{E}_{(x,y,r) \sim D} \left[ \left( \log \frac{\pi_\theta(y|x)}{\pi_{ref}(y|x)} - \frac{1}{\tau}(r(x, y) - \hat{V}(x)) \right)^2 \right]$$

Las características clave del método incluyen:

• Libre de Valor: No se entrena ninguna red de valor auxiliar; $\hat{V}(x)$ se calcula sobre la marcha a partir de la política de referencia y las recompensas observadas.
• Gradientes Estables: La actualización del gradiente depende únicamente de los parámetros de la política, evitando la retropropagación a través de un estimador de valor aprendido.
• Compatibilidad Offline: El método es totalmente offline y escala a grandes conjuntos de datos sin requerir bucles de RL online.

Contribuciones Clave

1. Propuesta de RPO: Un nuevo método de optimización de política de trayectoria única que reemplaza la estimación de valor con la normalización de recompensa empírica.
2. Análisis de DRO: Una crítica teórica y práctica de la Optimización de Recompensa Directa, resaltando su inestabilidad debido a la optimización conjunta de política y valor.
3. Derivación Formal: Una prueba que muestra cómo RPO aproxima la estructura de la política óptima aprovechando la formulación de la función de partición.
4. Evaluación Extensa: Experimentos exhaustivos en arquitecturas encoder-decoder (Flan-T5) y decoder-only (Mistral, LLaMA, Qwen), demostrando una superioridad consistente sobre SFT, KTO y DRO.

Resultados Experimentales

Los autores evaluaron RPO en el conjunto de datos UltraFeedback y en múltiples benchmarks (AlpacaEval, MT-Bench, IFEval, GSM8K).

Rendimiento Intra-dominio (UltraFeedback)

• Métricas: RPO superó consistentemente a SFT, KTO y DRO en métricas automáticas (BERTScore, ROUGE-L, Distinct-2) y seguridad (Toxicidad).
• LLM-as-a-Judge: RPO logró tasas de victoria más altas contra las líneas base cuando fue evaluado por GPT-4o y Claude-3.5.
• Estabilidad: El análisis de la divergencia KL mostró que RPO mantiene trayectorias suaves y controladas, mientras que DRO exhibió un fuerte crecimiento de la divergencia (especialmente en modelos decoder-only) y KTO mostró oscilaciones.
• Eficiencia: RPO demostró una convergencia más rápida y un menor tiempo de entrenamiento en comparación con DRO y KTO en todas las escalas de modelos.

Generalización Fuera de la Distribución (OOD)

• RPO mostró mejoras robustas pero dependientes de la tarea. Las ganancias fueron más pronunciadas en los benchmarks de seguimiento de instrucciones (AlpacaEval, IFEval).
• El rendimiento en tareas de razonamiento (GSM8K) y conversación de múltiples turnos (MT-Bench) mostró mejoras consistentes pero más moderadas, variando según la familia de modelos.

Estudios de Ablación

• Temperatura ($\tau$): El método es robusto ante un suavizado moderado ($\tau=5.0$), pero sufre una degradación significativa del rendimiento con un escalado muy agudo ($\tau=0.1$), lo que conduce a una optimización inestable.
• Normalización Empírica ($\hat{V}(x)$): Eliminar la normalización de recompensa a nivel de prompt causó una degradación consistente en la similitud semántica, la diversidad y la seguridad, confirmando su papel crítico en la estabilización de las escalas de recompensa.
• Densidad del Conjunto de Datos: El rendimiento decae a medida que disminuye el número de completaciones por prompt, siendo una configuración de densidad del 25% la que causa caídas sustanciales en la calidad. Esto resalta la dependencia del método en distribuciones de recompensa diversas por prompt para una estimación precisa de la partición.
• Robustez al Ruido: RPO permanece estable bajo ruido de recompensa moderado ($\sigma=0.1, 0.5$), pero el rendimiento declina a medida que aumenta la intensidad del ruido, aunque no colapsa incluso bajo una corrupción fuerte ($\sigma=1.0$).

Significado y Limitaciones

El artículo afirma que RPO ofrece una alternativa simple, escalable y estable a los métodos de alineación existentes. Al eliminar la necesidad de aprender la función de valor, reduce la complejidad y la varianza de la optimización, proporcionando al mismo tiempo supervisión de recompensa absoluta. Los autores posicionan a RPO como una dirección prometedora para la optimización de preferencias escalable, particularmente en escenarios donde los datos por pares son escasos pero la retroalimentación escalar está disponible.

Limitaciones reconocidas por los autores:

• Requerimientos de Datos: RPO depende de datos agrupados donde existen múltiples completaciones por prompt para estimar la función de partición de manera efectiva. El rendimiento decae si la diversidad a nivel de prompt es insuficiente.
• Calidad de la Recompensa: Aunque es robusto al ruido moderado, el método sigue dependiendo de la calidad de las anotaciones de recompensa; la corrupción severa impacta negativamente la alineación.
• Alcance de la Evaluación: La evaluación actual se centra principalmente en benchmarks de seguimiento de instrucciones y alineación. Los autores señalan que otros dominios, entornos multilingües y tareas de razonamiento de largo alcance requieren más estudio para evaluar plenamente la generalización.