Guiding Distribution Matching Distillation with Gradient-Based Reinforcement Learning

Resumen Técnico: GDMD

1. El Problema

Los modelos de difusión han revolucionado la síntesis de imágenes, pero su naturaleza iterativa los hace computacionalmente costosos, lo que dificulta la interactividad en tiempo real. Las técnicas de destilación (como Distribution Matching Distillation o DMD) buscan comprimir estos modelos en generadores de pocos pasos (few-step), pero a menudo sacrifican la calidad de la imagen por la velocidad de muestreo.

Intentos recientes han integrado Aprendizaje por Refuerzo (RL) en la destilación (ej. DMDR) para alinear el modelo con preferencias humanas. Sin embargo, estos enfoques presentan limitaciones críticas:

• Evaluación basada en muestras ruidosas: Los métodos existentes evalúan las imágenes generadas crudas ($x_0$) para calcular recompensas. En las etapas tempranas de la destilación, estas muestras son inestables y ruidosas, lo que genera señales de recompensa poco fiables.
• Conflictos de optimización: La fusión ingenua de los objetivos de RL y DMD crea conflictos de gradiente. El RL intenta optimizar la imagen final, mientras que DMD optimiza la alineación con la distribución del profesor, lo que a menudo lleva a inestabilidad o "hacking de recompensas" (reward hacking), donde el modelo explota la función de recompensa sin mejorar la calidad real.
• Dependencia de "Cold Start": Debido a la inestabilidad inicial, los métodos actuales requieren una fase de destilación pura (sin RL) para estabilizar el entrenamiento antes de introducir el RL.

2. Metodología: GDMD

Los autores proponen GDMD, un marco novedoso que redefine el mecanismo de recompensa en la destilación DMD-RL. En lugar de evaluar las imágenes generadas (píxeles), el método prioriza los gradientes de la destilación como la señal principal de optimización.

Los componentes clave del pipeline de GDMD son:

• Cambio de Paradigma (Gradiente vs. Muestra):

  • En lugar de calcular la recompensa sobre la imagen $x_0$, GDMD trata los gradientes derivados del objetivo DMD como "tensores objetivo implícitos".
  • Esto permite que los modelos de recompensa existentes (como CLIP, HPS, ImageReward) evalúen directamente la calidad de las actualizaciones de destilación, no solo el resultado final.

• Componentes Principales:

  1. Recolección de Gradientes Consciente de Destilación (DaGC): Utiliza la aleatoriedad inherente de DMD para recopilar gradientes en grupo. Emplea dos estrategias:
     • Basada en $t$: Estima puntuaciones con diferentes intensidades de ruido para cubrir la distribución temporal.
     • Basada en $\mu_{fake}^\phi$: Rastrea diferentes versiones de la función de puntuación "fake" durante el proceso de optimización.
  2. Puntuación de Gradiente Implícito (IGS): Transforma matemáticamente el gradiente DMD en un objetivo de pérdida ($x_{tar}$) que puede ser decodificado y evaluado por modelos de recompensa estándar. Esto evita la necesidad de muestreo SDE (Ecuaciones Diferenciales Estocásticas) costoso para la recolección de datos.
  3. Optimización de Preferencia Consciente de Negativo (NaPO): Inspirado en DiffusionNFT, utiliza políticas implícitas positivas y negativas para aprender preferencias. Define una probabilidad de optimalidad basada en la diferencia entre la recompensa del gradiente y la recompensa de la muestra actual, optimizando el campo de velocidad del modelo.

• Ventajas Clave:

  • Estabilidad temprana: Al optimizar sobre gradientes en lugar de muestras ruidosas, el entrenamiento es estable desde el inicio, eliminando la necesidad de una fase de "cold start".
  • Alineación de objetivos: Los gradientes de RL están alineados con la dirección de destilación original, evitando la divergencia y el hacking de recompensas.
  • Eficiencia: Elimina la necesidad de procesos de muestreo SDE adicionales para la recolección de datos.

3. Contribuciones Clave

1. Estrategia de Puntuación Basada en Gradientes: Pionera en el uso de gradientes de optimización como señal de recompensa primaria para la destilación DMD-RL, superando las limitaciones de la puntuación basada en muestras.
2. Pipeline Especializado: Diseño de un sistema completo que incluye DaGC, IGS y NaPO para aprovechar eficazmente la guía a nivel de gradiente.
3. Nuevo Estado del Arte (SOTA): Logra un entrenamiento conjunto DMD-RL que produce modelos de pocos pasos (4 pasos) que superan consistentemente a sus profesores de múltiples pasos y a métodos anteriores (DMD, DMDR) en métricas de calidad y alineación con el texto.

4. Resultados Experimentales

Los experimentos se realizaron en modelos basados en SDXL y SD3, utilizando recompensas entrenadas solo en Clip Score y HPS Score.

• Rendimiento Cuantitativo:

  • GDMD supera a todos los métodos existentes (LCM, Flash Diffusion, DMD, DMD2, DMDR) en métricas de calidad no vistas (Aesthetic Score, Pick Score, ImageReward).
  • En el benchmark GenEval, GDMD supera al modelo profesor (de 100 pasos) y a DMDR en todas las dimensiones (conteo, posición, atributos, colores).
  • En la tabla de resultados, GDMD (4 pasos) alcanza un ImageReward de 1.2702 (SD3), superando al profesor (1.0759) y a DMDR (1.0120).

• Rendimiento Cualitativo y Humano:

  • Alineación con el Prompt: GDMD muestra una alineación superior con las instrucciones del texto, reduciendo alucinaciones y errores de texto comparado con DMD y DMDR.
  • Calidad de Imagen: Mitiga problemas comunes como distorsión de extremidades, saturación excesiva y artefactos.
  • Preferencia Humana: En estudios de preferencia humana, GDMD (4 pasos) gana al modelo profesor (100 pasos) en un 55.1% de los casos en calidad de imagen y supera a DMDR y DMD+NFT en más del 60% de las comparaciones.

5. Significado e Impacto

GDMD representa un avance significativo en la síntesis de imágenes eficientes. Al resolver el conflicto fundamental entre la velocidad de la destilación y la calidad guiada por RL, demuestra que es posible entrenar generadores de 4 pasos que no solo igualan, sino que superan la calidad de los modelos de difusión tradicionales de 50-100 pasos.

La metodología de puntuación basada en gradientes ofrece una solución robusta a la inestabilidad del entrenamiento en etapas tempranas, eliminando la necesidad de fases de calentamiento complejas y permitiendo una escalabilidad más eficiente hacia modelos de generación en tiempo real de alta fidelidad.