Source-linked AI summary
Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models
Jingfeng Yao, Bin Yang, Xinggang Wang
TL;DR
Latent diffusion models face a trade-off in which higher-dimensional tokenizer latents improve reconstruction but impair generation and increase training demands. The paper aligns tokenizer latents with pre-trained vision foundation models through VA-VAE and combines it with LightningDiT, an enhanced DiT framework. The integrated system achieves FID 1.35 on ImageNet-256 and reaches FID 2.11 in 64 epochs, reporting over 21× faster convergence than the original DiT.
Problem
Higher-dimensional visual tokenizers improve reconstruction but degrade generation and require substantially greater training costs, while continuous VAE optimization lacks comprehensive analysis and effective solutions.
Method
VA-VAE aligns continuous tokenizer latents with pre-trained vision foundation models using VF Loss, while LightningDiT adds improved diffusion training strategies and transformer designs.
Results
The integrated system achieves SOTA FID 1.35 on ImageNet 256 generation and reaches FID 2.11 in 64 epochs, an over 21× convergence speedup versus the original DiT.
Takeaways & Limitations
Vision foundation model alignment expands the reconstruction-generation frontier and accelerates high-dimensional latent diffusion convergence without additional training costs.
Takeaways & Limitations
Lower-dimensional latent spaces can learn more reasonable distributions without additional supervisory signals, limiting the stated need for foundation-model guidance in that setting.
Abstract
from arXiv · showhide
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.
1. Introduction
Latent diffusion models face a reconstruction-generation trade-off: higher-dimensional visual tokens preserve more detail but make generation harder. The paper addresses this by aligning high-dimensional latents with vision foundation models and pairing the tokenizer with an enhanced DiT framework.
- Optimization dilemma: Higher-dimensional visual tokens improve detail reconstruction but significantly degrade generation performance under fixed diffusion-training compute.The tokenizer uses f for downsampling rate and d for feature dimension; results are evaluated on ImageNet 256×256.
- Optimization dilemma: Scaling model capacity can recover generation quality for higher-dimensional tokenizers, but requires substantially more training compute.The introduction cites Stable Diffusion 3 as an example of this strategy.
- Proposed explanation: The paper attributes the dilemma to unconstrained high-dimensional latent spaces whose representations become less spread out and harder to learn from scratch.Visualization shows more concentrated high-intensity areas as tokenizer feature dimensions increase.
- VA-VAE: VA-VAE aligns tokenizer latents with pre-trained vision foundation models, improving high-dimensional generation while preserving reconstruction capabilities.Its VF Loss is designed to regularize global and local feature structures while retaining latent capacity through a similarity margin.
- LightningDiT: LightningDiT enhances DiT with advanced diffusion training strategies and transformer architecture improvements to exploit VA-VAE’s potential.The integrated system is evaluated as a latent diffusion model combining VA-VAE with DiT.
- Results: The integrated system reaches FID 2.11 in 64 epochs, over 21× faster than the original DiT, and achieves SOTA FID 1.35 on ImageNet-256 generation.These milestones are reported alongside over 2.5× faster DiT training with high-dimensional tokenizers.
2. Related Work
Prior work establishes visual tokenizers and faster DiT variants, but continuous VAE systems still lack comprehensive analysis and effective solutions for their reconstruction-generation trade-off. This paper targets the tokenizer latent space rather than modifying diffusion training with auxiliary losses.
- Visual tokenizers: VAEs create continuous visual representations for compressing and reconstructing visual signals, while VQVAEs use codebooks for discrete representation learning.VQGAN extends the discrete approach with adversarial training.
- Continuous VAE systems: Continuous VAE latent diffusion systems improve reconstruction through higher tokenizer dimensionality but often degrade generation performance and require higher training costs.The related-work discussion identifies this as a parallel optimization challenge to discrete representation systems.
- Research gap: The literature lacks comprehensive analysis and effective solutions for optimizing continuous VAEs under this reconstruction-generation trade-off.The paper addresses this gap by structuring continuous VAE latents with vision foundation model representations.
- Diffusion Transformers: DiT is a scalable latent-diffusion implementation, but prior methods such as SiT, MDT, and MaskDiT target its slow convergence through diffusion or masking strategies.These methods provide context for the paper’s distinct tokenizer-focused approach.
- Positioning: The paper identifies the tokenizer as a major efficient-training bottleneck and accelerates convergence without modifying the diffusion model.Unlike auxiliary-loss approaches, its method incurs no additional diffusion-training parameters or computational costs.
- Relation to REPA: Although REPA also uses vision foundation models, it constrains DiT, whereas this work regulates the tokenizer to address both reconstruction and generation.The related-work comparison distinguishes their motivations and intervention points.
3. Align VAE with Vision Foundation Models
VA-VAE aligns a continuous visual tokenizer’s latent space with frozen vision foundation-model features through a plug-and-play VF loss. This alignment targets both pointwise features and relative feature distributions while adaptively balancing the new loss with reconstruction training.
- VA-VAE constrains the tokenizer’s latent space using features from a vision foundation model to improve its suitability for generative tasks.
- The VF loss aligns projected tokenizer latents with foundation-model representations without changing the VAE architecture or training pipeline.
- VF loss combines marginal cosine similarity and marginal distance-matrix similarity losses as a decoupled, plug-and-play module.
- Adaptive weighting balances VF and reconstruction losses by matching their encoder-gradient scales before backpropagation.
- The authors evaluate VF loss as a route toward the latent-diffusion Pareto frontier spanning reconstruction and generation.
4. Improved Diffusion Transformer
LightningDiT improves DiT convergence through training and architectural strategies, while VF loss addresses high-dimensional tokenizer trade-offs. The resulting system reaches strong ImageNet generation performance with substantially fewer training samples.
- LightningDiT targets DiT’s slow ImageNet convergence with an open-source, fast-converging training pipeline and advanced design strategies.
- Higher tokenizer dimensionality improves reconstruction but reduces generation quality, whereas VF loss improves high-dimensional generation with minimal reconstruction impact.
- LightningDiT uses computational, training, and architectural optimizations, including torch.compile and bfloat16 training for acceleration.
- FID-50k=7.13 is achieved on ImageNet class-conditional generation with SD-VAE at around 80 epochs, versus over 1400 epochs for original DiT and SiT models.
5. Experiments
Experiments evaluate reconstruction-generation trade-offs across tokenizers and model scales, then measure VF Loss and the integrated VA-VAE–LightningDiT system for convergence, scalability, and ImageNet generation.
- Reconstruction-generation trade-off: Increasing tokenizer dimension lowers reconstruction rFID but raises corresponding generation FID, demonstrating the reconstruction-generation optimization dilemma.
- Foundation-model alignment: VF Loss improves generative performance for high-dimensional f16d32 and f16d64 tokenizers across LightningDiT scales, but is unnecessary for lower-dimensional f16d16 tokenizers.
- Convergence: 2.54× and 2.76× convergence speedups are achieved with VF Loss on f16d32 and f16d64 tokenizers, respectively.
- Scalability: Beyond 1B parameters, f16d32 with VF DINOv2 increasingly separates from f16d16, demonstrating stronger scalability, while the unaligned high-dimensional tokenizer retains a substantial gap at 1.6B.
- System-level generation: FID 1.35 is achieved at 800 epochs, while FID 2.11 is reached at 64 epochs, representing over 21× faster convergence than the original DiT.
6. Ablations and Discussions
Ablations examine tokenizer architecture, foundation-model choices, and VF Loss components, while discussions connect latent-space uniformity with generative performance.
- Tokenizer and architecture: Replacing SD-VAE with VA-VAE reduces FID-50k from 7.13 to 4.29 when using patch size 1 instead of patch size 2.
- Foundation-model ablation: DINOv2, MAE, CLIP, and Segment Anything are evaluated as foundation models, with DINOv2 achieving the highest generative performance.
- VF Loss formulation: Removing margin cosine similarity, margin distance matrix similarity, or the margin from VF Loss causes performance degradation, validating the effectiveness of these components.
- Discussion: Discrete tokenizers exhibit reduced codebook utilization as codebook size increases, paralleling the continuous-tokenizer conflict between reconstruction and generation.
- Latent-space analysis: VF Loss improves high-dimensional latent-distribution uniformity, which shows a possible positive correlation with generative performance.
7. Conclusion
The paper addresses latent diffusion’s optimization dilemma with VA-VAE and LightningDiT, combining latent-space alignment with DiT improvements for reconstruction and generation.
- Contributions: VA-VAE uses VF Loss to align VAE latents with vision models through marginal cosine similarity and distance matrix losses, producing a more uniform distribution.
- Discussion: Table 6 evaluates feature-distribution uniformity and reports a possible positive correlation between uniformity and generative performance.
- Contributions: Up to 2.8× faster convergence is reported for VF Loss, while LightningDiT incorporates advanced training techniques and architectural improvements for faster DiT convergence.
- Results: The integrated system combines VA-VAE reconstruction quality of rFID=0.28 with LightningDiT to achieve FID 1.35 on ImageNet 256.