Resumen Técnico: DiNa-LRM – Modelado de Recompensa Latente Nativo de Difusión

1. Declaración del Problema

La alineación de modelos de difusión y de ajuste de flujo con preferencias humanas depende en gran medida de funciones de recompensa. Los enfoques actuales más avanzados utilizan predominantemente Modelos Visuales-Lingüísticos (VLM) como columnas vertebrales de recompensa. Aunque los VLM ofrecen capacidades discriminatorias robustas, presentan dos limitaciones críticas para la alineación de modelos de difusión:

1. Sobrecarga Computacional y de Memoria: Las recompensas basadas en VLM son costosas de consultar repetidamente durante la optimización, creando un cuello de botella para la eficiencia del entrenamiento.
2. Desajuste de Dominio: Los VLM suelen operar en el espacio de píxeles, mientras que los generadores de difusión latente se entrenan y optimizan en el espacio latente. Este desajuste de píxeles a latente complica la alineación, particularmente para los métodos de gradiente de recompensa, e introduce una sobrecarga innecesaria en el sistema.

Por el contrario, aunque los modelos de difusión han demostrado ser prometedores como columnas vertebrales discriminatorias, los intentos previos de utilizarlos como recompensas a menudo han estado vinculados a algoritmos específicos de optimización a nivel de paso o se han centrado en estados intermedios ruidosos sin establecer un marco de propósito general comparable a los VLM para puntuar muestras limpias.

2. Metodología: DiNa-LRM

Los autores proponen DiNa-LRM (Modelo de Recompensa Latente Nativo de Difusión), un marco que formula el aprendizaje de preferencias directamente sobre estados de difusión ruidosos dentro del espacio latente. El método consta de tres componentes principales:

A. Formulación de Preferencia Nativa de Difusión

En lugar de evaluar muestras limpias ($x_0$) directamente, DiNa-LRM extiende el modelo de Thurstone a estados de difusión ruidosos ($x_t$).

• Incertidumbre Calibrada por Ruido: Los autores reconocen que los juicios de preferencia sobre estados ruidosos son inherentemente menos ciertos a medida que aumenta el nivel de ruido. Introducen una incertidumbre de comparación calibrada por ruido $\sigma_u^2(t)$ que escala con el nivel de ruido de difusión $\sigma^2(t)$:
  $$\sigma_u^2(t) = k\sigma^2(t) + \sigma_u^2$$
  Esta formulación asegura que el objetivo de recompensa permanezca consistente con la variedad ruidosa encontrada durante el preentrenamiento de difusión. Actúa como un regularizador consciente de la incertidumbre, evitando que gradientes de alta varianza desestabilicen el entrenamiento en regímenes de alto ruido.
• Objetivo de Entrenamiento: El modelo se optimiza utilizando una pérdida de fidelidad sobre pares de preferencia $(x_t^+, x_t^-)$, donde la probabilidad de comparación está modulada por la incertidumbre dependiente del paso de tiempo.

B. Arquitectura de Recompensa Latente

DiNa-LRM se instaura sobre una columna vertebral de difusión latente preentrenada (por ejemplo, SD3.5-Medium) con una cabeza de puntuación ligera basada en consultas.

• Extracción de Características Multicapa: Se extraen características visuales y textuales de múltiples capas de la columna vertebral.
• Adaptación Condicionada al Paso de Tiempo: Las características se modulan mediante incrustaciones de pasos de tiempo utilizando FiLM (Modulación Lineal Específica de Características) para hacer que la cabeza de recompensa sea explícitamente consciente del nivel de ruido.
• Puntuación con Q-Former con Puerta: Un Q-Former con puerta agrega estas características multicapa y condicionadas al paso de tiempo para producir una puntuación de recompensa escalar. Esta arquitectura opera enteramente en el espacio latente del VAE, evitando la decodificación en el espacio de píxeles.

C. Escalado en Tiempo de Inferencia mediante Ensamblaje de Ruido

Para mejorar la robustez y reducir la sensibilidad a un paso de tiempo de evaluación específico, DiNa-LRM soporta ensamblaje de ruido.

• En lugar de puntuar una muestra limpia en un único nivel de ruido fijo, el modelo evalúa la muestra en múltiples pasos de tiempo $\{t_k\}$.
• Las características extraídas en estos diferentes pasos de tiempo se concatenan y procesan por la misma cabeza Q-Former para producir una puntuación agregada. Esto sirve como un mecanismo de escalado en tiempo de prueba "nativo de difusión".

3. Contribuciones Clave

1. Formulación de Preferencia Nativa de Difusión: El artículo extiende el modelo de preferencia de Thurstone a estados de difusión ruidosos mediante la introducción de una incertidumbre calibrada por ruido que escala con el nivel de ruido de difusión, cerrando la brecha distribucional entre el preentrenamiento y el aprendizaje de recompensas.
2. Escalado en Tiempo de Inferencia: Los autores proponen una estrategia de ensamblaje de ruido que agrega características de múltiples pasos de tiempo, proporcionando un mecanismo para el escalado en tiempo de prueba que aprovecha la trayectoria global de difusión en lugar de una única estimación puntual.
3. Validación Empírica: El trabajo demuestra que las recompensas nativas de difusión pueden lograr un rendimiento competitivo con jueces basados en VLM potentes mientras operan nativamente en el espacio latente.

4. Resultados Experimentales

Los autores evaluaron DiNa-LRM en varios puntos de referencia y escenarios de alineación:

• Precisión del Modelado de Recompensa:

  • DiNa-LRM supera sustancialmente a las líneas base existentes de recompensas basadas en difusión (por ejemplo, LRM-SD1.5, LRM-SDXL) en ImageReward, HPDv2, HPDv3 y GenAI Bench.
  • Logra un rendimiento competitivo con las recompensas basadas en VLM más avanzadas (por ejemplo, HPSv3, UnifiedReward) en promedio, con una brecha notable que se estrecha en conjuntos de datos específicos como HPDv2 y HPDv3.
  • Ensamblaje de Ruido: Aplicar el ensamblaje de ruido mejora aún más la precisión promedio (por ejemplo, de 71.49% a 72.48% en promedio), validando el beneficio de agregar evidencia a través de niveles de ruido.

• Alineación de Preferencias (ReFL):

  • En experimentos de alineación post-entrenamiento utilizando ReFL (Aprendizaje de Retroalimentación de Recompensa) en SD3.5-M, DiNa-LRM demostró una convergencia más rápida en puntuaciones proxy y métricas "doradas" mejoradas (PickScore retenido) en comparación con HPSv3.
  • Eficiencia: DiNa-LRM reduce significativamente los costos de recursos en comparación con las recompensas basadas en VLM. Reduce el uso máximo de VRAM en un 51.4%, los FLOPs de cálculo de recompensa en un 71.1% y los FLOPs de la fase de optimización en un 46.4%.

• Generalización:

  • La formulación se generaliza a través de diferentes columnas vertebrales de difusión (SDXL, FLUX.1-Dev, Z-Image-Turbo), superando consistentemente a las líneas base anteriores basadas en difusión.
  • Los experimentos de transferibilidad muestran que un modelo de recompensa entrenado en SD3.5-M puede alinear efectivamente SD3.5-L dentro del mismo espacio latente.

5. Significado y Afirmaciones

El artículo afirma que DiNa-LRM ofrece una alternativa práctica y eficiente a las recompensas de VLM en espacio de píxeles para la alineación de preferencias. Su significado principal radica en:

• Ventaja Estructural: Al operar nativamente en el espacio latente, elimina el desajuste de latente a píxel y la sobrecarga computacional de la decodificación, haciéndolo altamente compatible con métodos de alineación por gradiente de recompensa.
• Amigable con la Optimización: Proporciona una señal diferenciable y eficiente que mejora la dinámica del entrenamiento y reduce los requisitos de recursos sin sacrificar la calidad de la alineación.
• Escalabilidad: El mecanismo de ensamblaje de ruido proporciona una forma novedosa y nativa de difusión para escalar el rendimiento en tiempo de inferencia.

Los autores reconocen limitaciones, señalando que la recompensa está vinculada a espacios latentes específicos (limitando la transferencia entre columnas vertebrales sin reentrenamiento) y que el modelado en espacio latente puede subestimar ciertos artefactos a nivel de píxel, lo que potencialmente conduce a comportamientos de hackeo de recompensas (por ejemplo, inyección humana espuria o deriva estilística) durante la optimización de largo horizonte. Sugieren que el trabajo futuro debe centrarse en columnas vertebrales unificadas más fuertes e incorporar una regularización ligera en el espacio de píxeles.