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

数据派THU

于 2025-04-13 17:01:05 发布

阅读量145

点赞数

原文链接：https://mp.weixin.qq.com/s?__biz=MzI1MjQ2OTQ3Ng==&mid=2247655506&idx=2&sn=afcb6030eda6f73c56b00f9feac0feca&chksm=e8e8fe191aa33f7c01a1de444301bbe5e0d1aab48bebd246680b302b6e7ac3c72ca002a7a01e&scene=126&sessionid=0

版权

来源：极市平台  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) 组成，是一个即插即用模块，在不改变模型架构和训练流程的情况下解决优化困境。