重建与生成：潜在扩散模型中优化困境的调控
计算机视觉与模式识别
2025-03-11 v3 机器学习
摘要
具有 Transformer 架构的潜在扩散模型在生成高保真图像方面表现卓越。然而，近期研究揭示了这种两阶段设计中的优化困境：尽管增加视觉标记化器中每个标记的特征维度可提升重建质量，但为获得可 comparable 的生成性能，需大幅扩大扩散模型规模并增加训练迭代次数。因此，现有系统常为求次优解，或因标记化器中信息损失产生视觉伪影，或因计算成本高昂未能完全收敛。我们认为，这一困境源于在高维潜在空间中学习无约束表示的内在难度。为解决此问题，我们提出在训练视觉标记化器时，将潜在空间与预训练视觉基础模型对齐。我们提出的VA-VAE（Vision foundation model Aligned Variational AutoEncoder）显著扩展了潜在扩散模型的重建-生成边界，使扩散 Transformer（DiT）在高维潜在空间内实现更快收敛。为充分发挥VA-VAE的潜力，我们构建了增强型DiT基线，采用改进的训练策略和架构设计，称为LightningDiT。集成系统在ImageNet 256x256生成任务上实现了1.35的FID分数，同时仅需64个epoch即可达到2.11的FID分数，代表比原始DiT快21倍以上。模型和代码已公开：https://github.com/hustvl/LightningDiT。
引用
@article{arxiv.2501.01423,
title = {Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models},
author = {Jingfeng Yao and Bin Yang and Xinggang Wang},
journal= {arXiv preprint arXiv:2501.01423},
year = {2025}
}
备注
Models and codes are available at: https://github.com/hustvl/LightningDiT