ParetoSlider: Diffusion Models Post-Training for Continuous Reward Control

1. El Problema

El ajuste fino post-entrenamiento mediante Aprendizaje por Refuerzo (RL) se ha convertido en el estándar para alinear modelos generativos con preferencias humanas. Sin embargo, la mayoría de los métodos actuales dependen de una única recompensa escalar. Cuando existen múltiples criterios importantes (por ejemplo, adherencia al prompt vs. fidelidad a la imagen original en la edición de imágenes), la práctica predominante es la "escalarización temprana": combinar múltiples recompensas en una suma ponderada fija antes o durante el entrenamiento.

Esto genera dos problemas principales:

1. Punto de operación fijo: El modelo se entrena para un compromiso (trade-off) específico y estático. No puede adaptarse a diferentes preferencias del usuario en tiempo de inferencia sin reentrenar.
2. Inflexibilidad: Obliga al usuario a elegir un compromiso permanente, eliminando la capacidad de navegar dinámicamente entre objetivos en conflicto (como realismo fotográfico vs. estilo de boceto).

2. Metodología: ParetoSlider

Los autores proponen ParetoSlider, un marco de Aprendizaje por Refuerzo Multiobjetivo (MORL) diseñado para modelos de difusión y flujo (flow-matching). Su objetivo es entrenar un único modelo que aproxime toda la frente de Pareto, permitiendo a los usuarios controlar continuamente el equilibrio entre recompensas conflictivas mediante un vector de preferencias.

Componentes Clave:

• Condicionamiento por Vector de Preferencias ($\omega$):
  En lugar de entrenar modelos separados para cada combinación de pesos, ParetoSlider introduce un vector de preferencias $\omega = [\omega_1, ..., \omega_M]$ (donde $\sum \omega_m = 1$) como señal de entrada al modelo durante el entrenamiento. El modelo aprende a mapear cualquier $\omega$ a su punto óptimo correspondiente en la frente de Pareto.

• Estrategia de Escalarización Tardía (Late-Scalarization):
  Un desafío crítico es que las funciones de recompensa raw pueden tener escalas numéricas y varianzas muy diferentes, lo que podría "secuestrar" la señal de gradiente y anular la influencia del vector $\omega$.

  • Solución: El método calcula las ventajas (advantages) para cada recompensa individualmente y las normaliza por separado dentro de cada grupo de muestras.
  • Solo en el paso final de la función de pérdida, las pérdidas individuales se agregan utilizando los pesos $\omega$. Esto asegura que la optimización respete estrictamente la preferencia del usuario y no la escala arbitraria de las recompensas.

• Base de Entrenamiento (DiffusionNFT):
  El método se construye sobre DiffusionNFT, un marco de RL eficiente para modelos de flujo.

  • Utiliza un proceso de tres etapas: (1) Generación de un grupo de muestras bajo la misma condición y preferencia; (2) Evaluación por múltiples modelos de recompensa y cálculo de ventajas normalizadas por canal; (3) Actualización de la política mediante una pérdida de flujo de difusión ponderada por las preferencias.
  • Esto evita la necesidad de redes de valor aprendidas y reduce la varianza mediante la normalización relativa dentro del grupo.

• Arquitecturas de Condicionamiento:
  Se proponen mecanismos ligeros para inyectar $\omega$ en los transformadores de difusión (SD3.5, FluxKontext, LTX-2):

  • Inyección global: A través de la incrustación de tiempo (timestep embedding).
  • Modulación residual: Aplicada a los estados ocultos de las capas del transformador (ej. mediante AdaLN o capas residuales compartidas).
  • Estos módulos se inicializan cerca de cero para no perturbar el modelo preentrenado y se ajustan junto con adaptadores LoRA.

3. Contribuciones Clave

1. Control Continuo en Tiempo de Inferencia: Por primera vez en modelos de difusión, se logra un control suave y continuo sobre el equilibrio de recompensas competitivas sin necesidad de reentrenar ni mantener múltiples checkpoints.
2. Aproximación de la Frente de Pareto con un Solo Modelo: El modelo aprende a navegar todo el espectro de compromisos óptimos, superando la limitación de los métodos de escalarización temprana que quedan atrapados en un solo punto.
3. Estrategia de Escalarización Tardía: Demuestran que normalizar las ventajas por recompensa antes de la agregación ponderada es crucial para evitar el desbalance de escalas y lograr una adherencia fiel a las preferencias del usuario.
4. Versatilidad: El marco se valida exitosamente en tres tareas distintas: Síntesis de Imagen Texto-a-Imagen (T2I), Edición de Imagen Instrucción-a-Imagen (I2I) y Generación de Video Texto-a-Video (T2V).

4. Resultados

El método se evaluó sobre tres backbones de última generación: Stable Diffusion 3.5, FluxKontext y LTX-2.

• Comparación Cualitativa:
  • T2I: Muestra transiciones suaves entre realismo fotográfico y estilos como bocetos, arte vectorial plano, acuarela y anime.
  • I2I: Logra un equilibrio controlado entre la preservación de la identidad de la imagen original y la adherencia estricta a la instrucción de edición (ej. cambiar una mujer a un personaje de fantasía).
  • T2V: Navega entre estilos realistas y animación (estilo Disney/Pixar) en videos generados.
• Comparación Cuantitativa:
  • ParetoSlider iguala o supera el rendimiento de múltiples modelos de línea base entrenados por separado con pesos fijos (Fixed-Weights) en sus respectivos puntos de operación.
  • Supera a métodos de interpolación de modelos (como Rewarded Soups) y métodos de selección de Pareto estáticos (FlowMulti).
  • Utilizando el indicador de Hipervolumen (HV), ParetoSlider domina consistentemente a las líneas base, demostrando una mejor cobertura y convergencia hacia la frente de Pareto óptima.
• Estudios de Ablación:
  • Confirmaron que la escalarización tardía produce transiciones más suaves y fieles que la escalarización temprana o el método Tchebycheff suave.
  • La arquitectura de condicionamiento compartida (Shared Residual) mostró un mejor rendimiento que las variantes basadas en tokens o híbridas.

5. Significado e Impacto

ParetoSlider establece un nuevo paradigma para la alineación multiobjetivo en modelos generativos visuales. Su importancia radica en:

• Eficiencia: Elimina la necesidad costosa de entrenar y almacenar múltiples modelos para diferentes preferencias de usuario.
• Interfaz de Usuario Intuitiva: Permite a usuarios no expertos controlar complejos compromisos de recompensa mediante un simple "deslizador" (slider) en tiempo de inferencia.
• Escalabilidad: Ofrece una solución escalable para problemas de optimización multiobjetivo en IA generativa, permitiendo una adaptación dinámica a necesidades cambiantes sin reentrenamiento.

En resumen, ParetoSlider transforma la alineación de modelos de difusión de un proceso estático y rígido a uno dinámico y controlable, permitiendo una generación de contenido que se adapta fluidamente a las preferencias específicas y cambiantes del usuario.