Síntesis de imágenes de alta resolución con Latent Diffusion Models

Guía de lectura

Problema y motivación

Los modelos probabilísticos de difusión generan imágenes de alta fidelidad que cubren adecuadamente los modos de la distribución invirtiendo un proceso de ruido hacia adelante mediante una secuencia de autoencoders de eliminación de ruido. Sin embargo, su entrenamiento y evaluación directa en el espacio de píxeles RGB de alta dimensión es computacionalmente exigente. La optimización a menudo requiere cientos de días de GPU (por ejemplo, entre 150 y 1000 días de V100), y el muestreo secuencial de múltiples pasos puede tardar días en ejecuciones de evaluación estándar. Debido a que el modelado generativo basado en verosimilitud dedica una capacidad excesiva a detalles de alta frecuencia imperceptibles en la imagen antes de aprender estructuras semánticas, las evaluaciones directas en el espacio de píxeles provocan evaluaciones de gradiente y función sumamente ineficientes.

Ideas clave

• Descomposición en dos etapas: Separar la compresión perceptual (eliminar detalles imperceptibles) del modelado generativo semántico (aprender la estructura y la composición de la imagen).
• Modelos de Difusión Latente (LDMs): Entrenar modelos de difusión de eliminación de ruido dentro del espacio latente de un autoencoder entrenado con pérdidas perceptuales y adversarias.
• Compresión espacial moderada: A diferencia de los modelos autorregresivos que requieren una reducción espacial agresiva (por ejemplo, $f=16$ o $f=32$), los LDMs mantienen el sesgo inductivo espacial con una UNet convolucional 2D, lo que permite factores de submuestreo moderados como $f=4$ o $f=8$ que preservan los detalles finos.
• Condicionamiento universal mediante cross-attention: Generalizar las UNets para admitir entradas multimodales (descripciones de texto, cuadros delimitadores, mapas semánticos) empleando codificadores específicos de dominio mapeados en capas intermedias de la red mediante atención cruzada.

Método

1. Compresión perceptual de imágenes: Un codificador $\mathcal{E}$ proyecta una imagen $x\in {\mathbb{R}}^{H\times W\times 3}$ a una representación latente de menor dimensión $z=\mathcal{E}(x)\in {\mathbb{R}}^{h\times w\times c}$, donde el factor de submuestreo es $f=H/h=W/w=2^m$. Un decodificador reconstruye $\tilde{x}=\mathcal{D}(z)$. El modelo se regulariza mediante una penalización débil de divergencia KL hacia $\mathcal{N}(0,\mathbb{I})$ (KL-reg.) o mediante cuantización vectorial (VQ-reg.) absorbida en el decodificador.

2. Modelado latente generativo: Se entrena una UNet condicionada por tiempo ${ϵ}_{\theta }(z_t,t)$ sobre los latentes usando el objetivo reponderado de eliminación de ruido: $$L_{LDM}:={\mathbb{E}}_{\mathcal{E}(x),ϵ\sim \mathcal{N}(0,1),t}\left[\Vert ϵ-{ϵ}_{\theta }(z_t,t){\Vert }_2^2\right]$$

3. Condicionamiento mediante Cross-Attention: Para una entrada de condicionamiento $y$ (por ejemplo, texto, clases, distribuciones espaciales), se introduce un codificador de dominio ${\tau }_{\theta }(y)\in {\mathbb{R}}^{M\times d_{\tau }}$. Las características intermedias de la UNet ${\phi }_i(z_t)\in {\mathbb{R}}^{N\times d_{ϵ}^i}$ interactúan con el condicionamiento a través de cross-attention: $$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d}}\right)\cdot V$$ donde $Q=W_Q^{(i)}{\phi }_i(z_t)$, $K=W_K^{(i)}{\tau }_{\theta }(y)$ y $V=W_V^{(i)}{\tau }_{\theta }(y)$. El objetivo condicional es: $$L_{LDM}:={\mathbb{E}}_{\mathcal{E}(x),y,ϵ\sim \mathcal{N}(0,1),t}\left[\Vert ϵ-{ϵ}_{\theta }(z_t,t,{\tau }_{\theta }(y)){\Vert }_2^2\right]$$

Resultados

• Compromiso en la compresión perceptual: Los factores de submuestreo $f=4$ y $f=8$ logran el mejor equilibrio entre precisión de reconstrucción, calidad de síntesis y tasa de procesamiento. En ImageNet, una brecha de 38 puntos de FID separa al LDM-1 basado en píxeles del LDM-8 tras 2M pasos de entrenamiento con presupuestos computacionales idénticos.
• Generación incondicional: Se logró un FID de vanguardia de $5.11$ en CelebA-HQ ($256\times 256$), $4.98$ en FFHQ, $4.02$ en LSUN-Churches y $2.95$ en LSUN-Bedrooms.
• ImageNet condicional por clase: LDM-4 con classifier-free guidance alcanzó un FID de $3.60$ y un Inception Score de $247.67$, superando a ADM ($4.59$) y utilizando menos parámetros (400M frente a 608M) y considerablemente menos cómputo de entrenamiento (271 frente a 962 días de V100).
• Texto a imagen: Un modelo de 1.45B parámetros (LDM-KL-8-G) alcanzó un FID de $12.63$ en MS-COCO con 250 pasos de DDIM.
• Superresolución e inpainting: LDM-4 estableció un nuevo estado del arte en inpainting con un FID de $1.50$ en Places $512\times 512$ (y $9.39$ en imágenes con máscaras difíciles del 40–50%), operando a un rendimiento más de $2.7\times$ mayor que la difusión en el espacio de píxeles.

Limitaciones y preguntas abiertas

• Latencia de muestreo: Aunque es significativamente más rápido que los modelos de difusión basados en píxeles, el muestreo secuencial de pasos múltiples sigue siendo más lento que los generadores GAN de una sola pasada.
• Cuello de botella en la reconstrucción: Incluso con un submuestreo bajo ($f=4$), la compresión puede perder detalles de píxeles de alta frecuencia muy finos, lo que resulta ligeramente subóptimo para tareas críticas donde cada píxel es esencial.
• Impacto social y seguridad: Reducir las barreras de cómputo democratiza las herramientas creativas, pero también facilita la creación de deepfakes y el mal uso, además de heredar potencialmente los sesgos de conjuntos de datos de entrenamiento no curados a escala web.