聊聊 DiT 和 GenTron

1.介绍

随着 OpenAI 最新 Sora 的火爆出圈，Diffusion Transformer 无疑将成为继 Stable Diffusion 后最出圈的模型结构。Diffusion Transformer Model 是一种结合了 Transformer 架构和扩散模型的图像、视频及其他数据的生成模型。也就是说 Diffusion Transformer 继承了 Transformer 这种具有 scaling laws 优势的模型结构，具有知识随参数增加带来迸发效果的先天优势。

DiT模型：基于 transformer 架构的扩散模型，不仅将transformer成功应用在扩散模型，还探究了 transformer 架构在扩散模型上的 scalability 能力。
GenTron模型：基于Transformer 的扩散模型，在针对SDXL的人类评估中，GenTron在视觉质量方面取得了51.1%的胜率（19.8%的平局率），在文本对齐方面取得了42.3%的胜率（42.9%的平局率）。

2. DiT

论文：https://https://arxiv.org/pdf/2212.09748.pdf8.pdf

代码：https://github.com/facebookresearch/DiT

DiT 基于 Vision Transformer（ViT）架构，该架构操作于图像的 patch 序列。DiT 保留了 ViTs 的许多最佳实践。下图显示了完整的 DiT 架构概述，描述了 DiT 的前向传播，以及 DiT 类的设计空间组件。

图1

2.1 Patchify

DiT 的第一层是 “patchify”，它通过卷积和线性操作将输入的图像 patch 化，BCHW -> BNC。在 patchify 之后，对所有输入标记应用标准 ViT 的基于频率的位置编码（正弦-余弦版本）。如下图所示。

图2

2.2 Timestep t

Timestep t 是一个时间维度的 embedding，在 Timestep 中会根据预制的 frequency_embedding_size，对于预制的时间步数进行相对位置编码。再将相对位置编码后的 embedding 映射到隐藏层维度。

2.3 Label y

Label y 类似于 NLP 模型的词表映射，因为 DiT 模型只有 1000 类，所以 num_classes=1000，在实际代码中，如下代码发现还会增加 class_dropout_prob 类别。（这部分对于类别不清晰的部分做归纳）

self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)

Timestep t （N, D）和 Label y （N, D）的结果会进行相加得到 conditioning。

2.4 DiTBlock

conditioning 经过线性层生成，，，，，γ1，β1，α1，γ2，β2，α2 。如下代码所示。

shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)

如上图1所示，经过 depth 层 DiTBlock 迭代，(N, T, D) -> (N, T, D)。

2.5 FinalLayer

同上 DiTBlock 有 adaLN_modulation 模块，将 adaLN_modulation 的结果经过线性层得到最后的输出结果。(N, T, D) -> (N, T, patch_size 2 * out_channels)

shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)

2.6 unpatchify

将 FinalLayer 的结果 unpatchify （此处使用 reshape 和 einsum 做转换）到最后的输出噪声。

(N, T, patch_size 2 * out_channels) -> (N, out_channels, H, W)

2.7 vae 解码

通过 vae 解码得到最后的输出图片。

图3

2.8 实验

配置：

256 × 256 and 512 × 512 image resolution
ImageNet dataset
Initialize the final linear layer with zeros and otherwise use standard weight initialization techniques from ViT
AdamW
Learning rate of 1 × 10−4
No weight decay and a batch size of 256
Horizontal flips
Moving average (EMA): 0.9999

Scaling the DiT model improves FID at all stages of training

图4

Transformer Gflops are strongly correlated with FID

图5

Larger DiT models use large compute more efficiently

图6

Benchmarking class-conditional image generation on ImageNet 256×256

图7

3. GenTron

论文：https://arxiv.org/pdf/2312.04557.pdf7.pdf

代码：暂未找到

GenTron 是继 DiT 后的 Transformer扩散生成模型。将 Diffusion Transformers 扩展到文本输入，增加 cross attention 机制。GenTron 将参数规模从约 900M 扩展到超过 3B，生成图片质量显著提升。此外，GenTron 从文本扩展到到视频生成，引入了无运动引导以提高视频质量。GenTron 在视觉质量方面达到了 51.1%的胜率（平局率为 19.8%），在文本对齐方面达到了 42.3%的胜率（平局率为 42.9%）。

主要特点：

基于 Transformer 的T2I（文本到图像）生成与扩散模型进行了彻底而系统的调查。
探索了一种纯基于 Transformer 的扩散模型，用于T2V（文本到视频）生成。引入了无动作引导（motion-free guidance），这是一种创新技术，可以有效地对 T2I 生成模型进行微调，以产生高质量的视频。
实验结果表明，在人类评估中，GenTron 相比于SDXL有明显的优势。此外，GenTron在 T2I-CompBench 评估中相比现有方法展现出了更优秀的性能。

3.1 主体结构

图8

GenTron基于 DiT-XL/2 构建，它将形状为 32×32×4 的 Latent 特征转换为非重叠 tokens 序列，采用 2×2 的裁剪层。然后，这些 tokens 被送入一系列的 transformer 块。最后，应用标准的线性解码器将这些图像 tokens 转换为 Latent 空间。

文本的 cross attention
文本 pooling 后与 timestamp进行相加

3.2 Embedding integration

自适应层归一化（Adaptive layernorm (adaLN)）
主体和 DiT 的 adaLN 类似
跨注意力（Cross-attention）

图像特征充当 query，而文本 embedding 则充当 key 和 value。该设置允许通过注意力机制在图像特征和文本 embedding 之间进行直接交互

3.3 Scaling Up GenTron

为了探索大幅扩展模型规模的影响，作者开发了 GenTron 的高级版本，称为 GenTron-G/2。该模型扩展了三个关键方面：transformer块的数量（深度）、patch embeddings 的维度（宽度）和 MLP 的隐藏维度（MLP 宽度）。GenTron 模型的规格和配置详见下表1。

表1

3.4 GenTron-T2V结构

本文方法仅在每个 transformer 块中集成轻量级的时间自注意力（TempSelfAttn）层。如下图 9 所示，TempSelfAttn 层位于交叉注意力层之后和 MLP 层之前。

图9

作者通过在进入 TempSelfAttn 层之前更改交 cross attention 层的输出，然后在通过后将其重新改回其原始格式，修改了交叉注意力层的输出。这个过程可以正式表示为：

图10

其中，b、t、n、d 分别表示batch大小、帧数、每帧的patch数量和通道维度。

3.4.1 初始化

使用预训练的 T2I 模型作为初始化 T2I 和 T2V 模型之间共享层的基础。此外，对于新增的 TempSelfAttn 层，将输出 project 层的权重和偏差初始化为零。这确保在 T2V 微调阶段的开始，这些层产生零输出，与 shortcut 连接一起有效地作为恒等映射。

3.5 无运动引导（Motion-Free Guidance）

3.5.1 问题和挑战

作者观察到当前 T2V 扩散模型中存在一个显著现象，即每帧的视觉质量明显滞后于 T2I 模型。此外，分析还显示，在 T2V 模型经过微调后，其视觉质量明显下降，特别是与其原始的 T2I 对应物相比。注意到了这些问题通常存在于当前的 T2V 扩散模型中，不仅限于基于transformer的 T2V。

3.5.2 问题分析

作者推测 T2V 视觉质量滞后的现象主要源于两个因素：

公开可用的视频数据集在质量和数量上通常不及图像数据集，许多视频帧受到运动模糊和水印的影响，进一步降低了它们的视觉质量。
在视频微调期间专注于优化时间方面可能会无意中损害空间视觉质量，导致生成的视频总体质量下降。

3.5.3 解决方案

从数据方面考虑，采用联合图像-视频训练策略来缓解视频数据不足的问题。
将视频片段内的时间运动视为一种特殊的条件信号，类似于 T2I/T2V 扩散模型中的文本条件。基于这个类比，并提出了一种新方法，即无运动引导（MFG），灵感来自于无分类器引导，用于条件生成视频中运动信息的权重。作者使用一个单位矩阵以概率 pmotion free 抵消时间注意力。这个单位矩阵在图 9（无运动掩码）中呈现，其对角线填充为 1，而所有其他位置为 0。这种配置将时间自注意力限制在单个帧内。

在推理过程中，有文本和运动的条件。可以修改分数估计为：

图11

可以将解决方案 I 和 II 结合在一起，具体方法如下：当在训练步骤中省略运动时，加载一对图像文本，并将图像重复 T-1 次以创建一个伪视频。相反，如果包含运动，则加载视频片段并将其提取成 T 帧。

3.6 实验

配置：

AdamW
Learning rate of 1×10−4
Internal dataset, including a low-resolution (256×256) training with a batch size of 2048 and 500K optimization steps
High-resolution (512×512) with a batch size of 784 and 300K steps
Fully Sharded Data Parallel (FSDP) [68] and activation checkpointing (AC)
Video dataset: 34M videos
A short side of 512 pixels and a frame rate of 24 FPS
Batches of 128 video clips. Each clip comprises 8 frames, captured at a sampling rate of 4 FPS

Cross attention vs. adaLN-Zero

图12