U-ViT: A ViT Backbone for Diffusion Models

李天骥

Paper: All are Worth Words: A ViT Backbone for Diffusion Models
Code: https://github.com/baofff/U-ViT

本文介绍本组与清华大学朱军教授课题组以及北京智源研究院曹越研究员的合作工作：U-ViT:  A ViT Backbone for Diffusion Models
概括

最近，扩散概率模型(diffusion model)在图像生成领域大红大紫，出现了stable-diffusion，Imagen等一系列杰出的工作。扩散概率模型从2015年首次提出至今， 在概率建模方面取得了许多进展，其主干网络也得到了许多改进，从2015年Deep Unsupervised Learning using Nonequilibrium Thermodynamics一文的MLP，到2019年Song Yang在Generative Modeling by Estimating Gradients of the Data Distribution 一文中首次使用U-Net建模score-based model (即diffusion model)，后续DDPM ，ADM ，Imagen 等许多工作对U-Net进行了一系列改进。目前，绝大多数扩散概率模型的论文依然使用U-Net作为主干网络。
于此同时，在自然语言处理，以及计算机视觉领域，Transformer架构都展现出大杀四方的特质，不仅在各项任务上取得了很好的效果，也能很好的扩展到多模态学习上。基于这样的背景，U-ViT应运而生。U-ViT将扩散概率模型和Transformer结合，主要展现了以下两类能力：

• 图像生成的SOTA FID 分数
  
• 多模态数据的融合
  

目前，U-ViT已经被CVPR2023接收。
U-ViT网络结构



如上图所示，U-ViT延续了ViT的方法，将带噪图片划分为多个patch之后，将时间t，条件c，和图像patch视作token输入到Transformer block，同时在网络浅层和深层之间引入long skip connection。
Long skip connection

直观上理解，扩散概率模型中的噪声预测网络是像素级别的预测任务，对low-level feature敏感，long skip connection为连接low-level feature提供了快捷方式，所以有助于网络的训练。 对于网络主分支的特征h_m\in \mathbb{R}^{B,L,D}和来自long skip connection的特征 h_s \in \mathbb{R}^{B,L,D} ，其中B为batch_size，L为token 数目，D为每一个token的长度。作者探究了以下几种融合 h_m,h_s 的方式。  1. Linear(Concate([h_m, h_s],dim=-1))\\ 2. h_s + h_m\\ 3. h_m + Linear(h_s)\\ 4. Linear(h_s + h_m)\\ 5. No \ long \ skip \ connection  


可以看出，long skip connection对于图像生成的FID分数是至关重要的。
Patch size

假设patch size为 p ，图像分辨率为 I \times I ，则输入Transformer的token数目 L=\left( \frac{I}{p} \right)^2 ，patch size 减少为一半，token数目增加为4倍，在一般情况下，模型的GFLOPs 也增加为4倍。所以即使在相同参数量的情况下，patch size不同，训练或者推理所需的计算量也不同。作者探究了不同patch size对图像生成结果的影响。


The way to feed the time into the network

作者探究了两种将衡量噪声大小的时间t嵌入网络的方法：

• 将时间t视作token
  
• 类似于adaptive group normalization，在LayerNorm后插入时间t，即 AdaLN(h,t)=t_sLayerNorm(h) + t_b ，h是transfomer block的特征， t_s, t_b 是t经过线性投影后得到的特征。
  

实验

U-ViT的配置参数如下表所示，在本文后续的介绍中，例如，U-ViT-H/2 代表 U-ViT-Huge 模型，patch size 为2。


对于无条件生成，U-ViT在CIFAT10，CelebA数据集进行了实验；对于类别条件生成，U-ViT在ImageNet64，ImageNet256，ImageNet512数据集进行了实验； 对于文到图生成，U-ViT在MSCOCO数据集进行了实验。
高分辨率数据集，即ImageNet256，ImageNet512和MSCOCO（分辨率256），U-ViT跟随了Lantent diffusion的思路，使用Stable-diffusion 预训练好的图像AutoEncoder，把 256 \times 256 和 512 \times 512 分辨率的图片降采样到 32 \times 32 和 64 \times 64 ，然后建模图像特征空间的概率分布。


从上表可以看到，UViT在图像无条件生成以及类别条件生成上取得了和其他模型可比或者更优的FID，在ImageNet256数据集上取得了SOTA的FID。并且可以进一步证明，和用U-Net建模特征空间的Latent Diffusion 相比，在使用相同采样器（dpm_solver）和相同采样步数的情况下，U-ViT均能取得更优的表现。


作者同样尝试了用参数量、计算量相似的U-ViT替换掉Latent Diffusion 中的U-Net，进一步证明了U-ViT优越性。下图中U-ViT的实现跟随了ADM的实现，参数量为646M，FLOPs为135G；U-ViT参数量为501M，FLOPs为133G。


除了在像素空间证明了long skip connection的优越性（上一节），作者也在ImageNet256数据上进行实验，验证long skip connection在特征空间的有效性。


U-ViT也展现了杰出的多模态融合能力，在没有额外数据的情况下，U-ViT取得了MS-COCO数据集上text-to-image generation任务的SOTA FID。


利用在ImageNet512训练的U-ViT生成图片：


利用在ImageNet256训练的U-ViT生成图片：


在MSCOCO上训练的U-ViT和U-Net文生图对比：


总结

U-ViT是一种简单且通用的基于ViT的扩散概率模型的主干网络，U-ViT把所有输入，包括图片、时间、条件都当作token输入，并且引入了long skip connection。U-ViT在无条件生成、类别条件生成以及文到图生成上均取得了可比或者优于CNN的结果。为未来扩散模型中骨干网络研究提供见解，并有利于大规模跨模态数据集的生成建模。
U-ViT的一篇同期工作 DiT: Scalable Diffusion Models with Transformers 也提出了使用ViT代替U-Net的思想，不同的是DiT中没有引入long skip connection也依然取得了杰出的效果。笔者推测原因可能有：

• DiT 出色的Adaptive layer norm以及零初始化的设计能够有效提升生成质量；
  
• DiT 在建模特征空间表现良好，但在建模像素空间表现欠缺，可能在用扩散概率模型建模像素空间分布时long skip connection是至关重要的；
  
• 即使在建模特征空间上，DiT 没有long skip connection也能取得很好的效果，但long skip connection在加速收敛方面也起着关键的作用，如下图所示：
  

\begin{array}{|c|c|c|c|c|c|c|}    \hline Data & Model & Params(M) & Flops(G) & Batch Size & Training Iteration & Fid \\   \hline ImageNet256 (Latent Diffusion) & U-ViT-H/2 & 501 & 133 & 1024 & 50k & 2.29 \\    \hline ImageNet256 (Latent Diffusion)  & DiT-XL/2  & 675 & 118 & 256 & 700k & 2.27 \\   \hline  \end{array}  
欢迎阅读我们的论文 All are Worth Words: A ViT Backbone for Diffusion Models
作者列表：鲍凡，聂燊，薛凯文，曹越，李崇轩，苏航，朱军
PS: 转载请注明出处

海魂衫

请问可以以图像作为条件吗？