[论文解读] Your ViT is Secretly a Hybrid Discriminative-Generative Diffusion Model

Xiulong Yang, Sheng-Min Shih+3 more|arXiv (Cornell University)|Aug 16, 2022

Generative Adversarial Networks and Image Synthesis被引用 15

一句话总结

本文提出GenViT与HybViT，两种新颖模型，将扩散生成建模与视觉Transformer（ViT）统一，采用单一ViT主干网络同时实现图像生成与分类。通过将ViT整合至DDPM框架，该方法在两项任务中均实现最先进性能，训练过程稳定且无需MCMC，相较于先前的混合模型，计算效率更高。

ABSTRACT

Diffusion Denoising Probability Models (DDPM) and Vision Transformer (ViT) have demonstrated significant progress in generative tasks and discriminative tasks, respectively, and thus far these models have largely been developed in their own domains. In this paper, we establish a direct connection between DDPM and ViT by integrating the ViT architecture into DDPM, and introduce a new generative model called Generative ViT (GenViT). The modeling flexibility of ViT enables us to further extend GenViT to hybrid discriminative-generative modeling, and introduce a Hybrid ViT (HybViT). Our work is among the first to explore a single ViT for image generation and classification jointly. We conduct a series of experiments to analyze the performance of proposed models and demonstrate their superiority over prior state-of-the-arts in both generative and discriminative tasks. Our code and pre-trained models can be found in https://github.com/sndnyang/Diffusion_ViT .

研究动机与目标

• 探索单一视觉Transformer（ViT）是否可作为生成模型，挑战DDPM中传统UNet的使用。
• 通过使单一ViT同时执行图像分类与生成，弥合判别式建模与生成式建模之间的差距。
• 克服如JEM与JEM++等基于MCMC的混合模型所存在的训练不稳定与高计算成本问题。
• 证明基于ViT的扩散模型相较于基于CNN或GAN的替代方案，在可扩展性、稳定性与效率方面表现更优。
• 分析所提模型在对抗鲁棒性、不确定性校准与分布外（OOD）检测性能方面，相较于现有基准的表现。

提出的方法

• 将视觉Transformer架构直接集成至DDPM框架，用ViT替代标准UNet主干网络，用于扩散过程中的去噪。
• 通过变分下界（VLB）目标函数，训练ViT在反向扩散过程中预测噪声，从而构建纯生成模型GenViT。
• 设计HybViT作为混合模型，通过共享ViT特征空间与两个独立头，联合优化图像重建（通过扩散）与分类任务。
• 为生成与判别任务共享同一组ViT参数，实现参数效率与统计信息共享。
• 采用单一目标函数，结合重建损失与交叉熵损失，避免训练过程中使用昂贵的MCMC采样。
• 应用基于图像块的处理方式，结合可学习线性投影与多头自注意力机制，以建模潜在扩散空间中的长距离依赖关系。

实验结果

研究问题

• RQ1 单一视觉Transformer能否有效用作基于扩散的图像生成主干网络，替代标准UNet？
• RQ2 在共享ViT主干网络上联合训练生成与判别头，对图像生成与分类任务的性能有何影响？
• RQ3 所提方法在训练稳定性、收敛速度与最终性能方面，是否优于基于MCMC的混合模型如JEM与JEM++？
• RQ4 在所提模型中，图像块大小与模型维度对生成质量与分类准确率之间权衡的影响如何？
• RQ5 GenViT与HybViT在似然度、对抗鲁棒性、不确定性校准与OOD检测方面，相较于最先进模型表现如何？

主要发现

• HybViT在STL10上实现87.1%的测试准确率，Inception Score为6.90，FID为125.5，优于先前的混合模型。
• GenViT在模型维度为768时，在CIFAR-10上实现Inception Score 7.01与FID 126.6，展现出强大的生成性能。
• HybViT训练速度显著快于基于MCMC的模型：在单张GPU上，CIFAR-10训练仅耗时31.2小时，而JEM(K=20)需101.3小时。
• 采用6×6图像块大小的模型（HybViT-ps6）在准确率与生成质量之间达到最佳平衡，准确率为81.7%，IS为7.30。
• 与基线模型相比，HybViT在不确定性校准与OOD检测性能方面表现更优，表明其鲁棒性更强。
• 尽管结果优异，高分辨率图像的生成质量在图像块大小为6时趋于饱和，表明限制因素在于线性投影层，而非自注意力机制。