RLHF in an SFT Way: From Optimal Solution to Reward-Weighted Alignment

Resumen Técnico: Variational Alignment with Re-weighting (VAR)

1. El Problema

El Alineamiento de Modelos de Lenguaje (LLM) mediante Aprendizaje por Refuerzo con Retroalimentación Humana (RLHF) es fundamental para alinear los modelos con valores humanos (utilidad y seguridad). Sin embargo, los métodos actuales presentan desafíos significativos:

• Complejidad y Costo Computacional: Los métodos en línea basados en muestreo, como PPO (Proximal Policy Optimization) y GRPO, requieren múltiples modelos (política, referencia, crítico y recompensa) y son computacionalmente costosos e inestables.
• Inestabilidad en Métodos Offline: Las alternativas offline como DPO (Direct Preference Optimization) y KTO simplifican el proceso pero sufren de problemas de sobreajuste (over-fitting) e inestabilidad en el entrenamiento.
• Limitaciones de los Pesos Negativos: Métodos como DPO y A-LoL pueden asignar pesos negativos a respuestas no preferidas. Esto crea un paisaje de pérdida no acotado (unbounded loss landscape), lo que lleva a inestabilidad durante la optimización.
• Recorte (Clipping): Enfoques como PPO y A-LoL utilizan técnicas de recorte (clipping) en los pesos de importancia, lo que aplanan las distinciones de recompensa y dificultan la resolución de preferencias a granularidad fina.

2. Metodología: Variational Alignment with Re-weighting (VAR)

Los autores proponen VAR, un nuevo marco de alineamiento offline derivado de la inferencia variacional. En lugar de optimizar la política directamente mediante gradientes de RL complejos, VAR reformula el problema como una aproximación de la solución óptima teórica de RLHF.

• Fundamento Teórico:

  • Partiendo de la solución de forma cerrada de RLHF (ecuación de Boltzmann), el objetivo es minimizar la divergencia KL hacia adelante ($KL[\pi^* \parallel \pi_\theta]$) entre la política óptima $\pi^*$ y la política aprendida $\pi_\theta$.
  • Esto transforma el objetivo de alineamiento en una tarea de Ajuste Fino Supervisado (SFT) ponderado por recompensas.
  • La función de pérdida resultante es:
    $$\mathcal{L}_{VAR} = -\mathbb{E}_{\pi_{ref}} \left[ \frac{\exp(r(x, y)/\beta)}{Z(x)} \cdot \log \pi_\theta(y|x) \right]$$
    Donde el término $\frac{\exp(r(x, y)/\beta)}{Z(x)}$ actúa como un peso no negativo derivado de la recompensa.

• Innovaciones Clave:

  1. Pesos No Negativos: Al utilizar una transformación exponencial de la recompensa, los pesos son estrictamente positivos. Esto garantiza un paisaje de optimización estable y bien definido, eliminando los problemas de inestabilidad causados por pesos negativos.
  2. Estimación In-Batch de la Función de Partición ($Z(x)$): Calcular $Z(x)$ (la función de partición) es costoso. VAR introduce una técnica eficiente de estimación por lotes (in-batch) que aproxima $Z(x)$ utilizando las muestras disponibles dentro del mismo micro-lote de entrenamiento mediante estimación de Monte Carlo. Esto evita el muestreo en línea y permite una implementación escalable.
  3. Eliminación de Recorte: Al no depender de la estimación de ventajas o recortes de importancia, el método permite que las diferencias finas en las recompensas se traduzcan directamente en diferencias en los pesos de entrenamiento.

3. Contribuciones Clave

• Reformulación Variacional: Transforman el problema de RLHF en un problema de inferencia variacional sobre medidas positivas, conectando teóricamente la solución óptima de RL con un objetivo de SFT ponderado.
• Eficiencia Computacional: Logran un rendimiento comparable o superior a los métodos de RL en línea (como GRPO) pero con una complejidad de implementación similar a la de un SFT estándar.
• Estabilidad Mejorada: Demuestran que el enfoque de pesos no negativos y la estimación in-batch eliminan la inestabilidad de entrenamiento observada en DPO y A-LoL.
• Código Abierto: El código está disponible públicamente, facilitando la reproducibilidad.

4. Resultados Experimentales

Los autores evaluaron VAR en múltiples benchmarks (HHA, MMLU, GSM8k, HumanEval, BBH) utilizando modelos Llama y Qwen de diversos tamaños.

• Rendimiento en Alineamiento (HHA):
  • VAR superó consistentemente a DPO en métricas de "Ayuda" (Helpfulness) y "Inofensividad" (Harmlessness).
  • En promedio, VAR mostró una mejora del 7.16% sobre DPO en métricas de ayuda.
  • Funcionó bien tanto iniciando desde modelos base (pre-entrenados) como desde modelos SFT, demostrando robustez.
• Comparación con Métodos en Línea (GRPO):
  • VAR logró un rendimiento comparable o superior a GRPO (un método de RL en línea eficiente).
  • Velocidad: VAR es más de 5 veces más rápido que GRPO en términos de tiempo de entrenamiento por época, reduciendo drásticamente el costo computacional.
• Estabilidad de Entrenamiento:
  • Las curvas de validación mostraron que VAR tiene una convergencia más suave y estable en comparación con la volatilidad y declive temprano observados en DPO.
• Benchmarks Generativos:
  • En tareas de razonamiento y codificación (MMLU, GSM8k, HumanEval), VAR superó a DPO y A-LoL, manteniendo o mejorando las capacidades generales del modelo sin colapsar.
• Análisis de Longitud:
  • A diferencia de DPO, que tiende a generar secuencias excesivamente largas para maximizar recompensas (alucinación de longitud), VAR mantiene longitudes de salida más naturales y similares a la versión SFT, sin sacrificar la calidad.

5. Significado e Impacto

El trabajo VAR representa un avance significativo en la eficiencia y estabilidad del alineamiento de LLMs:

• Puente entre Eficiencia y Rendimiento: Demuestra que no es necesario recurrir a complejos procesos de RL en línea (como PPO/GRPO) para lograr un alineamiento de alta calidad. Se puede lograr un rendimiento de nivel RL mediante una modificación simple y estable del SFT.
• Simplificación de la Infraestructura: Al eliminar la necesidad de modelos críticos, muestreo en línea y técnicas de recorte, VAR reduce la barrera de entrada para implementar alineamiento de RLHF en entornos con recursos limitados.
• Solución a la Inestabilidad: Aborda uno de los problemas más persistentes en la alineación offline (la inestabilidad por pesos negativos), ofreciendo un marco matemáticamente más sólido basado en medidas positivas.

En conclusión, VAR propone un nuevo paradigma donde el RLHF se entiende y se implementa como un problema de inferencia variacional, logrando un equilibrio superior entre la calidad del alineamiento, la estabilidad del entrenamiento y la eficiencia computacional.