Reconstructie versus Generatie: Het Temmen van het Optimalisatie Dilemma in Latente Diffusie Modellen

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

January 2, 2025

Auteurs: Jingfeng Yao, Xinggang Wang

cs.AI

Samenvatting

Latente diffusiemodellen met Transformer-architecturen blinken uit in het genereren van hoogwaardige afbeeldingen. Recent onderzoek onthult echter een optimalisatiedilemma in dit tweestapsontwerp: het verhogen van de kenmerkdimensie per token in visuele tokenizers verbetert de reconstructiekwaliteit, maar vereist aanzienlijk grotere diffusiemodellen en meer trainingsiteraties om vergelijkbare generatieprestaties te bereiken. Als gevolg hiervan kiezen bestaande systemen vaak voor suboptimale oplossingen, waarbij ze ofwel visuele artefacten produceren door informatieverlies binnen tokenizers, ofwel niet volledig convergeren vanwege hoge berekeningskosten. Wij stellen dat dit dilemma voortkomt uit de inherente moeilijkheid van het leren van onbeperkte hoogdimensionale latente ruimtes. Om dit aan te pakken, stellen we voor om de latente ruimte af te stemmen op vooraf getrainde vision foundation-modellen bij het trainen van de visuele tokenizers. Ons voorgestelde VA-VAE (Vision foundation model Aligned Variational AutoEncoder) breidt aanzienlijk de reconstructie-generatiegrens van latente diffusiemodellen uit, waardoor een snellere convergentie van Diffusion Transformers (DiT) in hoogdimensionale latente ruimtes mogelijk is. Om het volledige potentieel van VA-VAE te benutten, bouwen we een verbeterde DiT-baseline met verbeterde trainingsstrategieën en architectuurontwerpen, genaamd LightningDiT. Het geïntegreerde systeem behaalt state-of-the-art (SOTA) prestaties op ImageNet 256x256 generatie met een FID-score van 1.35, terwijl het opmerkelijke trainingsefficiëntie aantoont door een FID-score van 2.11 te bereiken in slechts 64 epochs - wat een meer dan 21 keer snellere convergentie vertegenwoordigt in vergelijking met de oorspronkelijke DiT. Modellen en codes zijn beschikbaar op: 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.