Provably avoiding over-optimization in Direct Preference Optimization without knowing the data distribution

Enunciado del Problema

La Optimización de Preferencia Directa (DPO) se ha convertido en un estándar para el entrenamiento posterior de modelos de lenguaje grandes (LLM) debido a su simplicidad y capacidad para eludir el modelado explícito de recompensas. Sin embargo, se sabe que DPO sufre de sobre-optimización: a medida que avanza el entrenamiento, el modelo puede maximizar la pérdida de entrenamiento mientras degrada la calidad real de las respuestas generadas en relación con el modelo inicial. Este problema se ve exacerbado por el hecho de que la DPO estándar y el Aprendizaje por Refuerzo a partir de Retroalimentación Humana (RLHF) tratan la recompensa estimada como una verdad fundamental perfecta, ignorando la incertidumbre.

Las garantías teóricas para mitigar la sobre-optimización suelen depender del pesimismo, asegurando que la política aprendida sea robusta ante errores de estimación. Las variantes pesimistas existentes (por ejemplo, DPO regularizado con $\chi^2$, DPO+SFT) logran cotas teóricas sólidas que dependen únicamente de un coeficiente de concentrabilidad de política única ($C^\star$), lo que implica robustez incluso si el conjunto de datos no cubre todo el espacio de políticas. Sin embargo, estas garantías dependen críticamente de conocer la distribución generadora de datos ($\pi_{data}$) o tener acceso a una política de referencia bien cubierta por $\pi_{data}$. En la práctica, este supuesto se viola a menudo, como cuando se ajustan finamente modelos más pequeños con datos de preferencia generados por sistemas más grandes y propietarios (por ejemplo, GPT-4) donde $\pi_{data}$ es desconocido e inaccesible.

La pregunta central de investigación abordada es: ¿Podemos diseñar un algoritmo similar a DPO que sea demostrablemente robusto a la sobre-optimización sin conocimiento de la distribución generadora de datos $\pi_{data}$?

Metodología: PEPO

Los autores introducen PEPO (Optimización de Preferencia basada en Conjunto Pesimista), un algoritmo de un solo paso que logra el pesimismo mediante un conjunto de políticas entrenadas en subconjuntos de datos disjuntos, sin requerir conocimiento de $\pi_{data}$.

1. Entrenamiento del Conjunto:
El conjunto de datos de preferencias $D$ se divide en $L$ subconjuntos disjuntos. Para cada subconjunto $D_\ell$, se entrena una política ligera $\tilde{\pi}_\ell$ utilizando un objetivo pesimista de DPO. Este objetivo modifica la pérdida estándar de DPO desplazando horizontalmente la función sigmoide mediante un parámetro $\lambda_\ell$, reduciendo efectivamente el margen requerido para una probabilidad de victoria. Este desplazamiento induce una estimación conservadora (pesimista) de la brecha de recompensa.

2. Agregación Pesimista:
La política de salida $\pi_{out}$ se define como la solución a un problema de RLHF "virtual" que optimiza la recompensa del peor caso en todo el conjunto. La recompensa implícita para el conjunto se construye tomando el mínimo sobre los $L$ modelos de sus brechas de recompensa estimadas, ajustadas para la distribución de datos en los otros subconjuntos (para garantizar la independencia).
$$r^-(x, a) = \beta \min_{\ell \in [L]} \left( \log \frac{\tilde{\pi}_\ell(a|x)}{\pi_{ref}(a|x)} - \zeta_\ell(x) \right)$$
donde $\zeta_\ell(x)$ cuenta con la contribución de los datos fuera del $\ell$-ésimo subconjunto.

3. Muestreo Eficiente:
La solución de forma cerrada para $\pi_{out}$ implica una constante de normalización sobre todo el espacio de respuestas, lo cual es intratable. Para eludir esto, PEPO emplea muestreo por rechazo (Algoritmo 2). Cualquier miembro del conjunto $\tilde{\pi}_\ell$ sirve como distribución propuesta porque el numerador de $\pi_{out}$ está acotado por $\tilde{\pi}_\ell$. Para aplicaciones prácticas a gran escala, los autores también proponen una aproximación a nivel de token que aplica la agregación pesimista en cada paso de decodificación, evitando la necesidad de muestreo por rechazo a nivel de trayectoria y reduciendo significativamente la latencia de generación.

Contribuciones Clave

Contribuciones Teóricas:

• Complejidad de Muestra sin $\pi_{data}$: En el escenario tabular, PEPO logra garantías de complejidad de muestra que dependen únicamente del coeficiente de concentrabilidad de política única ($C^\star$). Esto contrasta con la DPO estándar, que sufre de un coeficiente de concentrabilidad de todas las políticas ($C_{all}$), y con las variantes existentes de DPO pesimista que requieren conocimiento de $\pi_{data}$.
• Tamaño Óptimo del Conjunto: El análisis caracteriza el tamaño óptimo del conjunto $L$ como el valor mínimo suficiente para garantizar el pesimismo con alta probabilidad. Los límites teóricos muestran que $L$ escala logarítmicamente con los tamaños del espacio de estados y acciones.
• Robustez: Los límites se mantienen sin asumir que el conjunto de datos cubre toda la clase de políticas, abordando la brecha del "$\pi_{data}$ desconocido" en trabajos teóricos anteriores.

Contribuciones Prácticas:

• Validación Empírica: PEPO se evalúa en modelos de 7B, 8B y 34B parámetros (Zephyr-7B, Llama-3.1-8B, Mistral-7B, Yi-34B) utilizando el conjunto de datos UltraFeedback.
• Mitigación de la Sobre-optimización: Los experimentos demuestran que, mientras que la DPO estándar muestra una caída brusca en las tasas de victoria tras ganancias iniciales (sobre-optimización), el rendimiento de PEPO se estabiliza y mejora.
• Superioridad sobre Líneas Base: PEPO supera consistentemente a $\chi^2$PO y DPO+SFT en configuraciones donde $\pi_{data} \neq \pi_{ref}$ (por ejemplo, UltraFeedback), donde las garantías teóricas de esas líneas base son inválidas.
• Eficiencia: La implementación a nivel de token logra un equilibrio favorable entre la tasa de victoria y el tiempo de generación, ofreciendo un rendimiento casi óptimo con una latencia significativamente menor que el muestreo por rechazo completo.

Resultados

• Sobre-optimización: La Figura 1 muestra que la tasa de victoria de DPO contra el modelo inicial alcanza un pico y luego disminuye, mientras que PEPO (con un $L$ apropiado) muestra una mejora sostenida. Se demuestra que la parada temprana basada en la pérdida de validación es poco fiable para seleccionar el mejor punto de control de DPO.
• Comparación: En la Figura 3, PEPO logra tasas de victoria más altas en AlpacaEval2 en comparación con $\chi^2$PO y SFT+DPO en todos los modelos probados. La brecha de rendimiento se estrecha solo cuando $\pi_{data} = \pi_{ref}$, confirmando que la ventaja de PEPO se deriva de su capacidad para manejar distribuciones de datos desconocidas.
• Ablación: Los experimentos confirman que $L$ es un hiperparámetro crítico; un valor demasiado pequeño no induce un pesimismo suficiente, mientras que un valor demasiado grande reduce los datos efectivos por modelo, degradando el rendimiento.

Significado y Afirmaciones

El artículo afirma presentar el primer algoritmo similar a DPO que evita demostrablemente la sobre-optimización sin conocimiento de la distribución generadora de datos. Esto es significativo porque cierra la brecha entre la robustez teórica (concentrabilidad de política única) y las restricciones prácticas ($\pi_{data}$ desconocido), un escenario común en la alineación de LLM en el mundo real donde los datos son generados por modelos propietarios o diversos.

Los autores posicionan a PEPO como una solución práctica que conserva la simplicidad de DPO (entrenamiento de un solo paso, sin modelo de recompensa explícito) mientras incorpora la seguridad teórica de la regularización pesimista. Reconocen que sus garantías teóricas se establecen actualmente en el escenario tabular, y extender estos resultados a clases generales de aproximación de funciones sin conocer $\pi_{data}$ sigue siendo un desafío abierto. Además, el método actual produce una política no paramétrica (mediante muestreo), lo cual es ideal para el paso final de una tubería de entrenamiento pero no para puntos de control intermedios, sugiriendo una dirección para trabajos futuros sobre salidas paramétricas.