Ricostruzione vs. Generazione: Addomesticare il Dilemma dell' Ottimizzazione nei Modelli di Diffusione Latente

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

January 2, 2025

Autori: Jingfeng Yao, Xinggang Wang

cs.AI

Abstract

I modelli di diffusione latente con architetture Transformer eccellono nella generazione di immagini ad alta fedeltà. Tuttavia, recenti studi rivelano un dilemma di ottimizzazione in questo design a due fasi: mentre aumentare la dimensione delle caratteristiche per token nei tokenizzatori visivi migliora la qualità della ricostruzione, richiede modelli di diffusione notevolmente più grandi e più iterazioni di addestramento per ottenere prestazioni di generazione comparabili. Di conseguenza, i sistemi esistenti spesso si accontentano di soluzioni subottimali, producendo artefatti visivi a causa della perdita di informazioni all'interno dei tokenizzatori o non riuscendo a convergere completamente a causa dei costi computazionali elevati. Sosteniamo che questo dilemma derivi dalla difficoltà intrinseca nell'apprendere spazi latenti ad alta dimensionalità non vincolati. Per affrontare ciò, proponiamo di allineare lo spazio latente con modelli di fondazione visiva preaddestrati durante l'addestramento dei tokenizzatori visivi. Il nostro VA-VAE proposto (Variational AutoEncoder Allineato al Modello di Fondazione Visiva) espande significativamente il confine di ricostruzione-generazione dei modelli di diffusione latente, consentendo una convergenza più rapida dei Transformer di Diffusione (DiT) in spazi latenti ad alta dimensionalità. Per sfruttare appieno il potenziale di VA-VAE, costruiamo un basale DiT potenziato con strategie di addestramento migliorate e design di architettura, denominato LightningDiT. Il sistema integrato raggiunge prestazioni all'avanguardia nella generazione di ImageNet 256x256 con un punteggio FID di 1.35, dimostrando un'efficienza di addestramento notevole raggiungendo un punteggio FID di 2.11 in soli 64 epoche - rappresentando un'accelerazione della convergenza di oltre 21 volte rispetto al DiT originale. Modelli e codici sono disponibili su: 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.