Unconditional Priors Matter! Improving Conditional Generation of Fine-Tuned Diffusion Models

1. El Problema: Degradación del Prior Incondicional

El éxito de los modelos de difusión condicional (como los que generan imágenes a partir de texto o ediciones de imágenes) se basa en la técnica de Classifier-Free Guidance (CFG). Durante el entrenamiento, el modelo aprende simultáneamente a predecir el ruido de forma condicional ($\epsilon_\theta(x_t, c)$) y de forma incondicional ($\epsilon_\theta(x_t, \emptyset)$), aplicando un dropout de la condición en un pequeño porcentaje de las iteraciones.

Los autores identifican un problema crítico: cuando un modelo base (como Stable Diffusion) se ajusta (fine-tuning) para una tarea específica (por ejemplo, síntesis de vistas 3D o edición de imágenes), la capacidad de la red para representar la distribución incondicional se degrada drásticamente. Esto ocurre porque:

1. El modelo comparte su capacidad limitada entre la nueva tarea condicional y la tarea incondicional.
2. El ancho de banda para el aprendizaje incondicional es muy bajo debido al bajo porcentaje de dropout.

Consecuencia: Esta "pobreza" en el ruido incondicional afecta negativamente la calidad de la generación condicional, ya que el CFG utiliza la diferencia entre ambos ruidos para guiar la muestra. Un ruido incondicional de baja calidad introduce errores en la guía, resultando en imágenes con colores saturados, distorsiones de forma o falta de realismo semántico.

2. Metodología: Reemplazo de Ruido Incondicional

La propuesta es una solución sin entrenamiento adicional (training-free) que consiste en sustituir el ruido incondicional degradado del modelo ajustado por un ruido incondicional proveniente de un modelo con un "prior" más rico y robusto.

La técnica propuesta:
En lugar de usar la formulación estándar de CFG:
$$\epsilon_\theta^{(\gamma)}(x_t, c) = \epsilon_\theta(x_t, \emptyset) + \gamma(\epsilon_\theta(x_t, c) - \epsilon_\theta(x_t, \emptyset))$$

Los autores proponen utilizar el ruido de un modelo base ($\psi$) o de cualquier otro modelo preentrenado con buenos conocimientos incondicionales:
$$\epsilon_{\theta, \psi}^{(\gamma)}(x_t, c) = \epsilon_\psi(x_t, \emptyset) + \gamma(\epsilon_\theta(x_t, c) - \epsilon_\psi(x_t, \emptyset))$$

Hallazgo sorprendente: El modelo de reemplazo no tiene que ser necesariamente el modelo base original del cual se derivó el modelo ajustado; puede ser cualquier modelo de difusión preentrenado (incluso con arquitecturas distintas, como pasar de un UNet a un DiT como PixArt-$\alpha$).

3. Contribuciones Clave

• Identificación de un fallo sistémico: Demuestran que la degradación del prior incondicional durante el fine-tuning es una causa principal de la baja calidad en modelos condicionales.
• Solución Plug-and-Play: Proponen un método de muestreo que no requiere reentrenar el modelo, lo que lo hace extremadamente eficiente y aplicable a modelos existentes.
• Versatilidad de modelos: Demuestran que el método funciona independientemente de la arquitectura (UNet o DiT) y de la modalidad (imagen, video o edición).

4. Resultados Experimentales

El método fue validado en cinco modelos de difusión condicional de gran escala:

1. Zero-1-to-3 (Síntesis de vistas 3D): Mejora significativamente la iluminación y reduce las distorsiones de forma (medido con LPIPS y PSNR).
2. Versatile Diffusion (Variaciones de imagen): Reduce la saturación excesiva y las distorsiones de objetos, mejorando el FID y la calidad visual.
3. DiT (Generación condicional a clases): Mejora el FID en tareas de clasificación, especialmente en datasets grandes.
4. DynamiCrafter (Texto/Imagen a Video): Mejora la consistencia temporal y la calidad estética de los videos generados.
5. InstructPix2Pix (Edición de imágenes): Logra ediciones más fieles a las instrucciones de texto manteniendo mejor la identidad del objeto original.

5. Significado y Conclusión

Este trabajo cambia la perspectiva sobre el entrenamiento de modelos condicionales. Sugiere que no es necesario que un modelo aprenda de forma conjunta la condición y la ausencia de ella si se dispone de un modelo externo que ya posea un conocimiento incondicional sólido.

Aunque el método introduce un ligero aumento en el uso de memoria (al cargar un segundo modelo) y un costo computacional marginal durante la inferencia, la mejora sustancial en la calidad de la generación justifica ampliamente el intercambio (trade-off). Es un avance fundamental para mejorar la utilidad de los modelos de difusión especializados sin necesidad de costosos procesos de reentrenamiento.