Rekonstruktion vs. Generierung: Die Beherrschung des Optimierungs-Dilemmas in Latenten Diffusionsmodellen

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

January 2, 2025

Autoren: Jingfeng Yao, Xinggang Wang

cs.AI

Zusammenfassung

Latente Diffusionsmodelle mit Transformer-Architekturen zeichnen sich durch die Fähigkeit aus, hochwertige Bilder zu generieren. Allerdings zeigen aktuelle Studien ein Optimierungsproblem in diesem zweistufigen Design: Eine Erhöhung der Merkmalsdimension pro Token in visuellen Tokenizern verbessert die Rekonstruktionsqualität, erfordert jedoch wesentlich größere Diffusionsmodelle und mehr Trainingsiterationen, um vergleichbare Generierungsleistungen zu erreichen. Folglich akzeptieren bestehende Systeme oft suboptimale Lösungen, entweder indem sie visuelle Artefakte aufgrund von Informationsverlust innerhalb der Tokenizer erzeugen oder nicht vollständig konvergieren aufgrund hoher Rechenkosten. Wir argumentieren, dass dieses Dilemma aus der inhärenten Schwierigkeit resultiert, unbeschränkte hochdimensionale latente Räume zu erlernen. Um dies anzugehen, schlagen wir vor, den latenten Raum mit vorab trainierten Vision-Grundlagenmodellen auszurichten, während die visuellen Tokenizer trainiert werden. Unser vorgeschlagenes VA-VAE (Vision-Grundlagenmodell-ausgerichteter Variationaler Autoencoder) erweitert signifikant die Rekonstruktions-Generierungs-Grenze von latenten Diffusionsmodellen und ermöglicht eine schnellere Konvergenz von Diffusions-Transformern (DiT) in hochdimensionalen latenten Räumen. Um das volle Potenzial von VA-VAE auszuschöpfen, bauen wir eine verbesserte DiT-Basislinie mit optimierten Trainingsstrategien und Architekturentwürfen auf, genannt LightningDiT. Das integrierte System erreicht eine State-of-the-Art (SOTA) Leistung bei der Generierung von ImageNet 256x256 mit einem FID-Score von 1,35 und zeigt eine bemerkenswerte Trainingseffizienz, indem es bereits nach 64 Epochen einen FID-Score von 2,11 erreicht - was eine über 21-fache Beschleunigung der Konvergenz im Vergleich zum ursprünglichen DiT darstellt. Modelle und Codes sind verfügbar unter: 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.