Sora = Diffusion + Transformer，爆火的背后是如何节约计算成本！

掘金安东尼

发布于 2024-02-27 08:55:46

6250

发布于 2024-02-27 08:55:46

Sora，说是 2024 年以来最备受瞩目的生成式模型，一点不为过，它的生成视频的能力、效果，令人咋舌。

许多文章对 Sora 进行了应用层、展示层的介绍，而对其背后原理的解读还不够。

本篇一起来揭开 Sora 背后的神秘面纱，接触它的核心技术。

直白来讲，Sora 是一个扩散模型，基于 DiT 技术构建。

就像我们知道 ChatGPT 是基于 Transformer 模型，一种基于自注意力机制的深度学习模型。

Sora 的 DiT 展开就是：Diffusion Transformer，Sora = Diffusion + Transformer，这种更进一步的扩散模型，不仅能够实现媲美 GAN 的图像生成质量，而且具有更好的扩展性和计算效率。

如果之前使用过、了解过 stable diffusion 的同学，对 Diffusion model 应该会有印象：

Diffusion model 是基于噪声的生成模型，通过逐步添加噪声来模拟数据分布，然后学习逆向过程去除噪声，从而生成新的数据。

在 Sora 中，DiT 能够合成高质量的图像、并且可以通过简单修改，还保证其较低的计算成本。

现在算力这么贵的背景下，计算成本真的太重要了~

一图胜千言：

DiT 包括 VAE 编码器、ViT、DDPM 和 VAE 解码器，这些组件：

VAE 编码器：VAE 是一种生成模型，用于将输入数据映射到潜在空间中的低维表示。
- 在 DiT 中的作用： DiT 使用 VAE 编码器将图像压缩到低维表示，以便在低维空间中训练 DDPM 模型，这有助于降低计算成本并提高效率。
ViT：ViT 是一种基于 Transformer 的图像分类模型，它将图像划分为均匀的图块，并使用自注意力机制来处理这些图块。
- 在 DiT 中的作用：DiT 借鉴了 ViT 的思想，将图像切分成多个大小为 p × p 的补丁，然后将其转换为长度为 T 的序列作为 Transformer 的输入。这使得 DiT 能够处理不同分辨率、持续时间和长宽比的视频和图像。

DDPM：DDPM 是一类深度生成模型，通过逐步添加噪声来模拟数据分布，然后学习逆向过程去除噪声，从而生成新的数据。
- 在 DiT 中的作用：DiT 基于 DDPM 技术构建，通过逐步添加噪声来模拟图像数据分布，然后学习逆向过程去除噪声，从而生成高质量的图像。
VAE 解码器：VAE 解码器将潜在空间中的低维表示映射回原始数据空间，从而生成重构的图像。
- 在 DiT 中的作用：DiT 使用 VAE 解码器将低维表示映射回图像空间，以生成最终的高质量图像。

这些组件共同构成了 DiT 模型，使其能够根据文本指令生成逼真且富有想象力的场景。

Patchify 技术也是不可绕开的关键词，是视频高质量的关键！

Patchify 是一种将图像切分成多个大小为 p × p 的补丁，并将其转换为长度为 T 的序列作为 Transformer 的输入的方法。

Sora 使用 Patchify 将图像分割成小块，然后将这些块转换为序列，以便输入到 Transformer 模型中。这样做的好处是，Sora 可以处理不同分辨率、持续时间和长宽比的视频和图像。

实验证明：训练模型规模越大、patch size (p) 越小 (即增加transformer的深度/宽度或增加输入token的数量)——可以显著提高视觉保真度

基于 patch 的表示方式使 Sora 能够训练不同分辨率、持续时间和高宽比的视频和图像。我们可以通过在适当大小的网格中安排随机初始化的 patches 来控制生成的视频的大小。

在 patchify 之后，输入 tokens 将由一系列 transformer blocks 进行处理。除了噪声图像输入，扩散模型有时还会处理附加的条件信息，如噪声时间步长 t、类标签 c、自然语言等。

DiT block 包含自注意力层、层规范层和前馈网络层。其中，有四种变体，如 in-context conditioning、cross-attention、adaptive layer norm (adaLN) 和 adaLN-Zero。

DiT（Diffusion Transformer） 的四种变体在实现两个额外嵌入的方式上有所不同：

In-context conditioning：将两个嵌入视为两个标记，合并到输入的标记中。类似于 ViT 中的 cls 标记。实现简单，几乎不引入额外的计算量。
Cross-attention block：将两个嵌入拼接成一个长度为 2 的序列，然后在 Transformer 块中插入一个交叉注意力层。条件嵌入作为交叉注意力层的键和值。这是目前文生成图像模型所采用的方式，但需要额外引入约 15% 的计算量。
Adaptive layer norm (adaLN) ：使用自适应层规范化（adaLN）来融合两个嵌入。adaLN 可以根据输入的不同样本自动调整规范化参数。
AdaLN-Zero：是 adaLN 的一种变体。与 adaLN 类似，但在计算规范化参数时，将其中一个嵌入的均值设为零。这种方式在保持性能的同时减少了计算成本。

这些变体在处理两个额外嵌入时具有不同的权衡，开发者可以根据具体任务和性能需求选择适合的变体。

transformer block 根据结构 (4种结构) 有 4 种模型大小 (S,B,L,XL），分别为 DiT-S、DiT-B、DiT-L 和 DiT-XL。

这四种配置在层数、隐藏维度和计算量上有所不同。

有能力的同学可基于以上开展训练：

在 ImageNet 数据集上训练类条件的 DiT 模型，分辨率为 256 × 256 和 512 × 512。使用 AdamW 优化器，学习率为 1 × 10^-4，权重衰减为 0，批量大小为 256。
扩散：使用预训练的变分自编码器 (VAE) 将图像编码为低维表示，然后在 VAE 的潜在空间中训练 DiT 模型。