0
点赞
收藏
分享

微信扫一扫

Java虚拟机 - JVM

(DiTs)Scalable Diffusion Models with Transformers

文章目录

引用: [1] Peebles W, Xie S. Scalable diffusion models with transformers[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023: 4195-4205.

论文链接: (ICCV2023) https://arxiv.org/abs/2212.09748

代码链接: https://github.com/facebookresearch/DiT

论文概述

Ho等[1]的工作首先提出了用于扩散模型的U-Net骨架。U-Net最初在像素级自回归模型和条件GANs[2]中取得了成功,它继承自PixelCNN++[3],并进行了一些更改。该模型是卷积模型,主要由 ResNet 模块组成。与标准U-Net相比,空间自注意力模块以较低的分辨率散布。然而,Ho等人对UNet的高级设计基本上保持不变。

因此,伯克利的William Peebles提出了一类新的基于transformer架构的扩散模型,即DiTs。DiTs将常用的U-Net主干替换为对latent patches进行操作的transformer。论文证明了U-Net的归纳偏置对扩散模型的性能并不重要。实验发现,具有较高Gflop的DiTs——通过增加深度/宽度或增加输入tokens的数量——始终具有较低的FID。除了具有良好的可扩展性外,largest DiT XL/2模型在类条件 ImageNet 512上的性能优于所有先前的扩散模型512×512和256×256的基准,在后者上实现了2.27的最先进FID。

Diffusion Transformers

高斯扩散模型假设了一个前向噪声处理过程,该过程逐渐将噪声应用于真实数据。通过应用重新参数化技巧,可以对带噪声数据 x t x_t xt进行采样。扩散模型经过训练,可以学习反转正向进程损坏的反向过程。反向过程模型是用 x 0 x_0 x0的对数似然的变分下界[4]训练的.

条件扩散模型将额外信息作为输入,例如类标签 c。用 c = NULL 评估扩散模型,是通过在训练过程中随机丢弃 C 并将其替换为学习的“空”嵌入来完成的;众所周知,与通用采样技术相比,无分类器引导可以显著改善样本,并且DiT模型也保持了这一趋势。

潜在扩散模型(LDM)通过两阶段方法解决了这个问题:

(1)学习一个自动编码器,该自动编码器使用学习编码器E将图像压缩为更小的空间表示;

(2)训练表示z = E(x)的扩散模型,而不是图像x的扩散模型(E被冻结)。然后,可以通过从扩散模型中采样表示 z 并随后使用学习的解码器 x = D(z) 将其解码为图像来生成新图像。

如图 2 所示,LDM 在使用 ADM 等像素空间扩散模型的一小部分 Gflops 时实现了良好的性能。由于我们关注计算效率,这使它们成为架构探索的一个有吸引力的起点。在本文中,将 DiT 应用于潜在空间,尽管它们也可以应用于像素空间而无需修改。这使得图像生成管道成为一种基于混合的方法;使用现成的卷积 VAE 和基于 transformer 的 DDPM。

由于重点是训练图像的DDPM(特别是图像的空间表示),因此DiT基于ViT架构,该架构对补丁序列进行操作。DiT 保留了 ViT 的许多最佳实践。图 3 显示了完整的 DiT 架构的概述。

DiT 的输入是空间表示 z(对于 256×256×3 图像,z 的形状为 32×32×4)。DiT 的第一层是“patchify”,它通过线性地将每个 Patch 嵌入到输入中,将空间输入转换为一系列 T tokens,每个标记的维度为 d。在patchify之后,我们将标准的**基于 ViT 频率的位置嵌入(正余弦版本)**应用于所有输入tokens。patchify 创建的 tokens T 数量由补丁大小超参数 p 决定。如图 4 所示,将 p 减半将使 T 翻倍,因此Transformer总 Gflops 至少翻了四倍。尽管它对 Gflops 有重大影响,但请注意,更改 p 对下游参数计数没有有意义的影响。

在 patchify 之后,输入tokens由一系列transformer块处理。除了噪声图像输入外,扩散模型有时还会处理其他条件信息,例如噪声时间步长 t类标签 c自然语言等。论文总共探索了四种transformer模块的变体,它们以不同的方式处理条件输入。这些设计对标准 ViT 模块设计进行了微小但重要的修改,所有模块的设计如图 3 所示。

  • In-context conditioning:DiT只需将 t 和 c 的向量嵌入作为输入序列中的两个附加标记。这类似于 ViT 中的 cls,它允许模型无需修改即可使用标准 ViT 块。在最后一个块之后,从序列中删除条件标记。这种方法为模型引入了可忽略不计的新 Gflops。

  • Cross-attention block将 t 和 c 的嵌入连接成一个长度为二的序列,与图像标记序列分开。对变压器模块进行了修改,在多头自注意力模块之后增加了一个多头交叉注意力层。交叉注意力为模型增加了最多的 Gflops,大约是 15% 的开销。

  • Adaptive layer norm (adaLN) block:继自适应归一化层在GANs和具有UNet backbone的扩散模型的广泛使用之后,论文探索了用自适应层范数(adaLN)替换transformer模块中的标准层范数层,而不是直接学习尺度和位移参数。adaLN 添加的 Gflops 最少,因此计算效率最高。

  • adaLN-Zero block:之前对 ResNets 的研究发现,将每个残差块初始化为恒等函数是有益的。论文探索了对 adaLN DiT 模块的修改,使得它也有同样的作用。除了回归 γ \gamma γ β \beta β,模型还回归了在DiT块内任何剩余连接之前立即应用的按维度缩放参数 α \alpha α

DiT应用 N 个 DiT 块,每个块都以隐藏维度大小 d 运行。在ViT之后,使用标准transformer配置,这些配置可以联合缩放N、d和注意力头。具体来说,使用四种配置:DiT-S、DiT-B、DiT-L 和 DiT-XL。涵盖了从 0.3 到 118.6 Gflops 的各种模型大小和浮点分配,用来衡量扩展性能。表 1 提供了配置的详细信息。

实验

ImageNet 数据集上以 256×256 和 512×512 的图像分辨率训练类条件潜在 DiT 模型。用零初始化最终的线性层,其他模块则使用ViT的标准权重初始化技术。使用0.0001的恒定学习率没有权重衰减,批大小为256。使用的唯一数据扩充是水平翻转。在训练过程中保持DiT权重的指数移动平均值(EMA),其衰减为0.9999**。**其余超参数都未调整。使用 Frechet 初始距离 (FID)来衡量缩放性能,这是评估图像生成模型的标准指标。在与以前的工作进行比较时,遵循惯例,使用 250 个 DDPM 采样步骤报告 FID-50K。

参考文献

[1] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In NeurIPS, 2020. 2, 3.

[2] Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei AEfros. Image-to-image translation with conditional adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1125–1134, 2017. 2.

[3] Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik PKingma. PixelCNN++: Improving the pixelcnn with discretized logistic mixture likelihood and other modifications.arXiv preprint arXiv:1701.05517, 2017. 2.

[4] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013. 3, 6.

举报

相关推荐

java虚拟机(JVM)

java虚拟机 (JVM)

Java虚拟机(JVM)

JVM(java虚拟机)

JVM虚拟机

jvm虚拟机

JVM虚拟机与Android虚拟机

0 条评论