Resumen Técnico: Modulación Consciente de la Calidad para Transformadores de Difusión

Declaración del Problema

Los modelos de difusión modernos de texto a imagen, particularmente los Transformadores de Difusión (DiT) como Stable Diffusion 3.5 (SD3.5), dependen de señales de condicionamiento —específicamente, embeddings de paso de tiempo (timestep) y embeddings de prompt agrupados (pooled)— para modular el proceso de eliminación de ruido mediante la Normalización de Capa Adaptativa (AdaLN). Si bien estas señales comunican eficazmente el nivel de ruido actual y la intención semántica, carecen de información consciente de la calidad. En consecuencia, el mecanismo de modulación no puede reaccionar directamente a los modos de falla específicos de la muestra, tales como la acumulación de errores, la interpretación incorrecta del prompt o las características alucinadas que emergen a medida que la imagen latente evoluciona. El condicionamiento base es independiente del estado latente actual ($\partial M_{base} / \partial x_t = 0$), lo que impide que el modelo corrija problemas de fidelidad basados en la calidad real de la imagen intermedia.

Metodología

Los autores proponen el Módulo de Representación de Calidad (QRM), un adaptador ligero basado en transformadores diseñado para inyectar señales conscientes de la calidad en el proceso de difusión sin modificar el backbone preentrenado.

Arquitectura

El QRM se implementa como un decodificador de transformador que predice desplazamientos de modulación aditivos ($M_{qrm}$) a los parámetros de la AdaLN base ($M_{base}$).

• Entradas: El módulo recibe la imagen latente actual ($x_t$), el embedding de prompt agrupado ($p$), el embedding del paso de tiempo ($t$) y los parámetros de modulación base ($M_{base}$).
• Procesamiento:
  1. La imagen latente $x_t$ se proyecta en tokens latentes mediante una convolución $1\times1$ aprendida y un agrupamiento (pooling) adaptativo.
  2. Los parámetros de modulación base se normalizan y se proyectan para formar tokens de consulta ($Z_0$).
  3. Se construye una memoria de clave-valor ($KV$) a partir de los tokens latentes, el prompt y los embeddings del paso de tiempo.
  4. A través de capas de atención cruzada y autoatención, los tokens de consulta atienden a la memoria $KV$ para generar una representación refinada.
  5. El resultado se proyecta de nuevo a las anchuras nativas de la AdaLN para producir el vector de desplazamiento $M_{qrm}$.
• Aplicación: La modulación final aplicada a los bloques del transformador es $M_t = M_{base,t} + M_{qrm,t}$.

Estrategia de Entrenamiento (ReFL)

El QRM se entrena utilizando una estrategia de Aprendizaje de Retroalimentación de Recompensa (ReFL):

• Componentes Congelados: El backbone de difusión (denoiser, codificadores de texto, VAE), el muestreador (sampler) y el modelo de recompensa permanecen congelados. Solo se actualizan los parámetros del QRM.
• Objetivo: El módulo aprende a maximizar una puntuación de recompensa escalar derivada de un modelo de recompensa congelado (por ejemplo, HPSv2.1 o CLIPScore) aplicado a una imagen decodificada de un paso de eliminación de ruido intermedio.
• Funciones de Pérdida: El artículo utiliza una pérdida de maximización de recompensa o una pérdida de búsqueda de margen, la cual empuja a la imagen modulada por el QRM a superar a la base por un margen específico.
• Selección de Paso de Tiempo: Para mejorar la eficiencia y centrarse en la formación estructural, el QRM se activa solo durante un subconjunto predefinido de pasos de tiempo de alto ruido (por ejemplo, $t \in \{50, \dots, 26\}$), en lugar de todo el esquema de difusión.

Inferencia

Durante la inferencia, el módulo QRM entrenado actúa como un adaptador plug-in. Predice desplazos de modulación basados en el estado latente actual y los añade a los parámetros base. Esto permite al modelo dirigir la trayectoria de eliminación de ruido hacia estados de mayor calidad sin reentrenar el backbone ni alterar el esquema de muestreo.

Contribuciones Clave

1. Arquitectura QRM: Introducción de un módulo de transformador ligero que aprende representaciones conscientes de la calidad a partir de las entradas existentes del modelo (latente, prompt, paso de tiempo) para refinar los parámetros de modulación.
2. Integración No Invasiva: Un mecanismo para aumentar la modulación del transformador de difusión con señales de calidad de la imagen latente sin modificar el backbone de difusión, el VAE o el esquema de muestreo subyacente.
3. Modulación Guiada por Recompensa: Demostración de que los desplazos de modulación predichos ($M_{qrm}$) pueden refinar eficazmente el proceso de eliminación de ruido en cada paso de tiempo, mejorando la adherencia al prompt y la fidelidad de la imagen.
4. Evaluación Exhaustiva: Estudios de ablación extensos que analizan los efectos del tamaño de la arquitectura QRM, los rangos de pasos de tiempo activos y las métricas de recompensa de entrenamiento en la calidad de la imagen.

Resultados Experimentales

Los autores evaluaron el QRM utilizando Stable Diffusion 3.5 (SD3.5) en el benchmark Parti-Prompts.

• Desempeño Cuantitativo:
  • Comparado con la base de SD3.5, el QRM entrenado con HPSv2.1 como señal de recompensa logró las puntuaciones más altas entre todos los métodos evaluados (incluyendo refinamiento de prompt, ajuste fino y enfoques ReNeg).
  • Mejoras: +0.20 en la puntuación Estética y +0.82 en la puntuación HPSv2.1.
  • Cuando se entrenó con CLIPScore, el QRM mejoró la alineación con la preferencia humana (HPSv2.1) manteniendo una calidad perceptual comparable.
• Análisis Cualitativo:
  • Las visualizaciones de las trayectorias de eliminación de ruido muestran que aplicar QRM durante los pasos de tiempo iniciales de alto ruido conduce a que las estructuras semánticas reconoccibles emerjan significativamente antes que en la base.
  • Aplicar QRM solo en las etapas medias o tardías produjo mejoras menores, confirmando que la estructura global se establece temprano en el proceso de difusión.
• Hallazgos de Ablación:
  • Capacidad: Los modelos QRM más grandes optimizados más fuertemente para la recompensa de entrenamiento específica (por ejemplo, HPSv2.1) mostraron ligeros compromisos en otras métricas (por ejemplo, CLIPScore).
  • Rango de Pasos de Tiempo: La modulación durante los pasos de tiempo tempranos ($50 \to 26$) fue consistentemente más efectiva para la alineación que la modulación en etapas posteriores.
  • Métrica de Recompensa: La elección del modelo de recompensa (CLIPScore vs. HPSv2.1) moldeó directamente los sesgos perceptuales aprendidos por el QRM.

Significancia y Reivindicaciones

El artículo afirma que las intervenciones a nivel de modulación proporcionan una vía escalable y efectiva para alinear los grandes modelos de difusión con las señales de preferencia humana. A diferencia de los métodos que requieren un ajuste fino costoso de todo el backbone (como LoRA o el ajuste fino completo), el QRM logra mejoras consistentes de calidad operando enteramente a través de actualizaciones de modulación aditivas condicionadas a la representación latente en evolución.

Los autores posicionan al QRM como un adaptador flexible y consciente de la calidad que puede ser dirigido por diferentes señales de recompensa y aplicado a diversas arquitecturas de transformadores de difusión. Reconocen que el QRM hereda los sesgos de los modelos de recompensa de entrenamiento y sugieren que el trabajo futuro podría explorar el entrenamiento de múltiples recompensas y la selección de pasos de tiempo adaptativos para mejorar aún más la robustez y la generalización. El trabajo demuestra que se pueden lograr ganancias significativas en la fidelidad de la imagen y la alineación con el prompt sin alterar los pesos preentrenados de los modelos de vanguardia.