DiT

发表于 分类于 阅读次数： 本文字数： 2.2k 阅读时长 ≈ 8 分钟

Scalable Diffusion Models with Transformers

作者是来自UC Berkeley的William Peebles和NYU的Saining Xie。论文引用[1]:Peebles, William S. and Saining Xie. “Scalable Diffusion Models with Transformers.” 2023 IEEE/CVF International Conference on Computer Vision (ICCV) (2022): 4172-4182.

### Time

Key Words

  • Diffusion Model

总结

  1. 作者探索了基于Transformer的一种新的diffusion model, 作者训练image的latent diffusion model，用transformer代替常用的U-Net的backbone，在latent patches上进行操作。作者通过前向传播复杂度(以Gflops)为指标，分析了Diffusion Transformer的可扩展性。作者发现，通过增加Transformer的深度/宽度，或增加输入tokens的数量，从而提高Gflops的DiTs，始终具有更低的FID分数，除了具备良好的可扩展性之外，最大的DiT-XL/2 models在ImageNet数据集上超过了所有的diffusion models

  2. ML正在经历由Transformer开启的革命，在过去的5年, NLP、CV和其它的领域的架构都被transformer取代，很多image-level的生成model也保持这一趋势，transformers见证了自回归model的广泛应用，它们在生成model中很少被采样，例如diffusion models在image-level的生成models中是最前沿的了，它们采样U-Net架构作为事实上的backbone。Ho等人的工作首次为diffusion model引入了U-Net，U-Net在像素级自回归模型和条件生成对抗网络（conditional GANs）中取得成功，其架构由 PixelCNN++继承而来，并进行了少量修改，这个model是卷积的，主要由ResNet blocks组成，相比于标准的U-Net，额外的spatial self-attention blocks，也插入到了lower resolution 部分，Dhariwal和Nichol 对U-Net的多种架构选择进行了消融实验，例如使用自适应归一化层来注入条件信息，以及调整卷积层的通道数。然而，Ho等人提出的U-Net高层设计基本保持不变。 作者旨在深入地解释diffusion models的架构选择的重要性，为了generative modeling的研究提供了经验性的baselines。作者展示了：U-Net的归纳偏置对于diffusion models的性能是不重要的，他们能被一些标准的设计例如transformers所替代，因此，diffusion models能够从最近的架构的发展中受益，通过集成best practices和其它domains的训练策略，保持良好的特性例如scalability，robustness和efficiency，一个标准的架构为cross-domain的研究提供了新的可能。

    在本文中，作者聚焦于基于transformer的新的diffusion models，称之为Diffusion models，DiTs拥有ViTs的best practicesx，相比于传统的conv net在visual recognition上展示出了更强的扩展性。更具体地，在网络复杂性和sample quality上，作者研究了transformers的scaling behavior，我们表明，通过在潜在扩散模型（Latent Diffusion Models, LDMs）框架下构建并基准测试DiT设计空间（其中扩散模型在变分自编码器的潜在空间中训练），我们能够成功地用Transformer替代传统的U-Net骨干网络。进一步研究表明，DiTs是适用于扩散模型的可扩展架构：网络复杂度（以Gflops衡量）与生成样本质量（以FID衡量）之间存在显著正相关性。通过简单地扩大DiT规模并在LDM中采用高容量骨干网络（118.6 Gflops），我们在类条件256×256 ImageNet生成基准测试中实现了2.27的FID，达到了当时最先进的性能。

  3. Diffusion 和score-based的生成模型，作为图像生成模型取得了成功，超过了GANs，DDPMs的改进由improved sampling方法驱动的，最有名的是classifier-free guidance，将diffusion models转换为预测noise而不是pixels，用级联的DDPMs pipelines，其中低分辨率基础扩散模型与上采样器并行训练，对于所有的diffusion models，Conv U-Nets是事实上的backbone架构的选择，当前的工作引入了一个新的，高效的架构基于DDPMs。本文对模型复杂度的分析主要通过理论Gflops，这使我们的研究与架构设计领域的主流方法保持一致——在该领域，Gfops（每秒十亿次浮点运算）被广泛用于衡量复杂度。在实际应用中，黄金标准复杂度指标仍存在争议，因其通常依赖于具体应用场景。Nichol 和 Dhariwal 改进扩散模型的开创性工作与本文最为相关——他们分析了 U-Net 架构类别的可扩展性和 Gfops 属性。而本文则聚焦于 Transformer 架构类别。

  4. Diffusion formulation：在引入架构之前，简要回顾一下DDPMs的基本概念，高斯diffusion models假设一个前向的noising process，逐渐对real data加噪声。diffusion models被训练用来反向过程，这个reverse process model用variational lower bound进行训练。 条件diffusion models将额外的信息作为输入，例如一个class label c，在这个例子中，classifier-free guidance用于鼓励sampling procedure来找到x，使得 \(logp(c|x)\) 是high。在高分辨率像素空间中直接训练扩散模型可能会导致计算成本过高，难以实现。Latent diffusion models用一个two-stage方法来解决这个问题： 学习一个autoencoder，将images压缩到一个更小的spatial representations中，训练一个representation的diffusion model，而不是image的diffusion model，新的images通过从diffusion model中采样一个representation z来生成，后将用学习到的decoder将其解码成image。潜在扩散模型（LDMs）在显著降低计算复杂度（以 Gfops 衡量）的同时实现了优异性能，其计算量仅为像素空间扩散模型（如 ADM）的一小部分。由于我们关注计算效率，这使得 LDM 成为架构探索的理想起点。本文将 DiTs 应用于潜在空间，尽管 DiTs 也可直接用于像素空间而无需修改。这种设计使我们的图像生成流水线成为一种混合式方法：我们采用现成的卷积 VAE 和基于 Transformer 的 DDPMs。

  5. 作者介绍了DiTs，作者尽可能地忠于原始的transformer架构，保持其scaling properties。

    • Patchify: DiT的输入是spatial representation z，DiT的第一个layer是patchify，将spatial input转换成T个tokens的sequences，每个是dimension d，通过线性映射每个patch。在patchify之后，作者用了标准的ViT frequency-based positional embeddings，tokens的数量T是由超参数的patch size p决定。
    • DiT block：在patchify之后，输入的tokens通过一系列的transformer blocks进行处理，除了noised image inputs，diffusion models还处理额外的conditional information，例如noise timesteps t，class labels c等，作者探索了四个transformer blocks的variants，这个设计引入了small，但是重要的modifications
    • In-context conditioning，作者简单地将t和c的embeddings作为additional tokens加载input sequence上，在最后的block，作者去掉了conditioning tokens。
    • Cross-attention block：将t和c的embeddings concate为一个Length-two的sequence，和image token sequence 分开，修改transformer block，包含一个额外的多头cross-attention layer。
    • Adaptive layer norm(adaLN) block：在GANs中的adaptive normalization layers的广泛使用，作者探索用adaLN代替transformer blocks中的norm layers，不是直接学习dimension-wise scale和shift parameters \(\gamma\) and \(\beta\)，作者从t和c的embedding vectors的sums中回归出它们，在探索的三个block designs中，adaLN增加了最少的Gflops，是最计算高效的，它也是唯一的conditioning 机制，将same function用于所有的tokens。
    • adaLN-Zero block：之前的ResNets上的工作表明：初始化每个residual block 为identity function是有益的，Diffusion U-Net用了类似的初始化策略，zero-initializing每个block的最后的conv lyaer，作者探索了adaLN DiT block的修改版本，除了回归 \(\gamma\) 和 \(\beta\)，作者也回归dimension-wise scaling 参数 \(\alpha\)。
    • 作者用了一系列的N个DiT blocks，每次都在hidden dimension size d 上进行操作，跟着ViT，作者用了标准的transformer configs，特别地，作者用了4个configs，DiT-S, DiT-B, DiT-L和DiT-XL，它们覆盖了很大的Model size和flop allocations，从0.3 到118.6 Gflops，使得能够衡量scaling performance。
    • Transformer Decoder：在最后的DiT block之后，作者需要将image tokens的序列解码成output noise prediction和一个output diagonal convariance prediction，这些输出和原始的spatial input有着相同的shape，作者用了一个标注的linar decoder来做这个，用最后的layer norm，线性地将每个token解码为一个 \(p \time p \time 2C\) 的tensor，\(C\) 是spatial input的channels数，最后，将所有的decoded tokens rearrange为它们原始的spatial layout，来得到预测的noise和convariance。

architecture \(Fig.1^{[1]}\) 左边训练一个conditional latent DiT 模型，将input 解构成patches输入到DiT blocks，右边是几种不同的variants