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

• URL: http://arxiv.org/abs/2501.01423v2
• Date: Mon, 06 Jan 2025 15:28:11 GMT
• Title: Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models
• Authors: Jingfeng Yao, Xinggang Wang,
• Abstract summary: We propose aligning the latent space with pre-trained vision foundation models when training the visual tokenizers. Our proposed VA-VAE significantly expands the reconstruction-generation frontier of latent diffusion models. We build an enhanced DiT baseline with improved training strategies and architecture designs, termed LightningDiT.
• Score: 34.15905637499148
• License:
• Abstract: 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.

Related papers

• Masked Autoencoders Are Effective Tokenizers for Diffusion Models [56.08109308294133]
  MAETok is an autoencoder that learns semantically rich latent space while maintaining reconstruction fidelity. MaETok achieves significant practical improvements, enabling a gFID of 1.69 with 76x faster training and 31x higher inference throughput for 512x512 generation.
  arXiv  Detail & Related papers  (2025-02-05T18:42:04Z)
• Exploring Representation-Aligned Latent Space for Better Generation [86.45670422239317]
  We introduce ReaLS, which integrates semantic priors to improve generation performance. We show that fundamental DiT and SiT trained on ReaLS can achieve a 15% improvement in FID metric. The enhanced semantic latent space enables more perceptual downstream tasks, such as segmentation and depth estimation.
  arXiv  Detail & Related papers  (2025-02-01T07:42:12Z)
• Visual Autoregressive Modeling for Image Super-Resolution [14.935662351654601]
  We propose a novel visual autoregressive modeling for ISR framework with the form of next-scale prediction. We collect large-scale data and design a training process to obtain robust generative priors.
  arXiv  Detail & Related papers  (2025-01-31T09:53:47Z)
• ARLON: Boosting Diffusion Transformers with Autoregressive Models for Long Video Generation [83.62931466231898]
  This paper presents ARLON, a framework that boosts diffusion Transformers with autoregressive models for long video generation. A latent Vector Quantized Variational Autoencoder (VQ-VAE) compresses the input latent space of the DiT model into compact visual tokens. An adaptive norm-based semantic injection module integrates the coarse discrete visual units from the AR model into the DiT model.
  arXiv  Detail & Related papers  (2024-10-27T16:28:28Z)
• LinFusion: 1 GPU, 1 Minute, 16K Image [71.44735417472043]
  We introduce a low-rank approximation of a wide spectrum of popular linear token mixers. We find that the distilled model, termed LinFusion, achieves performance on par with or superior to the original SD. Experiments on SD-v1.5, SD-v2.1, and SD-XL demonstrate that LinFusion enables satisfactory and efficient zero-shot cross-resolution generation.
  arXiv  Detail & Related papers  (2024-09-03T17:54:39Z)
• Alleviating Distortion in Image Generation via Multi-Resolution Diffusion Models and Time-Dependent Layer Normalization [26.926712014346432]
  This paper presents innovative enhancements to diffusion models by integrating a novel multi-resolution network and time-dependent layer normalization. Our method's efficacy is demonstrated on the class-conditional ImageNet generation benchmark, setting new state-of-the-art FID scores of 1.70 on ImageNet 256 x 256 and 2.89 on ImageNet 512 x 512.
  arXiv  Detail & Related papers  (2024-06-13T17:59:58Z)
• Binarized Diffusion Model for Image Super-Resolution [61.963833405167875]
  Binarization, an ultra-compression algorithm, offers the potential for effectively accelerating advanced diffusion models (DMs) Existing binarization methods result in significant performance degradation. We introduce a novel binarized diffusion model, BI-DiffSR, for image SR.
  arXiv  Detail & Related papers  (2024-06-09T10:30:25Z)
• StraIT: Non-autoregressive Generation with Stratified Image Transformer [63.158996766036736]
  Stratified Image Transformer(StraIT) is a pure non-autoregressive(NAR) generative model. Our experiments demonstrate that StraIT significantly improves NAR generation and out-performs existing DMs and AR methods.
  arXiv  Detail & Related papers  (2023-03-01T18:59:33Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.