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

具有Transformer架构的潜在扩散模型在生成高保真图像方面表现出色。然而，最近的研究揭示了这种两阶段设计中的一个优化困境：虽然在视觉分词器中增加每个标记的特征维度可以提高重建质量，但这需要显著更大的扩散模型和更多的训练迭代次数才能达到相当的生成性能。因此，现有的系统往往只能选择次优解决方案，要么因为分词器内的信息丢失而产生视觉伪影，要么由于高昂的计算成本而无法完全收敛。我们认为，这一困境源于学习无约束高维潜在空间的固有难度。为了解决这个问题，我们提出在训练视觉分词器时，将潜在空间与预训练的视觉基础模型对齐。我们提出的VA-VAE（Vision foundation model Aligned Variational AutoEncoder）显著扩展了潜在扩散模型的重建-生成前沿，使Diffusion Transformers (DiT) 在高维潜在空间中能够更快地收敛。为了充分利用VA-VAE的潜力，我们构建了一个增强的DiT基线，改进了训练策略和架构设计，称为LightningDiT。集成系统在ImageNet 256x256生成任务上达到了最先进的（SOTA）性能，FID得分为1.35，同时展示了卓越的训练效率，在仅64个epoch内就达到了2.11的FID得分——这比原始DiT的收敛速度提高了超过21倍。模型和代码可在以下链接获取： https://github.com/hustvl/LightningDiT。

该论文试图解决在使用Transformer架构的潜在扩散模型生成高保真图像时遇到的优化难题。具体来说，当增加视觉分词器中的每令牌特征维度以提高重建质量时，需要显著更大的扩散模型和更多的训练迭代才能达到相当的生成性能，导致现有系统要么因信息丢失而产生视觉伪影，要么因计算成本高昂而无法完全收敛。这是一个新问题，特别是在结合高维潜空间与两阶段设计（分词器和扩散模型）的情况下。

论文的关键思路是通过在训练视觉分词器时将潜空间与预训练的视觉基础模型对齐来解决这一难题。提出了一种称为VA-VAE（Vision foundation model Aligned Variational AutoEncoder）的方法，这种方法显著扩展了潜在扩散模型的重建-生成前沿，使得Diffusion Transformers (DiT) 在高维潜空间中更快地收敛。相比当前研究，此方法新颖之处在于它利用了预训练视觉模型的知识，从而解决了高维潜空间学习的固有困难。

该论文值得关注的地方包括：1) 提出了VA-VAE方法，大幅提升了ImageNet 256x256图像生成的质量，FID得分为1.35；2) 构建了一个增强版的DiT基线——LightningDiT，具有改进的训练策略和架构设计；3) 实验表明，在仅64个epoch内就能达到FID 2.11，比原DiT快超过21倍；4) 开源了模型和代码，方便后续研究。未来值得深入的研究方向包括探索其他视觉基础模型对齐的效果以及进一步优化训练效率。

近期相关研究包括：1)《Latent Diffusion Models for Image Generation》探讨了如何利用潜在扩散模型生成高质量图像；2)《Improving Generative Adversarial Networks with Vision Transformers》研究了视觉Transformer在GAN中的应用；3)《High-Fidelity Image Generation with Fewer Iterations through Optimized Architectures》讨论了通过优化架构减少迭代次数以提高图像生成保真度的方法。