文章地址:Scalable Diffusion Models with Transformers
简介
文章提出使用Transformers替换扩散模型中U-Net主干网络,分析发现,这种Diffusion Transformers(DiTs)不仅速度更快(更高的Gflops),而且在ImageNet 512×512和256×256的类别条件图片生成任务上,取得了更好的效果,256×256上实现了SOTA的FID指标(2.27)。
Transformers已经广泛应用于包括NLP、CV在内的机器学习的各个领域。然而,很多图片level的生成模型还坚持使用卷积神经网络,比如扩散模型采用的就是U-Net的主干网络架构。经过演化,扩散模型中的U-Net网络增加了稀疏的自注意力模块,此外 Dhariwal and Nichol 也尝试过在U-Net模型上的一些改变,比如通过增加适配的正则化层来注入条件信息和Channel数量。尽管如此,U-Net的顶层设计还是与原始U-Net相差无几。
文章的目标就是要揭开扩散模型架构选择的神秘面纱,提供一个强有力的baseline。文章发现U-Net并非不可替代,并且很容易使用诸如Transformers的结构替代U-Net,使用Transformers可以很好地保持原有的优秀特性,比如可伸缩性、鲁棒性、高效性等,并且使用新的标准化架构可能在跨领域研究上展现出更多的可能。文章从网络复杂度和采样质量两个方面对DiTs方法进行评估。
相关工作
Transformers
当前,Transformers架构已经应用在了文本、视觉、强化学习、元学习等多个领域,同时模型的大小、训练开销、数据量等也不断地上涨。在语言模型的启发下,有些工作在视觉任务上训练离散的codebook,这种架构可以同时应用在自回归模型和masked生成模型。本文将研究在扩散模型的主干网络上应用Transformers。
DDPMs
扩散模型是借鉴了物理学上的扩散过程,在生成模型上,分为正向和逆向的过程。正向过程是向信号中逐渐每步加少量噪声,当步数足够大时可以认为信号符合一个高斯分布。所以逆向过程就是从随机噪声出发逐渐的去噪,最终还原成原有的信号。
去噪过程一般采用UNet或者ViT,使用t步的结果和条件输入预测t-1步增加的噪声,然后使用DDPM可以得到t-1步的分布,经过多步迭代就可以从随机噪声还原到有实际意义的信号。如果使用原始DDPM速度会慢很多,所以很多工作如DDIM、FastDPM等工作实现了解码加速。
在图像的无条件生成任务上,扩散模型的性能已经超过了GANs,并且在有条件生成如文图生成任务上大放异彩。
架构复杂度
对于图片生成的迭代过程,我们可以使用参数量来衡量不同模型的复杂度。一般而言,参数量来评估模型复杂度不是很合适,因为参数量并不能代表模型的计算复杂度,比如当模型参数量相同时,图片分辨率不同会导致计算复杂度上较大的差异。所以文章采用Gflops来衡量模型架构的复杂度。
方法
扩散模型基础
前向过程是一个T步逐渐加噪的马尔科夫链,公式如下
给定前向扩散过程作为先验,扩散模型训练反转的过程,可以通过去除所加噪声从XT恢复成X0,并且每步的扩散过程都采样自特定的高斯分布,其期望和方差如下:
暂无评论内容