Reconstrucción vs. Generación: Dominando el Dilema de Optimización en Modelos de Difusión Latente

Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models

January 2, 2025

Autores: Jingfeng Yao, Xinggang Wang

cs.AI

Resumen

Los modelos de difusión latente con arquitecturas Transformer destacan en la generación de imágenes de alta fidelidad. Sin embargo, estudios recientes revelan un dilema de optimización en este diseño de dos etapas: mientras que aumentar la dimensión de características por token en los tokenizadores visuales mejora la calidad de reconstrucción, requiere modelos de difusión sustancialmente más grandes y más iteraciones de entrenamiento para lograr un rendimiento de generación comparable. En consecuencia, los sistemas existentes a menudo se conforman con soluciones subóptimas, ya sea produciendo artefactos visuales debido a la pérdida de información dentro de los tokenizadores o fallando en converger completamente debido a los costos computacionales elevados. Sostenemos que este dilema surge de la dificultad inherente en aprender espacios latentes de alta dimensionalidad sin restricciones. Para abordar esto, proponemos alinear el espacio latente con modelos de visión pre-entrenados al entrenar los tokenizadores visuales. Nuestro propuesto VA-VAE (AutoCodificador Variacional Alineado con Modelo de Fundación de Visión) expande significativamente la frontera de reconstrucción-generación de los modelos de difusión latente, permitiendo una convergencia más rápida de los Transformadores de Difusión (DiT) en espacios latentes de alta dimensionalidad. Para explotar todo el potencial de VA-VAE, construimos una línea de base DiT mejorada con estrategias de entrenamiento y diseños de arquitectura mejorados, denominada LightningDiT. El sistema integrado logra un rendimiento de estado del arte (SOTA) en generación de ImageNet 256x256 con una puntuación FID de 1.35, demostrando una eficiencia de entrenamiento notable al alcanzar una puntuación FID de 2.11 en solo 64 épocas, lo que representa una aceleración de la velocidad de convergencia de más de 21 veces en comparación con el DiT original. Los modelos y códigos están disponibles en: https://github.com/hustvl/LightningDiT.

English

Latent diffusion models with Transformer architectures excel at generating high-fidelity images. However, recent studies reveal an optimization dilemma in this two-stage design: while increasing the per-token feature dimension in visual tokenizers improves reconstruction quality, it requires substantially larger diffusion models and more training iterations to achieve comparable generation performance. Consequently, existing systems often settle for sub-optimal solutions, either producing visual artifacts due to information loss within tokenizers or failing to converge fully due to expensive computation costs. We argue that this dilemma stems from the inherent difficulty in learning unconstrained high-dimensional latent spaces. To address this, we propose aligning the latent space with pre-trained vision foundation models when training the visual tokenizers. Our proposed VA-VAE (Vision foundation model Aligned Variational AutoEncoder) significantly expands the reconstruction-generation frontier of latent diffusion models, enabling faster convergence of Diffusion Transformers (DiT) in high-dimensional latent spaces. To exploit the full potential of VA-VAE, we build an enhanced DiT baseline with improved training strategies and architecture designs, termed LightningDiT. The integrated system achieves state-of-the-art (SOTA) performance on ImageNet 256x256 generation with an FID score of 1.35 while demonstrating remarkable training efficiency by reaching an FID score of 2.11 in just 64 epochs--representing an over 21 times convergence speedup compared to the original DiT. Models and codes are available at: https://github.com/hustvl/LightningDiT.