Reconstrução vs. Geração: Domando o Dilema de Otimização em Modelos de Difusão Latente

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

January 2, 2025

Autores: Jingfeng Yao, Xinggang Wang

cs.AI

Resumo

Os modelos de difusão latente com arquiteturas Transformer destacam-se na geração de imagens de alta fidelidade. No entanto, estudos recentes revelam um dilema de otimização nesse design de dois estágios: aumentar a dimensão da característica por token nos tokenizadores visuais melhora a qualidade da reconstrução, mas exige modelos de difusão substancialmente maiores e mais iterações de treinamento para alcançar desempenho de geração comparável. Consequentemente, os sistemas existentes frequentemente se contentam com soluções subótimas, seja produzindo artefatos visuais devido à perda de informação nos tokenizadores ou falhando em convergir completamente devido aos custos computacionais elevados. Argumentamos que esse dilema decorre da dificuldade inerente em aprender espaços latentes de alta dimensão não restritos. Para lidar com isso, propomos alinhar o espaço latente com modelos de fundação visual pré-treinados ao treinar os tokenizadores visuais. Nosso VA-VAE proposto (Variational AutoEncoder Alinhado com Modelo de Fundação Visual) expande significativamente a fronteira de reconstrução-geração de modelos de difusão latente, permitindo uma convergência mais rápida dos Transformadores de Difusão (DiT) em espaços latentes de alta dimensão. Para explorar todo o potencial do VA-VAE, construímos uma linha de base DiT aprimorada com estratégias de treinamento e designs de arquitetura melhorados, denominada LightningDiT. O sistema integrado alcança desempenho de ponta (SOTA) na geração do ImageNet 256x256 com uma pontuação FID de 1.35, demonstrando notável eficiência de treinamento ao atingir uma pontuação FID de 2.11 em apenas 64 épocas - representando um aumento de mais de 21 vezes na velocidade de convergência em comparação com o DiT original. Modelos e códigos estão disponíveis em: 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.