广告位:
图像拼接论文精读专栏 —— 图像拼接领域论文全覆盖(包含数据集),省时省力读论文,带你理解晦涩难懂的论文算法,学习零散的知识和数学原理,并学会写图像拼接领域的论文(介绍、相关工作、算法、实验、结论、并附有参考文献,不用一篇一篇文章再找)
图像拼接论文源码精读专栏 —— 图像拼接有源码的论文全覆盖(有的自己复现),帮助你通过源码进一步理解论文算法,助你做实验,跑出拼接结果,得到评价指标RMSE、SSIM、PSNR等,并寻找潜在创新点和改进提升思路。
超分辨率重建专栏 —— 从SRCNN开始,带你读论文,写代码,复现结果,找创新点,完成论文。手把手教,保姆级攻略。帮助你顺利毕业,熟练掌握超分技术。
有需要的同学可以点上面链接看看。
文章目录
前言
论文题目:Scalable Diffusion Models with Transformers —— VLOGGER:具身化身合成的多模态扩散
论文地址:Scalable Diffusion Models with Transformers
项目地址(源码,效果展示等):https://www.wpeebles.com/DiT.html
ICCV 2023 Oral!纯Transformer的扩散模型!研究扩散模型必读文章!
Abstract
我们基于变压器架构探索了一类新的扩散模型。我们训练图像的潜在扩散模型,用对潜在补丁进行操作的转换器替换常用的 U-Net 主干。我们通过 Gflops 测量的前向传递复杂度透镜来分析我们的扩散变换器 (DiTs) 的可扩展性。我们发现具有较高 Gflop 的 DiTs——通过增加变压器深度/宽度或输入标记数量的增加——始终具有较低的 FID。除了具有良好的可扩展性属性外,我们最大的 DiT-XL/2 模型在类条件 ImageNet 512×512 和 256×256 基准测试中优于所有先前的扩散模型,后者实现了 2.27 的最新 FID。
1. Introduction
机器学习正在经历由变压器驱动的复兴。在过去的五年中,自然语言处理[8,42]、视觉[10]和其他几个领域的神经架构在很大程度上被变压器[60]所包含。然而,许多类别的图像级生成模型仍然坚持趋势,而转换器在自回归模型中得到了广泛的应用[3,6,43,47],但它们在其他生成性建模框架中较少采用。例如,扩散模型一直是图像级生成模型最新进展的前沿[9,46];然而,它们都采用卷积U-Net架构作为骨干的实际选择。
Ho等人[19]的开创性工作首次引入了扩散模型的U-Net主干。在最初在像素级自回归模型和条件GANs[23]中取得了成功之后,U-Net继承自PixelCNN++[52,58],变化很少。该模型是卷积的,主要由 ResNet [15] 块组成。与标准 U-Net [49] 相比,额外的空间自注意力块(这是变压器中必不可少的组件)以较低分辨率穿插。Dhariwal 和 Nichol [9] 消融了 UNet 的几种架构选择,例如使用自适应归一化层 [40] 为卷积层注入条件信息和通道计数。然而,Ho 等人的 UNet 的高级设计在很大程度上保持不变。
在这项工作中,我们旨在揭开扩散模型中架构选择的重要性,并为未来生成建模研究提供经验基线。我们表明 U-Net 归纳偏差对扩散模型的性能并不重要,并且它们可以很容易地替换为变压器等标准设计。因此,扩散模型得到了很好的证明,受益于架构统一的最新趋势——例如,通过继承其他领域的最佳实践和训练食谱,并保留可扩展性、鲁棒性和效率等有利属性。标准化架构还将为跨域研究开辟了新的可能性。
在本文中,我们专注于一类基于变压器的新型扩散模型。我们称它们为扩散变压器,简称DiTs。DiTs遵循视觉变形器(ViTs)[10]的最佳实践,已被证明比传统的卷积网络(如ResNet[15])更有效地扩展到视觉识别。
更具体地说,我们研究了变压器相对于网络复杂度与样本质量的缩放行为。我们表明,通过在潜在扩散模型 (LDM) [48] 框架下构建和基准测试 DiT 设计空间,其中扩散模型在 VAE 的潜在空间内进行训练,我们可以成功地将 U-Net 主干替换为变压器。我们进一步表明,DiTs 是扩散模型的可扩展架构:网络复杂度(由 Gflops 测量)与样本质量(由 FID 测量)之间存在很强的相关性。通过简单地缩放DIT并使用高容量主干(118.6 Gflops)训练LDM,我们能够在类条件256 × 256 ImageNet生成基准上获得2.27 FID的最新结果。
2. Related Work
Transformers。Transformer[60]已经取代了跨语言、视觉[10]、强化学习[5,25]和元学习[39]的领域特定架构。它们在增加语言域[26]中的模型大小、训练计算和数据以及通用自回归模型[17]和VITS[63]下表现出了显著的缩放特性。除了语言之外,转换器还经过训练以自回归预测像素 [6, 7, 38]。他们还接受了离散码本[59]的训练,因为自回归模型[11,47]。前者显示了高达20B参数的出色缩放行为[62]。最后,在DDPM中探索了变压器来合成非空间数据;例如,在DALL·E 2[41,46]中生成CLIP图像嵌入。在本文中,我们研究了变压器作为图像扩散模型的主干时的缩放特性。
去噪扩散概率模型(DDPMs)。扩散[19, 54] 和基于分数的生成模型 [22, 56] 尤其成功,因为图像的生成模型 [35,46,48,50],在许多情况下优于以前最先进的生成对抗网络 (GAN) [12]。在过去的几年里,DDPM 的改进主要是由改进的采样技术 [19,27,55] 推动的,最著名的是无分类器指导 [21],重新制定扩散模型以预测噪声而不是像素 [19],并使用级联 DDPM 管道,其中低分辨率基础扩散模型与上采样器并行训练 [9, 20]。对于上面列出的所有扩散模型,卷积 U-Nets [49] 是主干架构的事实选择。并行工作 [24] 引入了一种基于 DDPM 注意力的新颖、高效的架构;我们探索了纯变压器。
架构复杂性。在评估图像生成文献中的架构复杂性时,使用参数计数是相当常见的做法。一般来说,参数计数可能是图像模型的复杂性较差的代理,因为它们没有考虑到图像分辨率,这显著影响性能[44,45]。相反,本文中的大部分模型复杂性分析是通过理论 Gflops 的镜头。这为我们提供了架构设计文献,其中 Gflops 被广泛用于衡量复杂性。在实践中,黄金复杂性度量仍然是辩论的,因为它经常取决于特定的应用场景。Nichol 和 Dhariwal 改进扩散模型的开创性工作 [9, 36] 与我们最相关——在那里,他们分析了 U-Net 架构类的可扩展性和 Gflop 属性。在本文中,我们专注于 Transformer 类。
3. Diffusion Transformers
3.1. Preliminaries
扩散公式。在介绍我们的架构之前,我们简要回顾了理解扩散模型 (DDPM) 所需的一些基本概念 [19, 54]。高斯扩散模型假设一个前向噪声过程,它逐渐将噪声应用于真实数据 x0: q(xt|x0) =N (xt; √ ̄αtx0, (1 − ̄αt)I),其中常数 ̄αt 是超参数。通过应用重新参数化技巧,我们可以采样 xt = √ ̄αtx0 + √1− ̄αt푡t,其中 witt ∼ N (0, I)。
扩散模型被训练来学习反转正向过程损坏的反向过程:pθ (xt−1|xt) =N (μθ (xt), Σθ (xt)),其中神经网络用于预测pθ的统计数据。反向过程模型使用x0对数似然的变分下界[30]进行训练,简化为L(θ) =−p(x0|x1) +∑t DKL(q * (xt−1|xt, x0)||pθ (xt−1|xt)),不包括与训练无关的附加项。由于 q∗ 和 pθ 都是高斯的,因此 DKL 可以使用两个分布的均值和协方差进行评估。通过将μθ重新参数化为噪声预测网络푡θ,可以使用预测噪声푡θ(Xt)和地面真实采样高斯噪声푝t之间的简单均方误差来训练模型:Lsimple(θ)=||푡θ(Xt)−푝t||2 2。但是,为了训练具有学习反向过程协方差Σθ的扩散模型,需要优化完整的DKL项。我们遵循 Nichol 和 Dhariwal 的方法 [36]:用 Lsimple 训练 ¢θ,并用完整的 L 训练 Σθ。一旦训练了 pθ,就可以通过重新参数化技巧初始化 xtmax ∼ N (0, I) 和采样 xt−1 ∼ pθ (xt−1|xt) 来采样新图像。
分类器自由指导。条件扩散模型以额外的信息作为输入,例如类标签 c。在这种情况下,反向过程变为 pθ (xt−1|xt, c),其中 푡θ 和 Σθ 以 c 为条件。在这种情况下,可以使用无分类器指导来鼓励采样过程找到 x,使得 log p(c|x) 很高 [21]。通过贝叶斯规则,log p(c|x) ∝ log p(x|c) − log p(x),因此 ∇x log p(c|x) ∝ ∇x log p(x|c)−∇x log p(x)。通过将扩散模型的输出解释为评分函数,可以将DDPM采样过程引导对p(x|c)高x进行采样:口令口令θ (xt, c) =口令θ (xt,∅)+ s·口令x log p(x|c)∝口令θ (xt,∅)+s·(口令θ (xt, c)−口令θ (xt,∅)),其中s > 1表示引导的比例(注意s = 1恢复标准采样)。使用 c = ∅ 评估扩散模型是通过在训练期间随机丢弃 c 并用学习的“空”嵌入∅ 替换它来完成的。众所周知,无分类器指导比通用采样技术 [21, 35, 46] 产生显着改进的样本,并且趋势适用于我们的 DiT 模型。
潜在扩散模型。直接在高分辨率像素空间中训练扩散模型在计算上可能是令人望而却步的。潜在扩散模型(ldm)[48]通过两阶段方法解决了这个问题:(1)学习一个自动编码器,用学习过的编码器E将图像压缩成更小的空间表示;(2)训练表示z = E(x)的扩散模型,而不是图像x (E被冻结)。然后可以通过从扩散模型中采样表示 z 然后将其解码为具有学习解码器 x = D(z) 的图像来生成新图像。
如图 2 所示,LDM 在使用 ADM 等像素空间扩散模型的一小部分 Gflop 的同时取得了良好的性能。由于我们关注计算效率,这使得它们成为架构探索的一个有吸引力的起点。在本文中,我们将 DiT 应用于潜在空间,尽管它们也可以应用于像素空间而不进行修改。这使得我们的图像生成管道成为一种基于混合的方法;我们使用现成的卷积 VAE 和基于变压器的 DDPM。
3.2. Diffusion Transformer Design Space
我们引入了扩散变压器(DiTs),这是一种扩散模型的新架构。我们的目标是尽可能忠实于标准的变压器架构,以保持其缩放特性。由于我们的重点是训练图像的ddpm(特别是图像的空间表示),DiT基于视觉转换器(ViT)架构,该架构对补丁序列[10]进行操作。DiT保留了vit的许多最佳实践。图 3 显示了完整的 DiT 架构的概述。在本节中,我们将描述 DiT 的前向传递以及 DiT 类的设计空间的组件。
修补。DiT 的输入是空间表示 z(对于 256 × 256 × 3 图像,z 的形状为 32 × 32 × 4)。DiT 的第一层是“patchify”,它通过线性嵌入输入中的每个补丁,将空间输入转换为 T 个标记的序列,每个维度 d。在 patchify 之后,我们将标准的基于 ViT 频率的位置嵌入(正弦余弦版本)应用于所有输入标记。由patchify创建的令牌T的数量由patch size超参数p决定。如图 4 所示,减半 p 将四元组 T ,因此至少四倍的总变压器 Gflop。尽管它对 Gflop 有显着影响,但请注意更改 p 对下游参数计数没有有意义的影响。我们将 p = 2, 4, 8 添加到 DiT 设计空间中。
DiT块设计。在 patchify 之后,输入标记由一系列转换器块处理。除了噪声图像输入之外,扩散模型有时会处理额外的条件信息,例如噪声时间步 t、类标签 c、自然语言等。我们探索了四种不同的变压器块变体,它们以不同的方式处理条件输入。该设计引入了对标准ViT块设计的微小但重要的修改。所有块的设计如图3所示。
- 上下文条件。我们只需将 t 和 c 的向量嵌入附加到输入序列中的两个额外的标记,将它们视为与图像标记没有区别。这类似于 ViT 中的 cls 标记,它允许我们使用标准 ViT 块而无需修改。在最后一个块之后,我们从序列中删除条件标记。这种方法会给模型带来微不足道的新 Gflop。
- 交叉注意块。我们将 t 和 c 的嵌入连接成一个长度为 2 的序列,与图像标记序列分开。修改变压器块,以包含一个额外的多头交叉注意层,遵循多头自注意块,类似于Vaswani等人[60]的原始设计,也类似于LDM用于调节类标签的模型。交叉注意将最多的 Gflops 添加到模型中,大约 15% 的开销。
- 自适应层范数(adaLN)块。遵循 GAN [2, 28] 中自适应归一化层 [40] 和具有 UNet 主干 [9] 的扩散模型的广泛使用,我们探索用自适应层范数 (adaLN) 替换变压器块中的标准层范数层。我们没有直接学习维度缩放和移位参数γ和β,而是从t和c的嵌入向量之和回归它们。在我们探索的三个块设计中,adaLN 添加了最少的 Gflop,因此计算效率最高。它还是唯一仅限于将相同函数应用于所有标记的条件机制。- adaLN-Zero块。先前关于 ResNet 的工作发现,将每个残差块初始化为恒等函数是有益的。例如,Goyal 等人。发现每个块中最终批次范数比例因子 γ 的零初始化加速了监督学习设置中的大规模训练 [13]。扩散U-Net模型使用类似的初始化策略,在任何剩余连接之前,在每个块中初始化最终的卷积层。我们探索了 adaLN DiT 块的修改,该块相同。除了回归 γ 和 β,我们还回归维度缩放参数 α,这些参数在 DiT 块内的任何残差连接之前立即应用。我们初始化 MLP 以输出所有 α 的零向量;这将完整的 DiT 块初始化为恒等函数。与 vanilla adaLN 块一样,adaLNZero 在模型中添加了可以忽略不计的 Gflop。我们在 DiT 设计空间中包含了上下文、交叉注意力、自适应层范数和 adaLN-Zero 块。
模型大小。我们应用了一系列 N DiT 块,每个块以隐藏维度大小 d 运行。在 ViT 之后,我们使用标准的 Transformer 配置来联合缩放 N、d 和注意力头 [10, 63]。具体来说,我们使用了四个配置:DiT-S、DiT-B、DiT-L 和 DiT-XL。它们涵盖了广泛的模型大小和触发器分配,从 0.3 到 118.6 Gflops,使我们能够衡量缩放性能。表 1 给出了配置的详细信息。我们将 B、S、L 和 XL 配置添加到 DiT 设计空间中。
Transformer 解码器。在最终的 DiT 块之后,我们需要将我们的图像标记序列解码为输出噪声预测和输出对角协方差预测。这两个输出的形状都等于原始空间输入。我们使用标准线性解码器来做到这一点;我们应用最后一层范数(如果使用 adaLN 自适应),并将每个标记线性解码为 p×p×2C 张量,其中 C 是空间输入中的通道数到 DiT。最后,我们将解码的令牌重新排列到其原始空间布局中,得到预测的噪声和协方差。我们探索的完整 DiT 设计空间是补丁大小、变压器块架构和模型大小。
4. Experimental Setup
我们探索了 DiT 设计空间并研究了我们的模型类的缩放属性。我们的模型根据它们的配置和潜在补丁大小 p 命名;例如,DiT-XL/2 指的是 XLarge 配置和 p = 2。
训练。我们在ImageNet数据集[31]上训练256 × 256和512 × 512图像分辨率的类条件潜在DiT模型,这是一个高度竞争的生成建模基准。我们用零初始化最终的线性层,否则使用来自ViT的标准权重初始化技术。我们使用 AdamW [29,33] 训练所有模型。我们使用1 × 10−4的恒定学习率,没有权值衰减,批大小为256。我们使用的唯一数据增强是水平翻转。与 ViT [57, 61] 的许多先前工作不同,我们没有发现训练 DiT 实现高性能所需的学习率预热或正则化。即使没有这些技术,在所有模型配置中训练都非常稳定,我们没有观察到训练转换器时常见的任何损失峰值。在生成建模文献中的常见做法之后,我们在衰减为 0.9999 的情况下在训练期间保持 DiT 权重的指数移动平均 (EMA)。所有报告的结果都使用 EMA 模型。我们在所有 DiT 模型大小和补丁大小中使用相同的训练超参数。我们的训练超参数几乎完全保留在 ADM 中。我们没有调整学习率、衰减/预热计划、Adam β1/β2 或权重衰减。
扩散。我们使用来自稳定扩散[48]的现成预训练变分自编码器(VAE)模型[30]。VAE 编码器的下采样因子为 8——给定一个形状为 256 × 256 × 3 的 RGB 图像 x,z = E(x) 的形状为 32 × 32 × 4。在本节中的所有实验中,我们的扩散模型在这个 Z 空间中运行。在从我们的扩散模型中采样一个新的潜在后,我们使用 VAE 解码器 x = D(z) 将其解码为像素。我们从 ADM [9] 保留扩散超参数;具体来说,我们使用 atmax = 1000 个线性方差调度,范围从 1×10−4 到 2 × 10−2,ADM 的协方差参数化 Σθ 及其嵌入输入时间步长和标签的方法。
评估指标。我们使用 Fr ́echet Inception Distance (FID) [18] 来衡量缩放性能,这是评估图像生成模型的标准指标。我们遵循与之前工作的比较时的惯例,并使用 250 个 DDPM 采样步骤报告 FID-50K。众所周知,FID对小的实现细节[37]很敏感;为了确保准确的比较,本文报告的所有值都是通过导出样本并使用ADM的TensorFlow评估套件[9]获得的。本节中报告的 FID 数字不使用无分类器指导,除非另有说明。我们还报告了 Inception Score [51]、sFID [34] 和 Precision/Recall [32] 作为次要指标。
计算。我们在 JAX [1] 中实现所有模型,并使用 TPU-v3 pod 对其进行训练。DiT-XL/2 是我们最计算密集型模型,在 TPU v3-256 pod 上训练大约 5.7 次迭代/秒,全局批量大小为 256。
5. Experiments
DiT块设计。我们训练了四个最高的 Gflop DiT-XL/2 模型,每个模型都使用不同的块设计——上下文(119.4 Gflops)、交叉注意(137.6 Gflops)、自适应层范数(adaLN,118.6 Gflops)或 adaLN-zero(118.6 Gflops)。我们在训练过程中测量 FID。图 5 显示了结果。adaLN-Zero 块产生的 FID 低于交叉注意力和上下文条件,同时是最有效的。在 400K 训练迭代中,adaLN-Zero 模型实现的 FID 几乎是上下文模型的一半,这表明条件机制严重影响了模型质量。初始化也很重要——adaLNZero,它将每个 DiT 块初始化为恒等函数,显着优于 vanilla adaLN。对于本文的其余部分,所有模型都使用 adaLN-Zero DiT 块。
缩放模型大小和补丁大小。我们训练了 12 个 DiT 模型,扫描模型配置 (S, B, L, XL) 和补丁大小 (8, 4, 2)。请注意,与其他配置相比,DiT-L 和 DiT-XL 在相对 Gflop 方面彼此明显更接近。图 2(左)概述了每个模型 Gflops 及其在 400K 训练迭代中的 FID。在所有情况下,我们发现增加模型大小和减少补丁大小会产生显着改善的扩散模型。
图 6(顶部)展示了随着模型大小的增加,FID 如何变化,补丁大小保持不变。在所有四个配置中,通过使转换器更深更广,在所有训练阶段都获得了 FID 的显着改进。类似地,图 6(底部)显示了 FID,因为补丁大小减少,模型大小保持不变。我们再次通过简单地缩放 DiT 处理的令牌数量来在整个训练过程中观察到相当大的 FID 改进,保持参数近似固定。
DiT Gflops 对提高性能至关重要。图 6 的结果表明,参数计数并不能唯一地确定 DiT 模型的质量。由于模型大小保持不变,补丁大小减小,变压器的总参数有效不变(实际上,总参数略有下降),只有 Gflops 增加。这些结果表明,缩放模型 Gflopsis 实际上是提高性能的关键。为了进一步研究这一点,我们在图 8 中针对模型 Gflops 在 400K 训练步骤中绘制了 FID-50K。结果表明,当总 Gflops 相似时,不同的 DiT 配置获得了相似的 FID 值(例如,DiT-S/2 和 DiT-B/4)。我们发现模型 Gflops 和 FID-50K 之间存在很强的负相关,这表明额外的模型计算是提高 DiT 模型的关键因素。在图 12(附录)中,我们发现这种趋势适用于其他指标,例如 Inception 分数。
更大的 DiT 模型的计算效率更高。在图 9 中,我们将 FID 绘制为所有 DiT 模型的总训练计算函数。我们将训练计算估计为模型 Gflops · batch size · 训练步骤 · 3,其中 3 因子粗略地将后向传递近似为计算量的两倍作为前向传递。我们发现,与为更少步骤训练的更大的 DiT 模型相比,小型 DiT 模型(即使训练时间更长)最终也会变得计算效率。类似地,我们发现除了补丁大小之外相同的模型即使在控制训练 Gflop 时也具有不同的性能配置文件。例如,在大约 1010 Gflops 之后,XL/4 优于 XL/2。
可视化缩放。我们在图 7 中可视化了缩放对样本质量的影响。在 400K 训练步骤中,我们使用相同的起始噪声 xtmax、采样噪声和类标签从 12 个 DiT 模型中的每一个中采样图像。这让我们可以直观地解释缩放如何影响 DiT 样本质量。事实上,缩放模型大小和令牌的数量可以显着提高视觉质量。
5.1. State-of-the-Art Diffusion Models
256×256 ImageNet。按照我们的缩放分析,我们继续训练我们最高的 Gflop 模型 DiT-XL/2,用于 7M 步。我们在图 1 中展示了来自模型的样本,并与最先进的类条件生成模型进行了比较。我们在表 2 中报告了结果。当使用无分类器指导时,DiT-XL/2 优于所有先前的扩散模型,将 LDM 实现的先前最佳 FID-50K 降低到 2.27。图 2(右)显示 DiT-XL/2(118.6 Gflops)相对于潜在空间 U-Net 模型(如 LDM-4(103.6 Gflops)的计算效率很高,并且比像素空间 U-Net 模型(如 ADM(1120 Gflops)或 ADM-U(742 Gflops)更有效。我们的方法实现了所有先前生成模型的最低 FID,包括之前最先进的 StyleGANXL [53]。最后,我们还观察到,与LDM-4和LDM-8相比,DiT-XL/2在所有测试的无分类器引导尺度上实现了更高的召回值。当只训练了2.35M步(类似于ADM)时,XL/2仍然优于所有先前的扩散模型,FID为2.55.
512×512 ImageNet。我们在 ImageNet 上训练一个新的 DiT-XL/2 模型,分辨率为 512 × 512,用于 3M 次迭代,超参数与 256 × 256 模型相同。补丁大小为 2,此 XL/2 模型在修补 64 × 64 × 4 输入潜在(524.6 Gflops)后总共处理 1024 个标记。表 3 显示了与最先进方法的比较。XL/2 在这种分辨率下再次优于所有先前的扩散模型,将 ADM 实现的先前最佳 FID 提高了 3.85 到 3.04。即使令牌数量的增加,XL/2 仍然计算效率很高。例如,ADM 使用 1983 Gflops 和 ADM-U 使用 2813 Gflops; XL/2 使用 524.6 Gflops。我们在图 1 和图附录中展示了高分辨率 XL/2 模型的样本。
5.2. Scaling Model vs. Sampling Compute
扩散模型是独一无二的,因为它们可以通过在生成图像时增加采样步骤的数量来在训练期间使用额外的计算。鉴于模型 Gflops 对样本质量的影响,在本节中,我们研究通过使用更多的采样计算,较小的模型计算 DiT 是否可以胜过更大的模型。我们在 400K 训练步骤后计算所有 12 个 DiT 模型的 FID,每张图像使用 [16, 32, 64, 128, 256, 1000] 个采样步骤。主要结果如图10所示。考虑使用1000个采样步骤的DiT-L/2和使用128步的DiT-XL/2。在这种情况下,L/2 使用 80.7 Tflops 对每张图像进行采样; XL/2 使用 5 倍的计算量——15.2 Tflops——对每张图像进行采样。尽管如此,XL/2 具有更好的 FID-10K(23.7 对 25.9)。一般来说,由于缺乏模型计算,放大采样计算无法补偿。
6. Conclusion
我们引入了扩散变压器(DiTs),这是一种简单的基于变压器的扩散模型骨干,其性能优于之前的U-Net模型,继承了变压器模型类的出色缩放特性。鉴于本文有希望的缩放结果,未来的工作应该继续将 DiT 扩展到更大的模型和令牌计数。DiT 也可以探索为文本到图像模型(如 DALL·E 2 和 Stable Diffusion)的下降主干。
A. Additional Implementation Details
我们在表 4 中包含有关我们所有 DiT 模型的详细信息,包括 256 × 256 和 512 × 512 模型。在图 13 中,我们报告了 DiT 训练损失曲线。最后,我们还在表6中包含了来自ADM和LDM的DDPM U-Net模型的Gflop计数。
**DiT 模型细节。**为了嵌入输入时间步长,我们使用 256 维频率嵌入 [9],然后是一个维度为等于转换器隐藏大小和 SiLU 激活的两层 MLP。每个 adaLN 层将时间步和类嵌入的总和馈送到 SiLU 非线性和输出层,输出神经元等于 4× (adaLN) 或 6× (adaLN-Zero) 变压器隐藏大小。我们在核心变压器[16]中使用GELU非线性(用tanh近似)。
通道子集上的无分类器指导。在我们使用无分类器指导的实验中,我们仅将指导应用于潜在变量的前三个通道,而不是所有四个通道。在调查后,我们发现三通道引导和四通道引导在简单地调整比例因子时给出了相似的结果(就 FID 而言)。具体来说,以 (1 + x) 尺度为 (1 + 3 4 x) 的四通道引导似乎相当好地近似为(例如,尺度为 1.5 的三通道引导给出了 2.27 的 FID-50K,尺度为 1.375 的四通道引导给出了 2.20 的 FID-50K。有趣的是,将指导应用于元素子集仍然可以产生良好的性能,我们将其留给未来的工作来进一步探索这种现象。
B. Model Samples
我们展示了两个 DiT-XL/2 模型在 512 × 512 和 256 × 256 分辨率下的样本,分别训练了 3M 和 7M 步。图 1 和图 11 显示了两个模型中选择的样本。图 14 到 33 显示了两个模型在一系列无分类器指导尺度和输入类标签(使用 250 个 DDPM 采样步骤和 ft-EMA VAE 解码器生成)中的未整理样本。与之前使用指导的工作一样,我们观察到更大的尺度增加了视觉保真度并降低样本多样性。
C. Additional Scaling Results
缩放对FID以外的指标的影响。在图 12 中,我们展示了 DiT 量表对一套评估指标的影响——FID、sFID、Inception Score、Precision 和 Recall。我们发现,我们在主要论文中的 FID 驱动分析可以推广到其他指标——在每个指标上,放大的 DiT 模型的计算效率更高,模型 Gflop 与性能高度相关。特别是,Inception 分数和精度从模型规模的增加中受益匪浅。
缩放对训练损失的影响。我们还在图 13 中检查了规模对训练损失的影响。增加 DiT 模型 Gflops(通过 Transformer 大小或输入标记的数量)会导致训练损失更快地减少并在较低的值下饱和。这种现象与语言模型观察到的趋势一致,其中放大转换器展示了改进的损失曲线以及下游评估套件的性能 [26]。
D. VAE Decoder Ablations
在我们的实验中,我们使用了现成的、预训练的 VAE。VAE 模型(ft-MSE 和 ft-EMA)是原始 LDM“f8”模型(仅微调解码器权重)的微调版本。我们使用ft-MSE解码器监控第5节中缩放分析的指标,并使用ft-EMA解码器进行表2和表3中报告的最终指标。在本节中,我们消融了三个VAE解码器的不同选择;LDM使用的原始解码器和稳定扩散使用的两个微调解码器。由于编码器在模型之间是相同的,因此可以在不重新训练扩散模型的情况下交换解码器。表 5 显示了结果;当使用 LDM 解码器时,XL/2 继续优于所有先前的扩散模型。