VA-VAE：视觉基础模型对齐，突破潜在扩散模型重建与生成瓶颈

AI资讯
1

华中科大提出VA-VAE，通过视觉基础模型对齐，解决潜在扩散模型重建与生成之间的优化难题，LightningDiT在ImageNet 256x256上取得SOTA，FID 1.35，训练效率提升21倍。

原文标题：CVPR 2025 满分论文！重建 vs 生成：解决扩散模型中的优化难题

原文作者：数据派THU

冷月清谈：

该文章介绍了一种名为VA-VAE的新方法，旨在解决潜在扩散模型中重建质量与生成性能之间的优化难题。现有方法在提升重建质量时往往需要更大的扩散模型和更多的训练迭代，导致计算成本高昂或效果不佳。VA-VAE的核心思想是在训练视觉词元分析器时，将其潜在空间与预先训练的视觉基础模型对齐，从而显著扩展潜在扩散模型的重建生成边界，使高维潜在空间中的Diffusion Transformers (DiT) 能够更快地收敛。文章还介绍了基于VA-VAE构建的LightningDiT模型，该模型在ImageNet 256x256生成任务上取得了SOTA性能，同时显著提升了训练效率。VA-VAE通过VF损失（包括边缘余弦相似度损失和边缘距离矩阵相似度损失）优化潜在空间，无需改变模型架构和训练流程。实验结果表明，使用VF损失可以显著提升模型收敛速度，并在重建和生成任务中均表现出色。

怜星夜思：

1、VA-VAE方法中，将视觉词元分析器的潜在空间与预训练的视觉基础模型对齐，为什么能解决重建和生成之间的优化难题？这种对齐操作背后的原理是什么？
2、文章中提到LightningDiT在短短64个epoch内就达到了很高的FID得分，训练效率提升显著。除了VA-VAE的贡献，LightningDiT在训练策略和架构设计上做了哪些优化？这些优化对于加速收敛起到了什么作用？
3、VA-VAE方法依赖于预训练的视觉基础模型，那么选择不同的视觉基础模型会对最终的生成效果产生什么影响？未来是否可以探索使用自监督学习得到的视觉基础模型，或者训练一个专门为VA-VAE量身定制的视觉基础模型？

原文内容

来源：极市平台  Deeeep Learning  ‍‍‍‍
本文共1400字，建议阅读5分钟
本文提出了一种名为VA-VAE的方法，通过将视觉词元分析器的潜在空间与预训练的视觉基础模型对齐，解决了潜在扩散模型中重建与生成之间的优化难题。


本文提出了一种名为VA-VAE的方法，通过将视觉词元分析器的潜在空间与预训练的视觉基础模型对齐，解决了潜在扩散模型中重建与生成之间的优化难题。基于该方法构建的LightningDiT模型在ImageNet 256x256生成任务上取得了最佳性能，FID得分1.35，并在64个epoch内达到2.11的FID得分，显著提升了训练效率。

重建vs 生成：解决扩散模型中的优化难题

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

作者：Jingfeng Yao, Xinggang Wang

作者单位：华中科技大学Paper：https://arxiv.org/abs/2412.04852

Code：https://github.com/hustvl/LightningDiT


01 介绍


两阶段的潜在扩散模型中存在优化难题：在visual tokenizer中增加每个标记的特征维度，虽能提升重建质量，但要达到相近的生成性能，却需要大得多的扩散模型和更多训练迭代。因此，现有系统常常只能采用次优解决方案，要么因tokenizer中的信息丢失而产生视觉伪影，要么因计算成本高昂而无法完全收敛。作者认为这种困境源于学习不受约束的高维潜在空间的固有困难。

为了解决这一问题，作者建议在训练视觉词元分析器时，将潜在空间与预先训练的视觉基础模型对齐。提出的VA-VAE（视觉基础模型结合变分自动编码器）显著扩展了潜在扩散模型的重建生成边界，使高维潜在空间中的Diffusion Transformers(DiT) 能够更快地收敛。为了充分发挥VA-VAE的潜力，构建了一个增强型DiT基线，改进了训练策略和架构设计，称为LightningDiT。在ImageNet 256x256 生成上实现了最佳 (SOTA) 性能，FID得分为1.35，同时在短短64个epoch内就达到了2.11的FID得分，展现了卓越的训练效率——与原始DiT相比，收敛速度提高了21倍以上。

相关工作

  • 可视化生成的tokenizer

visual tokenizer包括以变分自编码器（VAE）为代表的连续型和 VQVAE、VQGAN 等离散型。离散型词元分析器虽然能提高重建保真度，但编码对照本利用率低下，对生成性能产生不利影响。连续型tokenizer通过增加词元分析器的特征维度会提高重建质量，但会降低生成性能，还需要大幅增加训练成本，当前缺乏对连续型 VAE 优化的有效解决方案。

  • 扩散Transformer的快速收敛

扩散Transformer（DiT）目前是潜在扩散模型最常用的实现方式，存在收敛速度慢的问题，往研究提出多种加速方法，本文则从优化视觉词元分析器学习的潜在空间入手，在不修改扩散模型的情况下实现更快收敛，并对 DiT 进行了训练策略和架构设计优化。

02 方法


  • 网络架构

VA-VAE基于VQGAN模型架构，通过视觉基础模型对齐损失（VF损失）优化潜在空间。VF损失由边缘余弦相似度损失(Marginal Cosine Similarity Loss) 和边缘距离矩阵相似度损失(Marginal Distance Matrix Similarity Loss) 组成，是一个即插即用模块，在不改变模型架构和训练流程的情况下解决优化困境。


  • 边际余弦相似度损失

将视觉标记器编码器输出的图像潜在特征投影后，与冻结的视觉基础模型输出的特征计算余弦相似度，通过 ReLU 函数和设置边际值，使相似度低于边际值的特征对损失有贡献，从而聚焦于对齐差异较大的特征对。


  • 边际余弦相似度损失

将视觉标记器编码器输出的图像潜在特征投影后，与冻结的视觉基础模型输出的特征
计算余弦相似度，通过ReLU函数和设置边际值，使相似度低于边际值的特征对损失有贡献，从而聚焦于对齐差异较大的特征对。


03 结果


visual tokenizer采用LDM的架构和训练策略，使用VQGAN网络结构和KL损失，训练三种不同的f16标记器（无VF损失，VF损失（MAE），VF损失（DINOv2））。生成模型采用LightningDiT，在ImageNet上以256分辨率训练，设置相关参数和训练策略。

使用 VF 损失在f16d32和f16d64收敛速度提升明显



8种不同tokenizer的重建和生成的评估



与现有扩散模型对比



可视化效果


图片
图片831×387 523 KB

编辑：王菁‍‍‍



关于我们

数据派THU作为数据科学类公众号，背靠清华大学大数据研究中心，分享前沿数据科学与大数据技术创新研究动态、持续传播数据科学知识，努力建设数据人才聚集平台、打造中国大数据最强集团军。



图片
图片1080×564 45.8 KB


新浪微博：@数据派THU

微信视频号：数据派THU

今日头条：数据派THU

2

会不会是用了更大的batch size和更强的优化器？更大的batch size可以提高GPU的利用率，减少训练时间；更强的优化器，比如AdamW，可以更快地找到最优解。

3

使用自监督学习得到的视觉基础模型是一个很有潜力的方向。自监督学习不需要人工标注数据，可以利用海量的无标签数据进行训练，从而得到更鲁棒、更通用的视觉特征表示。这或许可以避免对特定数据集的过度依赖。

4

从信息论的角度看，对齐操作降低了潜在空间的“复杂度”或者“信息熵”。预训练的视觉基础模型已经包含了大量关于图像结构的先验知识，将其融入tokenizer的训练，本质上是引入了一种正则化，使得tokenizer学习到的latent space更具有可解释性和可控性，从而更容易被diffusion model利用。

5

这让我想到了迁移学习。预训练的视觉基础模型就像是一个经验丰富的老师，它已经掌握了很多图像相关的知识。VA-VAE就像是让学生（tokenizer）去学习老师的经验，这样学生就不用从零开始摸索，可以更快更好地掌握知识，并在此基础上进行创新（生成）。

6

选择不同的视觉基础模型肯定会对生成效果产生影响。不同的视觉基础模型关注的图像特征可能不同，例如，DINOv2更注重图像的语义信息，而MAE可能更注重图像的细节信息。因此，选择合适的视觉基础模型需要根据具体的生成任务进行调整。

7

我猜可能用了更激进的数据增强策略，比如Mixup或者CutMix。这些方法通过在训练过程中混合不同的图像，可以有效增加训练数据的多样性，提高模型的泛化能力，从而加速收敛。

8

文中并没有详细说明LightningDiT的训练策略和架构设计优化，这部分信息需要查阅原始论文。不过，根据经验，加速DiT训练通常会涉及以下几个方面：1) 更高效的attention机制，例如线性attention或稀疏attention；2) 更强的正则化手段，防止过拟合；3) 更好的学习率调度策略，例如warmup和cosine decay；4) 混合精度训练，利用半精度浮点数加速计算。

9

训练一个专门为VA-VAE量身定制的视觉基础模型听起来很有意思！可以针对diffusion model的特点，设计一个能够提取diffusion model所需特征的视觉基础模型，从而实现更好的对齐效果和生成性能。

10

VA-VAE的核心在于让tokenizer学习到的latent space更“规整”，更接近视觉基础模型已经学到的图像特征表示。当latent space的结构更合理，diffusion model就更容易在这个空间中进行生成，而不需要耗费大量的计算资源来“理解”这个latent space。这就像是，如果你的房间本来就很干净整洁（对齐的潜在空间），那么你稍微整理一下（生成）就能保持干净；但如果你的房间乱成一团（未对齐的潜在空间），那么你需要花费很多精力才能把它打扫干净。