Taming Sampling Perturbations with Variance Expansion Loss for Latent Diffusion Models

Resumen Técnico: Taming Sampling Perturbations with Variance Expansion Loss

1. El Problema: La Paradoja de la Robustez en el Espacio Latente

Los Modelos de Difusión Latente (LDMs) han establecido el estado del arte en la generación de imágenes de alta fidelidad. Tradicionalmente, la investigación se ha centrado en dos pilares para el tokenizador (codificador/decodificador):

1. Precisión de reconstrucción: Minimizar el error al reconstruir la imagen original.
2. Alineación semántica: Organizar el espacio latente para que coincida con representaciones semánticas preentrenadas.

La observación crítica: Los autores descubrieron una paradoja contraintuitiva. Modelos que logran una reconstrucción casi perfecta y una pérdida de difusión (diffusion loss) muy baja, a menudo generan imágenes visualmente inferiores en comparación con modelos que tienen un error de reconstrucción ligeramente mayor.

Causa raíz:

• Los tokenizadores basados en $\beta$-VAE (comúnmente usados) tienden a colapsar la varianza latente ($\sigma^2 \to 0$) para minimizar el error de reconstrucción.
• Esto crea un manifold latente excesivamente compacto (casi determinista).
• Durante el muestreo de difusión, se introducen perturbaciones estocásticas (ruido). En un espacio latente "colapsado" o demasiado estrecho, incluso pequeñas perturbaciones empujan las muestras fuera del manifold de datos válido.
• Consecuencia: El decodificador falla al interpretar estos puntos fuera del manifold, resultando en degradación visual o fallos de generación, a pesar de que el modelo de difusión haya aprendido bien la trayectoria de denoising.

2. Metodología: Pérdida de Expansión de Varianza (VE Loss)

Para abordar este problema, los autores proponen un enfoque que prioriza la robustez del espacio latente frente a las perturbaciones de muestreo, sin sacrificar la fidelidad de reconstrucción.

Análisis Teórico:

• Se demuestra matemáticamente que la pérdida de reconstrucción induce un gradiente que empuja la varianza latente hacia cero ($\partial L_{rec}/\partial \sigma \approx 2\sigma T(\mu)$).
• El término de regularización KL estándar (usado en VAEs) no es necesario en LDMs, ya que el proceso de difusión aprende la distribución a priori por sí mismo. De hecho, el KL fuerza una alineación con una gaussiana fija que puede dañar la calidad de reconstrucción.

Propuesta: Variance Expansion (VE) Loss
En lugar de usar KL, se introduce una función de pérdida diseñada para contrarrestar el colapso de la varianza:

1. Objetivo: Mantener una varianza latente ($\sigma^2$) moderada y adaptativa que absorba las perturbaciones estocásticas del muestreo.
2. Formulación:
   • Se introduce un término $L_{var}(\sigma) = \frac{1}{\sigma^2 + \delta}$ que penaliza variances demasiado pequeñas (empujando $\sigma$ a aumentar).
   • Se añade un término de regularización $L_{reg} = e^{|z|-\tau}$ para controlar la magnitud global y evitar una dispersión excesiva.
   • Pérdida Total: $L = L_{rec} + \lambda_1 L_{var} + \lambda_2 L_{reg}$.
3. Mecanismo de Equilibrio: Existe una relación adversarial natural. La pérdida de reconstrucción quiere $\sigma$ pequeño (precisión), mientras que la VE Loss quiere $\sigma$ mayor (robustez). El modelo aprende una varianza adaptativa óptima: lo suficientemente grande para resistir el ruido de difusión, pero lo suficientemente pequeña para mantener la fidelidad de la imagen.

3. Contribuciones Clave

1. Identificación de un factor crítico: Se revela que la falta de robustez del espacio latente frente a las perturbaciones de muestreo es un factor determinante (y a menudo ignorado) en la calidad de generación de LDMs.
2. Nueva función de pérdida: Propuesta de la VE Loss, una solución simple pero efectiva que elimina la necesidad de regularización KL y fomenta un espacio latente robusto y adaptativo.
3. Validación empírica y teórica: Demostración de que este enfoque mejora consistentemente la calidad de generación en diversas arquitecturas (DiT, LightningDiT, etc.) y tokenizadores, manteniendo una reconstrucción competitiva.

4. Resultados Experimentales

Los experimentos se realizaron en el conjunto de datos ImageNet (256x256) utilizando varias arquitecturas de base (DiT, SiT, LightningDiT).

• Comparación de Tokenizadores:
  • Al reemplazar la pérdida KL estándar con VE Loss en tokenizadores Vanilla y VA-VAE, se observó una mejora consistente en la calidad de generación (medida por FID).
  • En el caso de VA-VAE, el método propuesto logró un FID de 12.89 con solo 16 épocas de entrenamiento del tokenizador, superando al modelo base entrenado durante 50 épocas (FID 15.82).
• Comparación con el Estado del Arte (SOTA):
  • El modelo propuesto ("Ours") alcanzó un gFID de 1.18 en 530 épocas, superando a métodos SOTA como VA-VAE (1.35), REG (1.36) y RAE (1.41), que requieren 800 épocas o más.
  • También superó a modelos autoregresivos avanzados como LlamaGen y VAR en términos de FID.
• Métricas de Reconstrucción:
  • Aunque hubo un ligero compromiso teórico en la reconstrucción pura (PSNR ligeramente menor en configuraciones estrictas), las métricas generales (rFID, LPIPS, SSIM) se mantuvieron altamente competitivas, demostrando que la ganancia en generación no se logra a costa de una reconstrucción inutilizable.

5. Significado e Impacto

Este trabajo cambia la perspectiva sobre el diseño de espacios latentes para modelos generativos:

• Cambio de Paradigma: Sugiere que la "compacidad" extrema (baja varianza) no es siempre deseable. Un espacio latente debe tener cierta "flexibilidad" o "dispersión" controlada para ser robusto frente a la naturaleza estocástica de la difusión.
• Simplicidad Efectiva: La solución no requiere arquitecturas complejas ni grandes cambios en el modelo de difusión, sino una modificación en la función de pérdida del tokenizador.
• Estabilidad: Ofrece una vía para estabilizar el entrenamiento de LDMs de larga duración, evitando los fallos de muestreo que ocurren cuando el espacio latente es demasiado rígido.

En conclusión, la VE Loss proporciona un mecanismo principiado para equilibrar la fidelidad de reconstrucción y la robustez al ruido, llenando un vacío crítico en la investigación de modelos de difusión latente y estableciendo un nuevo estándar para la generación de imágenes de alta calidad.