摘要

基于transformer结构构建了一类新的扩散模型，DiT。即使增加transformer的深度/宽度或者增加输入tokens的数量，依旧能够保持更低的FID。最大的DiT-XL/2模型在ImageNet 512×512和256×256基准上性能优于所有先前的扩散模型，在后者上实现了2.27的FID。

方法

预备知识

1.扩散模型

前向加噪过程：对真实数据逐步添加噪声：

$x_{0}: q(x_t|x_0)={\cal N} ( x_{t} ; \sqrt{\bar{\alpha}_{t}} x_{0}, ( 1-\bar{\alpha}_{t} ) {\bf I} )$ $x_{t}=\sqrt{\bar{\alpha}_{t}} x_{0}+\sqrt{1-\bar{\alpha}_{t}} \epsilon_{t}, \quad where\quad\epsilon_{t} \sim{\mathcal{N}} ( 0, \mathbf{I} )$

对扩散模型进行训练，学习正向加噪过程的反向过程：

$p_{\theta}\left(x_{t-1} \mid x_{t}\right)=\mathcal{N}\left(\mu_{\theta}\left(x_{t}\right), \Sigma_{\theta}\left(x_{t}\right)\right)$

反向过程模型通过 $x_{0}$ 对数似然的变分下界进行训练：

$\mathcal{L}(\theta)=-p(x_0|x_1)+\sum_t\mathcal{D}_{KL}(q^*(x_{t-1}|x_t,x_0)||p_\theta(x_{t-1}|x_t))$

模型训练的损失函数为：

$\mathcal{L}_{simple}(\theta)=||\epsilon_\theta(x_t)-\epsilon_t||_2^2$

用 $\mathcal{L}_{simple}$ 训练 $\epsilon_{\theta}$ ，用完整的 $\mathcal{L}$ 训练 $\sum_{\theta}$ 。只要 $p_{\theta}$ 训练确定了，新的图像就能够通过初始化 $x_{t_{\mathrm{max}}}\sim\mathcal{N}(0,\mathbf{I})$ 采样得到，并采样 $x_{t-1}\sim p_\theta(x_{t-1}|x_t)$ 。

2.无分类器引导

条件扩散模型采用额外的信息作为输入，例如一个分类标签 $c$ 。在这种情况下，反向过程变成了 $p_\theta(x_{t-1} | x_t, c)$ 。在这种设定下，无分类引导能够用来鼓励采样过程找到一个 $x$ 使得 $\log p(c|x)$ 很大。根据贝叶斯规则，$\log p(c|x)\propto\log p(x|c)-\log p(x)$ ，因此 $\nabla_x\log p(c|x)\propto\nabla_x\log p(x|c)-\nabla_x\log p(x)$ 。通过将扩散模型的输出解释为得分函数，可以通过以下方式引导DDPM采样过程对具有高 $p(x|c)$ 的 $x$ 进行采样：$\hat{\epsilon}_\theta(x_t,c)=\epsilon_\theta(x_t,\emptyset)+s\cdot \nabla_{x}\log p(x|c)\propto\epsilon_{\theta}(x_{t},\emptyset)+s\cdot(\epsilon_{\theta}(x_{t},c)-\epsilon_{\theta}(x_{t},\emptyset))$ 。在训练的时候随机丢弃 $c$ 并替换为可学习的 “null” embedding $\emptyset$ 。

3.潜在扩散模型

在高分辨率像素空间中直接训练扩散模型计算成本很高。潜在扩散模型训练了一个autoencoder将图像压缩到更小的空间表示中，然后在潜在表示上训练扩散模型。通过生成潜在表示，可以采用解码器恢复出原图像。

Diffusion Transformer 设计空间

DiT

1.Patchify

DiT的输入是一个空间表示 $z$ 。DiT的第一层是patchify，将空间输入转换为 $T$ 个tokens的序列，每个token的维度为 $d$ 。每一个patch线性嵌入到输入中。然后对所有的输入tokens应用标准ViT的基于频率的位置嵌入。tokens的数量由patch size的超参数 $p$ 决定。

DiT-1