Resumen Técnico: SSDD: Decodificador de Difusión de Paso Único para una Tokenización de Imágenes Eficiente

Planteamiento del Problema

Los modelos generativos de imágenes de última generación dependen de tokenizadores para comprimir señales visuales de alta dimensión en representaciones latentes compactas. Si bien los tokenizadores estándar actuales, como los Autoencoders Variacionales regularizados por KL (KL-VAE), son efectivos, enfrentan un compromiso fundamental entre la distorsión (similitud a nivel de píxel) y el desplazamiento de distribución (realismo perceptual y diversidad). Los decodificadores deterministas minimizan la distorsión pero a menudo fallan al modelar la compleja distribución de imágenes realistas, lo que resulta en salidas borrosas. Por el contrario, los decodificadores generativos pueden modelar esta distribución, pero típicamente requieren un muestreo iterativo (inferencia lenta) o dependen de pérdidas adversarias (GAN), las cuales son difíciles de escalar y estabilizar.

Los decodificadores basados en difusión existentes ofrecen un enfoque fundamentado para modelar la distribución condicional de las imágenes dados los latentes, pero sufren dos limitaciones principales:

1. Inferencia Lenta: Requieren múltiples pasos de eliminación de ruido, lo que crea un cuello de botella para los modelos generativos descendentes.
2. Inestabilidad de Entrenamiento: Lograr resultados de vanguardia a menudo requiere pérdidas GAN, que son difíciles de escalar y pueden provocar ejecuciones de entrenamiento divergentes.

Metodología

Los autores presentan SSDD (Single-Step Diffusion Decoder), una arquitectura novedosa diseñada para superar estas limitaciones combinando una arquitectura híbrida escalable, un entrenamiento libre de GAN y una estrategia de destilación de paso único.

1. Arquitectura Híbrida Escalable

En lugar de depender únicamente de U-Nets convolucionales (comunes en la difusión en el espacio de píxeles) o Transformers puros (comunes en el modelado latente), SSDD utiliza una arquitectura basada en U-ViT.

• Estructura: Emplea un esqueleto U-Net con cuatro niveles de bloques ResNet convolucionales, pero reemplaza el bloque de atención central con un Transformer que opera sobre tokens que representan parches de píxeles de $8 \times 8$.
• Condicionamiento: El modelo se condiciona a la representación latente $z$ mediante el sobremuestreo (upsampling) de la cuadrícula latente y la concatenación con la entrada con ruido. Además, mejora el condicionamiento mediante Normalización de Grupo Adaptativa (AdaGN) y Normalización de Capa Adaptativa (AdaLN) dentro de los bloques tanto de ResNet como de Transformer.
• Escalabilidad: La arquitectura escala eficientemente desde 13.4M (SSDD-S) hasta 153.8M (SSDD-XL) de parámetros aumentando los anchos de canal y la profundidad del transformer.

2. Esquema de Entrenamiento Libre de GAN

SSDD elimina la necesidad de entrenamiento adversario combinando tres funciones de pérdida específicas:

• Pérdida de Flow Matching: El objetivo principal utiliza el emparejamiento de flujo de transporte óptimo (optimal transport flow matching) para predecir el campo de velocidad ($v = x - \epsilon$), entrenado con una pérdida $L_2$.
• Pérdida Perceptual (LPIPS): Añadida para guiar la generación hacia áreas perceptualmente correctas, alineando características de alto nivel.
• Pérdida REPA: Una pérdida de regularización de características utilizando características de DINOv2-B para estabilizar las representaciones del transformer y acelerar el entrenamiento.
• Entrenamiento de Múltiples Etapas: El modelo se entrena en etapas: primero en recortes de $128 \times 128$ con aumentación de redimensionamiento aleatorio, seguido de un ajuste fino (fine-tuning) a las resoluciones objetivo ($256 \times 256$).

3. Destilación de Paso Único

Para abordar el cuello de botella de la velocidad de inferencia, SSDD emplea una estrategia de destilación:

• Configuración Maestro-Estudiante: Un decodificador maestro de múltiples pasos (entrenado con 8 pasos) genera reconstrucciones de alta calidad y diversas. Un decodificador estudiante es luego ajustado para reproducir estas salidas en un solo paso.
• Preservación de la Pérdida: A diferencia de la destilación estándar que podría depender únicamente de la regresión $L_2$, la destilación de SSDD preserva el objetivo de entrenamiento completo. El estudiante minimiza las pérdidas de flow-matching y LPIPS calculadas contra las salidas de múltiples pasos del maestro.
• Resultado: Esto permite que el modelo herede las propiedades perceptuales y generativas del maestro de múltiples pasos mientras opera con el rendimiento (throughput) de un modelo de paso único.

4. Codificadores Compartidos

Los autores demuestran que los decodificadores SSDD pueden entrenarse con un codificador compartido (específicamente el codificador regularizado por KL estándar de Rombach et al., 2022). Esto permite que SSDD sirva como un reemplazo directo para los decodificadores VAE existentes sin requerir el reentrenamiento del codificador o la creación de un nuevo espacio latente para cada tamaño de decodificador.

Contribuciones Clave

1. Primer Decodificador de Difusión de Paso Único Libre de GAN: SSDD es la primera arquitectura en lograr una tokenización continua de vanguardia sin pérdidas adversarias, asegurando la estabilidad del entrenamiento a través de diversas escalas de modelos.
2. Arquitectura Híbrida: La integración de sesgos inductivos convolucionales con capacidades generativas de transformer en un decodificador de espacio de píxeles produce propiedades de escalado superiores en comparación con los U-Nets puros.
3. Destilación Eficiente: El método de destilación propuesto comprime con éxito el comportamiento de la difusión de múltiples pasos en un solo paso con una degradación mínima de la calidad, mejorando significativamente el rendimiento.
4. Compatibilidad de Espacio Latente Compartido: La capacidad de utilizar codificadores pre-entrenados estándar hace que SSDD sea aplicable de inmediato a los flujos de trabajo generativos existentes.

Resultados Experimentales

Evaluado en ImageNet-1k (256$\times$256) y COCO, SSDD demuestra mejoras significativas sobre KL-VAE y otros decodificadores de difusión:

• Calidad de Reconstrucción (configuración f8c4):
  • SSDD-S (13.4M params) mejora el FID de reconstrucción (rFID) de 0.87 (KL-VAE) a 0.46.
  • SSDD-M (48M params, tamaño similar a KL-VAE) logra un rFID de 0.39.
  • SSDD supera a todos los baselines en métricas perceptuales (LPIPS, DreamSim) mientras mantiene métricas de bajo nivel competitivas (PSR, SSIM), aunque prioriza la calidad perceptual sobre la distorsión de píxel perfecto.
• Rendimiento (Throughput):
  • SSDD-S logra un rendimiento 1.4$\times$ mayor que KL-VAE ofreciendo una mejor calidad.
  • Comparado con los decodificadores de difusión de múltiples pasos (ej. $\epsilon$-VAE), SSDD ofrece hasta 10$\times$ más velocidad de decodificación en niveles de rFID similares.
• Generación Descendente:
  • Al usarse para entrenar Diffusion Transformers (DiT), SSDD mejora el FID de generación (gFID) en comparación con KL-VAE.
  • Usando una mayor relación de compresión (f16c4), SSDD permite un muestreo 3.8$\times$ más rápido para el modelo DiT preservando la calidad de generación.
• Generalización: Los modelos SSDD generalizan bien a COCO y resoluciones más altas (hasta 1024$\times$1024) sin necesidad de reentrenamiento desde cero, y pueden decodificar latentes de codificadores congelados como DiTo y DC-AE.

Significado

El artículo afirma que SSDD establece un nuevo estándar para la tokenización de imágenes al resolver la tensión entre velocidad, estabilidad y fidelidad generativa. Al eliminar la dependencia de las GAN y el muestreo iterativo, SSDD proporciona un reemplazo directo y escalable para KL-VAE que permite el entrenamiento de modelos generativos de mayor calidad y más rápidos. El trabajo demuestra que la tokenización basada en difusión puede ser tanto eficiente como estable, allanando el camino para futuros modelos generativos a gran escala sin la sobrecarga computacional de la inferencia de múltiples pasos o la inestabilidad del entrenamiento adversario.