[基础] DiT: Scalable Diffusion Models with Transformers

名称

DiT: Scalable Diffusion Models with Transformers
时间：23/03
机构：UC Berkeley && NYU

TL;DR

提出首个基于Transformer的Diffusion Model，效果打败SD，并且DiT在图像生成任务上随着Flops增加效果会降低，比较符合scaling law。后续sora的DM也使用该网络架构。

Method

网络结构整体参考LDM，只不过将latent diffusion中的UNet替换为ViT，image与latent space之间编解码复用现成的基于卷积的VAE。DiT网络结构如下图所示，作者尝试了多种DiT blocks来编码condition信息，比如，cross-attention, in-context conditioning(直接concat with embeding tokens)，最终发现adaLN-Zero block效果最好。

adaLN-Zero block

adaLN的全称是adaptive layer norm。layer norm是"逐样本"将均值方差替换为可学习参数beta/gamma的方法，而这里adaptive指得是额外学习一个逐channel的参数alpha。Zero指得是将residual连接的前一层MLP的norm中alpha初始值设定为0，起到Identity的作用，个人理解，这么做相当于在网络学习前期去掉了一些层，降低了前期学习复杂度加速收敛。更详细的解释参考论文

Patchify

类似于ViT中的图片切块出embedding的操作，只不过这里专门起了一个名字。如下图所示，DiT的输入尺寸是一个I x I的保留spatial shape的noise latent feature，被切成大小为pxp的patch，再将patch抽成embedding，把所有embedding组成一个序列的tokens。