通过对齐视觉基础模型,提升扩散模型的重建与生成效率。
Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models
- 用预训练视觉模型对齐潜在空间,缓解高维特征训练难题。
- 在ImageNet上实现1.35的FID,64轮内达2.11,提速超21倍。
- 适合追求高效高质图像生成的研究者与开发者。
基于Transformer架构的潜在扩散模型在生成高质量图像方面表现优异。然而,近期研究揭示其两阶段设计存在优化困境:增加视觉编码器的每标记特征维度虽能提升重建质量,却需更大扩散模型和更多训练迭代才能达到相当的生成性能。因此,现有系统常妥协于次优方案,或因编码器信息损失导致视觉伪影,或因计算成本过高而无法充分收敛。我们认为此困境源于学习无约束高维潜在空间的固有难度。为此,我们提出在训练视觉编码器时对齐预训练视觉基础模型,构建VA-VAE(视觉基础模型对齐变分自编码器)。该方法显著拓展了潜在扩散模型的重建-生成边界,使扩散Transformer(DiT)在高维潜在空间中实现更快收敛。为充分发挥VA-VAE潜力,我们设计改进版DiT基线LightningDiT,采用更优训练策略与结构。集成系统在ImageNet 256x256生成任务上达到1.35的FID,仅用64轮即达2.11,相较原始DiT实现超过21倍的收敛加速。代码与模型已开源。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。