Diffusion Model 深入剖析

最近AI生成艺术领域非常火热，从 Midjourney 到 Stable Diffusion，不管你是绘画高手还是艺术小白，只要输入想要绘制内容的描述或者基础图像，就可以生成富有艺术感的画作！ 这些风格各异、以假乱真的AI生成图像背后,离不开 Diffusion Model 。之前文章《Stable Diffusion原理详解》中我对 Diffusion Model 做了简要的介绍，本文将深入到 Diffusion Model 内部，深入剖析 Diffusion Model 的工作原理以及它是如何生成图像的。

概述

图1. 扩散模型原理概要

Diffusion Model的训练可以分为两部分：

1. 正向扩散过程 → 为图像添加噪声。
2. 反向扩散过程 → 去除图像中的噪声。

正向扩散过程

图2. 正向扩散过程

正向扩散过程向输入图像 $x_0$ 逐步加入高斯噪声，一共 $T$ 步。该过程将产生一系列噪声图像样本 $x_1,\dots ,x_T$。

当 $T\to \infty$ 时，最终的结果将变成一张完全噪声图像，就好像它是从各向同性高斯分布中采样的一样。

但是我们可以使用一个闭合公式在特定的时间步长 $t$ 直接对有噪声的图像进行采样，而不是设计一种算法来迭代地向图像添加噪声。

正向扩散可以用如下公式描述：
$$q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt{1-{\beta }_t}{x}_{t-1},{\beta }_{t}I)\text{\hspace{1em}(1)}$$
其中 $t$ 是时间帧（从 0 到 $T$），$x_t$ 是从真实数据分布 $q(x)$ 中采样的数据样本（例如 $x_0\sim q(x)$），${\beta }_t$ 是 variance schedule，$0\le {\beta }_t\le 1$，且 ${\beta }_0$ 较小，${\beta }_T$ 较大。$I$ 是单位矩阵。

公式推导

可以使用重参数化技巧（Reparameterization Trick）推导出闭合公式。

如果 $z\sim \mathcal{N}(\mu ,{\sigma }^2)$ ，那么 $z$ 可以写成 $z=\mu +\sigma \epsilon$ 的形式，其中 $\epsilon \sim \mathcal{N}(0,1)$。这就是 重参数化技巧 。

利用这个技巧，我们可以将采样图像 $x_t$ 表示为如下形式：
$$x_t=\sqrt{1-{\beta }_t}{x}_{t-1}+\sqrt{{\beta }_t}{\epsilon }_{t-1}\text{\hspace{1em}(2)}$$
这样我们就可以递归地展开它得到闭式公式：
$$\begin{array}{rlr}{x}_t& =\sqrt{1-{\beta }_t}{x}_{t-1}+\sqrt{{\beta }_t}{\epsilon }_{t-1}& \dots \dots \epsilon \sim \mathcal{N}(0,I)\\ & =\sqrt{{\alpha }_t}\boxed{x_{t-1}}+\sqrt{1-{\alpha }_t}{\epsilon }_{t-1}& \dots \dots 设{\alpha }_t=1-{\beta }_t\\ & =\sqrt{{\alpha }_t}\boxed{(\sqrt{{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{\epsilon }_{t-2})}+\sqrt{1-{\alpha }_t}{\epsilon }_{t-1}& \dots \dots 递归展开x_{t-1}\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\boxed{\sqrt{{\alpha }_t(1-{\alpha }_{t-1})}{\epsilon }_{t-2}+\sqrt{1-{\alpha }_t}{\epsilon }_{t-1}}& \dots \dots 乘法分配律乘开\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\boxed{\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\bar{\epsilon }}_{t-2}}& \dots \dots 怎么突然得到这个结果？\\ & ⋮& \\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}\dots {\alpha }_1}{x}_0+\sqrt{1-{\alpha }_t{\alpha }_{t-1}\dots {\alpha }_1}\epsilon & \\ & =\sqrt{\overline{{\alpha }_t}}{x}_0+\sqrt{1-{\bar{\alpha }}_t}\epsilon & \dots \dots \overline{{\alpha }_t}=\prod _{i=1}^t{\alpha }_i\end{array}$$

注意：所有 $\epsilon$ 都是独立同分布的标准正态随机变量。

这里使用不同的符号和下标区分它们很重要，因为它们是独立的并且它们的值在采样后可能不同。

上面的公式推导最难理解的是第4行到第5行，很多人卡在这一步止步不前，下面我将给出详细的推导步骤并解释它是如何工作的：

我们设：
$$\begin{array}{rl}\sqrt{{\alpha }_t(1-{\alpha }_{t-1})}{\epsilon }_{t-2}& =X\\ \sqrt{1-{\alpha }_t}{\epsilon }_{t-1}& =Y\end{array}$$
即
$$x_t=\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\underset{X}{\underbrace{\sqrt{{\alpha }_t(1-{\alpha }_{t-1})}{\epsilon }_{t-2}}}+\underset{Y}{\underbrace{\sqrt{1-{\alpha }_t}{\epsilon }_{t-1}}}$$
应用从重参数化技巧
$$\begin{array}{rl}0+\sqrt{{\alpha }_t(1-{\alpha }_{t-1})}{\epsilon }_{t-2}& ⟹X\sim \mathcal{N}(0,{\alpha }_t(1-{\alpha }_{t-1})I)\\ 0+\sqrt{1-{\alpha }_t}{\epsilon }_{t-1}& ⟹Y\sim \mathcal{N}(0,(1-{\alpha }_t)I)\end{array}$$
设 $Z=X+Y$ ，我们知道，如果 $X\sim \mathcal{N}({\mu }_X,{\sigma }_X^2),Y\sim \mathcal{N}({\mu }_Y,{\sigma }_Y^2)$，则 $Z\sim \mathcal{N}({\mu }_X+{\mu }_Y,{\sigma }_X^2+{\sigma }_Y^2)$

代入 $X$ 和 $Y$ 的实际数值可得
$$\begin{array}{rl}{\mu }_X& =0{\mu }_Y=0\\ & \\ {\sigma }_X^2+{\sigma }_Y^2& ={\alpha }_t(1-{\alpha }_{t-1})+(1-{\alpha }_t)\\ & =\cancel{{\alpha }_t}-{\alpha }_t{\alpha }_{t-1}+1-\cancel{{\alpha }_t}\\ & =1-{\alpha }_t{\alpha }_{t-1}\end{array}$$
所以 $Z\sim \mathcal{N}(0,(1-{\alpha }_t{\alpha }_{t-1})I)$，应用重参数化技巧即可得到：
$$X+Y=Z\sim \mathcal{N}(0,(1-{\alpha }_t{\alpha }_{t-1})I)=0+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\bar{\epsilon }}_{t-2}=\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\bar{\epsilon }}_{t-2}$$
这就得到了第五行的结果。

重复以上步骤，最终我们将得到一个仅取决于输入图像 $x_0$ 的公式：
$$x_t=\sqrt{\overline{{\alpha }_t}}{x}_0+\sqrt{1-{\bar{\alpha }}_t}\epsilon \text{\hspace{1em}(3)}$$
这样我们就可以使用公式（3）在任何时间帧 $t$ 直接对 $x_t$ 进行采样，这极大提高了正向扩散过程的执行效率。

逆向扩散过程

图3. 逆向扩散过程

逆向扩散过程可以用下面的公式描述：
$$q(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\tilde{\mu }}_t(x_t,x_0),{\tilde{\beta }}_{t}I)$$
与正向过程不同，我们不能使用 $q(x_{t-1}|x_t)$ 来逆转噪声，因为它很难处理（无法计算）。

因此，我们需要训练神经网络 $p_{\theta }(x_{t-1}|x_t)$ 来近似 $q(x_{t-1}|x_t)$：
$$p_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),\sum _{\theta }(x_t,t))$$
近似值 $pθ(xₜ₋₁|xₜ) $服从正态分布，其均值和方差需要满足：
$$\{\begin{array}{ll}{\mu }_{\theta }(x_t,t)& :={\tilde{\mu }}_t(x_t,x_0)\\ {\sum }_{\theta }(x_t,t)& :={\tilde{\beta }}_{t}I\end{array}$$

损失函数

我们可以将损失定义为负对数似然：
$$\text{Loss}=-\log (p_{\theta }(x_0))$$
其中 $p_{\theta }(x_0)$ 依赖于 $x_1,x_2,\dots ,x_T$ ，因此处理起来很棘手。

不难发现，这里的设置与变分下界中的设置非常相似。因此我们可以绕开棘手的损失函数本身，转而优化变分下界。通过优化可计算下界，我们可以间接优化棘手的损失函数。

下面是变分下界的推导和展开：
$$\begin{array}{rl}-\log p_{\theta }(x_0)& \le -\log p_{\theta }(x_0)+D_{KL}(q(x_{1:T}|x_0)||p_{\theta }(x_{1:T}|x_0))\\ ⋮& \\ -\log p_{\theta }(x_0)& \le {\mathbb{E}}_q\left[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\right]\\ ⋮& \\ -\log p_{\theta }(x_0)& \le {\mathbb{E}}_q[\underset{L_T}{\underbrace{D_{KL}(q(x_T|x_0)||p_{\theta }(x_T))}}+\sum _{t=2}^T\underset{L_{t-1}}{\underbrace{D_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))}}\underset{L_0}{\underbrace{-\log p_{\theta }(x_0|x_1)}}]\end{array}$$
其中 ${\mathbb{E}}_q\left[\log \frac{q(x_{1:T}|x_0)}{p_{\theta }(x_{0:T})}\right]$ 就是变分下界；展开后的公式由3部分构成，我将其分别命名为$L_T,L_{t-1},L_0$。下面重点解释一下这3部分。

$L_T$ : 常数项

$$L_T=D_{KL}(q(x_T|x_0)||p_{\theta }(x_T))$$

由于 $q(x_T|x_0)$ 没有可学习的参数，$p_{\theta }(x_T)$ 只是一个高斯噪声概率，因此这一项在训练期间是一个常数，可以忽略。

$L_{t-1}$ : 逐步去噪项

$$L_{t-1}=D_{KL}(q(x_{t-1}|x_t,x_0)||p_{\theta }(x_{t-1}|x_t))$$

这一项对目标去噪步骤 $q$ 和近似去噪步骤 $p_{\theta }$ 进行比较。这里通过以 $x_0$ 为条件，让 $q(x_{t-1}|x_t,x_0)$ 变得易于处理。我们分别来看 $q$ 和 $p_{\theta }$
$$\begin{gathered} q(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};\tilde{\mu }(x_t,x_0),{\tilde{\beta }}_{t}I) \\ {\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}\cdot {\beta }_t \\ \tilde{\mu }(x_t,x_0)=\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{x}_0 \\ ⋮ \\ \tilde{\mu }(x_t)=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\epsilon }_t) \end{gathered}$$
经过一系列的推导，$q(x_{t-1}|x_t,x_0)$ 的均值 ${\tilde{\mu }}_t$ 如上所示。其中 $x_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}{\epsilon }_t)$ 。

为了逼近目标去噪步骤 $q$，我们只需要使用神经网络来逼近其均值。因此，我们将近似均值 ${\mu }_{\theta }$ 设置为与目标均值 ${\tilde{\mu }}_t$ 相同的形式（使用可学习的神经网络 ${\epsilon }_{\theta }$）：
$$\begin{array}{rl}\tilde{\mu }(x_t)& =\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}\boxed{{\epsilon }_t})\\ {\mu }_{\theta }(x_t,t)& =\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}\boxed{{\epsilon }_{\theta }(x_t,t)})\end{array}$$
目标均值和近似均值之间的比较可以使用均方误差 (MSE) 来完成：
$$\begin{array}{rl}{L}_t& ={\mathbb{E}}_{x_0,\epsilon }[\frac{1}{2{\sigma }_t^2}||{\tilde{\mu }}_t(x_t)-{\mu }_{\theta }(x_t,t)|{|}^2]\\ & ={\mathbb{E}}_{x_0,\epsilon }[\frac{1}{2{\sigma }_t^2}||\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\epsilon }_t)-\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\epsilon }_{\theta }(x_t,t))|{|}^2]\\ & ={\mathbb{E}}_{x_0,\epsilon }[\frac{(1-{\alpha }_t{)}^2}{2{\alpha }_t(1-{\bar{\alpha }}_t){\sigma }_t^2}||{\epsilon }_t-{\epsilon }_{\theta }(x_t,t)|{|}^2]\end{array}$$
上面公式中 $\frac{(1-{\alpha }_t{)}^2}{2{\alpha }_t(1-{\bar{\alpha }}_t){\sigma }_t^2}$ 是个常数，可以忽略掉，因此简化后的逐步去噪损失为：
$$L_t^{\text{simple}}={\mathbb{E}}_{t\sim [1,T],x_0,{\epsilon }_t}[||{\epsilon }_t-{\epsilon }_{\theta }(x_t,t)|{|}^2]$$
实践中，通过忽略加权项并简单地将目标噪声和预测噪声与 MSE 进行比较，可以获得更好的结果。

因此，事实证明，为了逼近所需的去噪步骤 $q$ ，我们只需要使用神经网络 ${\epsilon }_{\theta }$ 来逼近噪声 ${\epsilon }_t$。

$L_0$ : 重构项

这是最后一步去噪的重构损失，在训练过程中可以忽略，原因如下：

• 可以使用 $L_{t-1}$ 中的相同神经网络对其进行近似。
• 忽略它会使样本质量更好，且更易于实施。

简化损失函数

上面分别解释了 $L_T,L_{t-1},L_0$ 。我们可以发现 $L_T$ 和 $L_0$ 都可以忽略，那么我们的损失函数就可以简化为:
$$\begin{gathered} L_{\text{simple}}={\mathbb{E}}_{t,x_0,\epsilon }[||{\epsilon }_t-{\epsilon }_{\theta }(x_t,t)|{|}^2] \\ {x}_t=\sqrt{\overline{{\alpha }_t}}{x}_0+\sqrt{1-{\bar{\alpha }}_t}\epsilon \end{gathered}$$

U-Net 模型

数据集

在每轮迭代：

• 为每个训练样本（图像）选择一个随机时间步长 $t$。
• 将高斯噪声（对应于 $t$）应用于每个图像。
• 将时间步长转换为嵌入（向量）。

训练

官方给出的训练算法如下：

下面详细解释一下训练步骤是如何工作的：

逆向扩散

我们可以使用上述算法从噪声中生成图像。 下图是具体说明：

注意，在最后一步中，我们只是简单地输出学习到的均值 ${\mu }_{\theta }(x_1,1)$，而不向其添加噪声。

总结

最后对本文的要点做一个总结：

• 扩散模型分为正向扩散和逆向扩散两部分。
• 正向扩散可以使用闭合的公式来完成。
• 可以使用经过训练的神经网络完成逆向扩散。
• 为了逼近去噪步骤 $q$，我们只需要使用神经网络 ${\epsilon }_{\theta }$ 来近似噪声 ${\epsilon }_t$。
• 对简化损失函数的训练产生更好的样本质量。