Diffusion Model 深入剖析-Toy模板网

这篇具有很好参考价值的文章主要介绍了Diffusion Model 深入剖析。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

Diffusion Model 深入剖析

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

Diffusion Model 深入剖析

概述

Diffusion Model 深入剖析

图1. 扩散模型原理概要

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

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

正向扩散过程

Diffusion Model 深入剖析

图2. 正向扩散过程

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

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

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

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

公式推导

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

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

利用这个技巧，我们可以将采样图像 $x_t$ 表示为如下形式：
$x_t = \sqrt{1-\beta_t}x_{t-1}+\sqrt{\beta_t}\varepsilon_{t-1} \tag{2}$
这样我们就可以递归地展开它得到闭式公式：
$\begin{aligned} x_t &= \sqrt{1-\beta_t}x_{t-1}+\sqrt{\beta_t}\varepsilon_{t-1}&\dots\dots \varepsilon \sim \mathcal{N}(0, I)\\ &=\sqrt{\alpha_t}\:\fcolorbox{red}{white}{$x_{t-1}$}+\sqrt{1-\alpha_t}\varepsilon_{t-1}&\dots\dots 设\alpha_t = 1-\beta_t\\ &=\sqrt{\alpha_t}\:\fcolorbox{red}{white}{$\Big(\sqrt{\alpha_{t-1}}x_{t-2}+\sqrt{1-\alpha_{t-1}}\varepsilon_{t-2}\Big)$}+\sqrt{1-\alpha_t}\varepsilon_{t-1}&\dots\dots 递归展开x_{t-1}\\ &=\sqrt{\alpha_t\alpha_{t-1}}x_{t-2}+\:\fcolorbox{red}{white}{$\sqrt{\alpha_t(1-\alpha_{t-1})}\varepsilon_{t-2}+\sqrt{1-\alpha_t}\varepsilon_{t-1}$}&\dots\dots 乘法分配律乘开\\ &=\sqrt{\alpha_t\alpha_{t-1}}x_{t-2}+\:\fcolorbox{red}{white}{$\sqrt{1-\alpha_t\alpha_{t-1}}\bar{\varepsilon}_{t-2}$}&\dots\dots \textcolor{red}{怎么突然得到这个结果？}\\ &\quad\vdots\\ &=\sqrt{\alpha_t\alpha_{t-1}\dots\alpha_1}x_0 + \sqrt{1-\alpha_t\alpha_{t-1}\dots\alpha_1}\varepsilon\\ &=\sqrt{\bar{\alpha_t}}x_0+\sqrt{1-\bar{\alpha}_t}\varepsilon&\dots\dots \bar{\alpha_t}=\prod_{i=1}^t\alpha_i \end{aligned}$

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

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

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

我们设：
$\begin{aligned} \sqrt{\alpha_t(1-\alpha_{t-1})}\varepsilon_{t-2} &= X\\ \sqrt{1-\alpha_t}\varepsilon_{t-1} &= Y \end{aligned}$
即
$x_t=\sqrt{\alpha_t\alpha_{t-1}}x_{t-2}+\underbrace{\sqrt{\alpha_t(1-\alpha_{t-1})}\varepsilon_{t-2}}_X+\underbrace{\sqrt{1-\alpha_t}\varepsilon_{t-1}}_Y$
应用从重参数化技巧
$\begin{aligned} 0 + \sqrt{\alpha_t(1-\alpha_{t-1})}\varepsilon_{t-2} &\implies X \sim \mathcal{N}(0, \alpha_t(1-\alpha_{t-1})I)\\ 0 + \sqrt{1-\alpha_t}\varepsilon_{t-1} &\implies Y \sim \mathcal{N}(0, (1-\alpha_{t})I) \end{aligned}$
设 $Z = X + Y$ ，我们知道，如果 $\sim \mathcal{N}(\mu_X, \sigma_X^2), \quad Y \sim \mathcal{N}(\mu_Y, \sigma_Y^2)$ ，则 $\sim \mathcal{N}(\mu_X+\mu_Y, \sigma_X^2+\sigma_Y^2)$

代入 $X$ 和 $Y$ 的实际数值可得
$\begin{aligned} \mu_X &= 0 \qquad \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{aligned}$
所以 $\sim \mathcal{N}(0, (1-\alpha_t\alpha_{t-1})I)$ ，应用重参数化技巧即可得到：
$\sim \mathcal{N}(0, (1-\alpha_t\alpha_{t-1})I)=0+\sqrt{1-\alpha_t\alpha_{t-1}}\:\bar\varepsilon_{t-2}=\sqrt{1-\alpha_t\alpha_{t-1}}\:\bar\varepsilon_{t-2}$
这就得到了第五行的结果。

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

逆向扩散过程

Diffusion Model 深入剖析

图3. 逆向扩散过程

逆向扩散过程可以用下面的公式描述：
$q(x_{t-1}|x_{t}) = \mathcal{N}(x_{t-1};\tilde\mu_t(x_t, x_0),\tilde\beta_tI)$
与正向过程不同，我们不能使用 $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_\theta
(x|xₜ) $服从正态分布，其均值和方差需要满足：
$\begin{cases} \mu_\theta(x_t, t) &:= \tilde\mu_t(x_t, x_0)\\ \sum_\theta(x_t, t) &:= \tilde\beta_tI \end{cases}$

损失函数

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

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

Diffusion Model 深入剖析

下面是变分下界的推导和展开：
$\begin{aligned} -\log p_\theta(x_0) &\le -\log p_\theta(x_0) + D_{KL}\big(q(x_{1:T}\:|\:x_0)\:||\:p_\theta(x_{1:T}\:|\:x_0)\big)\\ \vdots\\ -\log p_\theta(x_0) &\le \mathbb{E}_q\Big[\log \frac{q(x_{1:T}|x_0)}{p_\theta(x_{0:T})}\Big]\\ \vdots\\ -\log p_\theta(x_0) &\le \mathbb{E}_q\Big[\underbrace{D_{KL}(q(x_T|x_0)||p_\theta(x_T))}_{L_T} + \sum_{t=2}^T \underbrace{D_{KL}(q(x_{t-1}|x_t,x_0)||p_\theta(x_{t-1}|x_t))}_{L_{t-1}} \underbrace{-\log p_\theta(x_0|x_1)}_{L_0}\Big] \end{aligned}$
其中 $\mathbb{E}_q\Big[\log \frac{q(x_{1:T}|x_0)}{p_\theta(x_{0:T})}\Big]$ 就是变分下界；展开后的公式由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$
$q(x_{t-1}|x_t,x_0) = \mathcal{N}(x_{t-1};\tilde\mu(x_t, x_0), \tilde\beta_tI)\\ \tilde\beta_t = \frac{1-\bar\alpha_{t-1}}{1-\bar\alpha_t} \sdot \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\\ \vdots \\ \tilde\mu(x_t) = \frac{1}{\sqrt{\alpha_t}}\Big(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\varepsilon_t\Big)$
经过一系列的推导， $q(x_{t-1}|x_t,x_0)$ 的均值 $\tilde\mu_t$ 如上所示。其中 $x_0 = \frac{1}{\sqrt{\bar\alpha_t}}\Big(x_t-\sqrt{1-\bar\alpha_t}\varepsilon_t\Big)$ 。

为了逼近目标去噪步骤 $q$ ，我们只需要使用神经网络来逼近其均值。因此，我们将近似均值 $\mu_\theta$ 设置为与目标均值 $\tilde{\mu}_t$ 相同的形式（使用可学习的神经网络 $\varepsilon_\theta$ ）：
$\begin{aligned} \tilde\mu(x_t) &= \frac{1}{\sqrt{\alpha_t}}\Big(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\boxed{\varepsilon_t}\Big)\\ \mu_\theta(x_t,t) &= \frac{1}{\sqrt{\alpha_t}}\Big(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\boxed{\varepsilon_\theta(x_t,t)}\Big) \end{aligned}$
目标均值和近似均值之间的比较可以使用均方误差 (MSE) 来完成：
$\begin{aligned} L_t &= \mathbb{E}_{x_0,\varepsilon}\Big[\frac{1}{2\sigma_t^2} ||\tilde{\mu}_t(x_t)-\mu_\theta(x_t, t)||^2 \Big]\\ &= \mathbb{E}_{x_0,\varepsilon}\Big[\frac{1}{2\sigma_t^2} ||\frac{1}{\sqrt{\alpha_t}}\big(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\varepsilon_t\big) - \frac{1}{\sqrt{\alpha_t}}\big(x_t-\frac{1-\alpha_t}{\sqrt{1-\bar\alpha_t}}\varepsilon_\theta(x_t,t)\big)||^2 \Big]\\ &= \mathbb{E}_{x_0,\varepsilon}\Big[\frac{(1-\alpha_t)^2}{2\alpha_t(1-\bar\alpha_t)\sigma_t^2}||\varepsilon_t-\varepsilon_\theta(x_t,t)||^2\Big] \end{aligned}$
上面公式中 $\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,\varepsilon_t}\Big[||\varepsilon_t-\varepsilon_\theta(x_t,t)||^2\Big]$
实践中，通过忽略加权项并简单地将目标噪声和预测噪声与 MSE 进行比较，可以获得更好的结果。

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

$L_0$ : 重构项

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

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

简化损失函数

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