It's Never Too Late: Noise Optimization for Collapse Recovery in Trained Diffusion Models

1. Planteamiento del Problema

Los modelos de difusión texto-a-imagen contemporáneos (por ejemplo, SDXL-Turbo, Flux.1) sufren de colapso de modos. Al muestrear múltiples imágenes desde el mismo prompt de texto utilizando diferentes semillas aleatorias (ruido i.i.d.), los modelos a menudo producen salidas casi idénticas con variación visual mínima.

• Limitaciones actuales: Las soluciones existentes dependen de:
  • Mecanismos de guía: Dirigen el modelo mediante condicionamiento modificado (por ejemplo, Guía de Partículas, DPP), lo que a menudo sacrifica la calidad por diversidad.
  • Métodos basados en búsqueda: Generan un gran grupo de candidatos y los podan (por ejemplo, "Mejor de N"), lo cual es computacionalmente costoso y consume mucha memoria.
  • Ingeniería de prompts: Modifican la entrada de texto, lo que no puede capturar fácilmente todas las variaciones en el espacio de salida.
• Pregunta central: ¿Podemos optimizar directamente el ruido inicial ($x_0$) para recuperarnos del colapso de modos y generar imágenes diversas y de alta calidad sin alterar el modelo preentrenado ni el prompt del usuario?

2. Metodología

Los autores proponen un marco de optimización de ruido de extremo a extremo que trata el vector de ruido inicial como un parámetro aprendible. El proceso implica:

A. Objetivo de Optimización de Ruido

Dado un prompt fijo $c$ y un modelo de difusión congelado $g_\theta$, el método optimiza un lote de muestras de ruido inicial $\{x_0^{(i)}\}_{i=1}^B$. La función objetivo minimiza una pérdida $\mathcal{L}$ compuesta por tres términos:

1. Recompensa de Calidad ($r_s$): Asegura que las imágenes individuales se adhieran al prompt y mantengan alta fidelidad (por ejemplo, usando CLIPScore o HPSv2).
2. Métrica de Diversidad ($v_\mathcal{B}$): Maximiza la disimilitud entre las imágenes del lote. El artículo explora diversas métricas:
   • Por pares: Parches de características DINOv2, DreamSim, LPIPS, Histogramas de Color.
   • A nivel de conjunto: Procesos de Puntos Determinantes (DPP) y Puntuación Vendi, que miden el número efectivo de elementos únicos en un conjunto.
3. Regularización ($\text{reg}$): Evita que el ruido optimizado se desvíe demasiado de la distribución a priori $N(0, I)$. Esto se logra penalizando las desviaciones de la distribución $\chi_d$ de la norma del ruido.

La optimización se realiza mediante retropropagación a través del muestreador congelado $g_\theta$.

B. Estrategia de Inicialización de Ruido (Ruido Rosa)

Una idea clave del artículo es el análisis de la evolución del ruido. Durante la optimización, los autores observaron que los cambios más significativos ocurren en los componentes de baja frecuencia del espectro del ruido, mientras que las altas frecuencias permanecen relativamente estables.

• Hipótesis: Las imágenes naturales siguen un espectro de potencia $1/f$ (más potencia en bajas frecuencias). El ruido blanco estándar es espectralmente plano.
• Solución: Los autores introducen la inicialización con Ruido Rosa. Transforman el ruido blanco al dominio de la frecuencia mediante FFT, aplican un factor de escalado de potencia de $f^{-\alpha/2}$ (donde $\alpha \approx 0.2$) y transforman de nuevo.
• Beneficio: Esta inicialización alinea el ruido de partida con las propiedades estadísticas de las imágenes naturales y la dinámica de optimización, lo que conduce a una convergencia más rápida y una mayor diversidad.

C. Modos de Optimización

• Optimización por Lotes: Optimiza conjuntamente un conjunto de imágenes (por ejemplo, 4) para maximizar la diversidad a nivel de conjunto.
• Optimización Secuencial: Genera imágenes una por una, optimizando cada una para que difiera de las salidas generadas previamente. Esto reduce la sobrecarga de memoria y permite generar grandes conjuntos diversos.

3. Contribuciones Clave

1. Optimización de Ruido para Diversidad: Demuestra que optimizar directamente el ruido inicial es una estrategia superior para recuperarse del colapso de modos en comparación con los métodos basados en guía o búsqueda.
2. Objetivos de Diversidad a Nivel de Conjunto: Identifica que DPP y la Puntuación Vendi son más efectivos que las métricas simples por pares para optimizar conjuntos diversos, ya que evitan que la métrica "promedio por pares" sea manipulada haciendo que una imagen sea muy diferente mientras se mantienen las demás similares.
3. Inicialización con Ruido Rosa: Propone una estrategia simple y libre de parámetros (usando ruido rosa) que mejora consistentemente la diversidad en diferentes modelos (SDXL-Turbo, Flux.1) y líneas base, incluso sin optimización adicional.
4. Compensación Calidad-Diversidad: Muestra que su método puede aumentar la diversidad sin sacrificar la calidad de la imagen ni la adherencia al prompt, superando la compensación típica observada en los métodos basados en guía.

4. Resultados Experimentales

El método fue evaluado en GenEval, DPG-Bench (prompts complejos) y T2I-CompBench utilizando modelos como SDXL-Turbo y Flux.1 [schnell].

• Rendimiento Cuantitativo:
  • Diversidad: Supera al muestreo i.i.d. y al método de búsqueda más avanzado [Parmar et al., 2025] en todas las métricas (DINO, DreamSim, LPIPS, Puntuación Vendi).
  • Calidad: Mantiene o mejora ligeramente las puntuaciones de calidad de imagen (HPSv2, CLIPScore) en comparación con las líneas base.
  • Eficiencia: Con la inicialización de ruido rosa, el método alcanza puntuaciones de diversidad más altas que el método de búsqueda de línea base en significativamente menos iteraciones (por ejemplo, 6–8 iteraciones frente a 9–12 para ruido blanco).
• Resultados Cualitativos:
  • Genera variaciones diversas en la pose del objeto, color, estilo y fondo, manteniendo una adherencia estricta a prompts complejos.
  • Maneja con éxito prompts complejos (por ejemplo, "Una escena pintoresca de otoño...") donde el muestreo i.i.d. falla en producir variación.
• Estudio con Usuarios: Los evaluadores humanos prefirieron las salidas del método propuesto (especialmente aquellas optimizadas con DPP/Vendi) sobre el muestreo i.i.d. y la línea base de búsqueda en un 85–91% de los casos.

5. Significado

Este trabajo cambia el paradigma de la generación de diversidad en modelos de difusión, pasando de estrategias de muestreo (esperar una buena semilla) o guía (modificar la ruta de generación) a la optimización de entrada.

• Aplicabilidad General: El enfoque es agnóstico al modelo y funciona tanto en modelos de difusión destilados por pasos (rápidos) como en modelos estándar.
• Eficiencia: Al optimizar el ruido directamente, evita la enorme sobrecarga de memoria de generar y almacenar cientos de candidatos para su poda.
• Perspectiva sobre el Espacio Latente: El descubrimiento de que los componentes de ruido de baja frecuencia impulsan la diversidad proporciona una nueva comprensión de cómo los modelos de difusión interpretan el ruido inicial, sugiriendo que la "estructura" se establece temprano en el dominio de la frecuencia.

En conclusión, el artículo presenta una solución robusta, eficiente y de alta calidad al problema del colapso de modos, demostrando que "nunca es demasiado tarde" para optimizar el ruido de entrada y desbloquear el potencial de diversidad completo de los modelos de difusión entrenados.