Reconstruction vs. Génération : Apprivoiser le dilemme de l'optimisation dans les modèles de diffusion latente

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

January 2, 2025

Auteurs: Jingfeng Yao, Xinggang Wang

cs.AI

Résumé

Les modèles de diffusion latente avec architectures Transformer excellent dans la génération d'images haute fidélité. Cependant, des études récentes révèlent un dilemme d'optimisation dans cette conception en deux étapes : augmenter la dimension des caractéristiques par jeton dans les tokenizers visuels améliore la qualité de la reconstruction, mais nécessite des modèles de diffusion substantiellement plus grands et davantage d'itérations d'entraînement pour atteindre des performances de génération comparables. Par conséquent, les systèmes existants optent souvent pour des solutions sous-optimales, produisant soit des artefacts visuels dus à une perte d'information dans les tokenizers, soit échouant à converger complètement en raison des coûts de calcul élevés. Nous soutenons que ce dilemme découle de la difficulté inhérente à apprendre dans des espaces latents de haute dimension non contraints. Pour y remédier, nous proposons d'aligner l'espace latent avec des modèles de fondation visuelle pré-entraînés lors de l'entraînement des tokenizers visuels. Notre VA-VAE proposé (AutoEncodeur Variationnel Aligné avec le modèle de fondation visuelle) élargit significativement la frontière de reconstruction-génération des modèles de diffusion latente, permettant une convergence plus rapide des Transformateurs de Diffusion (DiT) dans des espaces latents de haute dimension. Pour exploiter pleinement le potentiel de VA-VAE, nous construisons une ligne de base DiT améliorée avec des stratégies d'entraînement et des conceptions architecturales améliorées, appelée LightningDiT. Le système intégré atteint des performances de pointe sur la génération ImageNet 256x256 avec un score FID de 1,35, tout en démontrant une efficacité d'entraînement remarquable en atteignant un score FID de 2,11 en seulement 64 époques - représentant une accélération de la vitesse de convergence de plus de 21 fois par rapport au DiT original. Les modèles et les codes sont disponibles sur : 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.