【Paper Notes】Denoising Diffusion Probabilistic Models

Kross Sun

已于 2022-11-21 21:26:48 修改

阅读量852

点赞数 1

文章标签：人工智能

于 2022-11-07 12:16:26 首次发布

本文链接：https://blog.csdn.net/S_Kross/article/details/127727964

版权

近来Diffusion Model在数据生成领域取得了很重要的进展，比如DALLE-2，和google的imagen等，效果非常惊人。因此有必要学习一下这个方法。

Diffusion Model的基本过程

Diffusion Model 有两个基本过程，分为前向（forward diffusion process) 和逆向（reverse diffusion process).

前向过程，是向图像 $\bold{x}_0$ 中不断加入噪声 $\bold{\epsilon}\sim \mathcal{N}(0,1)$ 。直到数据完全退化成高斯噪声 $\bold{x}_T$ 。

逆向过程，是向由噪声 $\bold{x}_T$ ,不断预测加入的噪声，一步一步恢复出 $\bold{x}_0$ 的过程。

我们可以用公式化的描述来理解这个过程。

Forward

$\bold{x}_t = \sqrt{1-\beta_t} \bold{x}_{t-1} +\sqrt{\beta_t} \bold{\epsilon}_t, \epsilon_t \sim \mathcal N(\bold{0}, \bold{I}) \ \ \ \ \ \ \ \ \ \ \ (1)$

注意等式右边两项的参数的平方和是1，这是一个有意的设计。原因是？请先思考，文末有答案。

并且， $\beta_t$ 一般都取比较小的数字，比如0.001，这是为了保证其逆向过程也近似为正态分布。

Reverse

$\bold{x}_{t-1} = \frac{1}{\sqrt{1-\beta_t}} \bold{x}_t - \frac{\sqrt{\beta_t}} {\sqrt{1-\beta_t}} \epsilon_{\theta}(\bold{x}_t,t) \ \ \ \ \ \ (2)$

在reverse过程中，我们一般通过一个UNet来对 $\bold{\epsilon}_\theta$ 进行建模。

训练目标函数

清晰了diffusion model的基本过程，那么就会出现一个问题，如何训练这个UNet？

最简单的方法就是，根据Forward过程，采样得到 $\bold{x}_t, \bold{x}_{t-1}$ ,然后直接求解出 $\bold{\epsilon}_t$ ,如下所示：

$L_{obj} = || \bold{\epsilon}_t - \bold{\epsilon}_{\theta}(\bold{x}_t, t)||^2=|| \frac{\bold{x}_t-\sqrt{1-\beta_t}\bold{x}_{t-1}}{\sqrt{\beta_t}} - \bold{\epsilon}_{\theta}(\bold{x}_t, t)||^2 \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ (3)$

但是这个过程显然不够高效，因为每次训练之前，都得执行一遍完整的Forward Process。因此需要考虑更高效的训练方式，如果可以直接用 $\bold{x}_0$ 训练，不用执行forward过程是最好的。

递推关系

这里可以先推导 $\bold{x}_0,\bold{x}_t$ 的关系。

首先令 $\alpha_t = 1 - \beta_t, \bar \alpha_t = \prod \limits_{i=1}^t \alpha_i$ , 根据公式（1）可以得到：

$\begin{align*} \bold{x}_t &= \sqrt{\alpha_t} \bold{x}_{t-1} + \sqrt{1-\alpha_t}\bold{\epsilon}_{t} \\ &= \sqrt{\alpha_t}( \sqrt{\alpha_{t-1}} \bold{x}_{t-2} + \sqrt{1-\alpha_{t-1}}\bold{\epsilon}_{t-1} ) + \sqrt{1-\alpha_t}\bold{\epsilon}_{t} \\ &= \sqrt{\alpha_t \alpha_{t-1}} \bold{x}_{t-2} + \sqrt{\alpha_t - \alpha_t \alpha_{t-1}} \epsilon_{t-1} + \sqrt{1-\alpha_{t}}\epsilon_{t} \\ &= \sqrt{\alpha_t \alpha_{t-1}} \bold{x}_{t-2} + \sqrt{1-\alpha_t \alpha_{t-1}} \bar \epsilon_{t} \end{align*}$

其中，最后一步是根据两个高斯分布相加，还是高斯分布，并且其均值是两个分布的均值之和，新的方差是两个分布方差之和。

根据上式，可以得到

$\bold{x}_t = \sqrt{\bar \alpha_t} \bold{x}_0 + \sqrt{1 - \bar \alpha_t} \bar \epsilon_t, \bar \epsilon_t \sim \mathcal N(\bold{0},\bold{I}) \ \ \ \ \ \ \ \ (4)$

训练目标简化

我们将（4）式直接带入（3）式，可以得到

$L_{obj} = || \epsilon_t - \epsilon_{\theta} (\sqrt{\bar\alpha_t}\bold{x}_{0}+\sqrt{1-\bar \alpha_{t}}\bar \epsilon_{t}, t)||^2$

这样的话，其实每次只需要采样 $\epsilon_{t}, \bar \epsilon_{t}$ 就可以（其实有点问题， $\epsilon_{t},\bar\epsilon_{t}$ 可能不互相独立），但是这需要采样两次，还没有达到最优，最优情况下只需要采样一次，并且需要采样两个高维随机变量，可能会导致训练不稳定，难以收敛。因此还需要对式子进行化简。

我们采用 $\bold{x}_t \mapsto \bold{x}_{t-1} \mapsto \bold{x}_0$ 来进行递推。

可以得到，

$L_{obj} = || \epsilon_t - \epsilon_{\theta} (\sqrt{\bar\alpha_t}\bold{x}_{0}+\sqrt{\alpha_t - \bar\alpha_{t}} \bar \epsilon_{t-1} + \sqrt{1-\alpha_t}\epsilon_t, t)||^2$

到这里如何进一步化简呢？

我们令

$\begin{align*} \sqrt{1-\bar\alpha_t}\mathbf{w}& = \sqrt{\alpha_t - \bar \alpha_t} \bar \epsilon_{t-1} + \sqrt{1-\alpha_t}\epsilon_t , \mathbf{w} \sim \mathcal N(\bold{0}, \bold{I}) \\ \sqrt{1-\bar\alpha_t} \mathbf{s}& = \sqrt{\alpha_t - \bar \alpha_t} \bar \epsilon_{t-1} - \sqrt{1-\alpha_t}\epsilon_t , \mathbf{s} \sim \mathcal N(\bold{0}, \bold{I}) \end{align*}$

因此，可以得到，注意 $\bold{w},\bold{s}$ 并不独立？

$L_{obj} = ||\frac{\sqrt{1-\bar\alpha_t}}{2\sqrt{1-\alpha_t}}(\bold{w}-\bold{s}) - \epsilon_{\theta}(\sqrt{\bar\alpha_t}\bold{x}_0 + \sqrt{1-\bar\alpha_t}\bold{w},t) ||^2$

又因为，在计算损失时，我们通常考虑期望，即

$L_{obj} = E_{\bold{w}, \bold{s}}\big[||\frac{\sqrt{1-\bar\alpha_t}}{2\sqrt{1-\alpha_t}}(\bold{w}-\bold{s}) - \epsilon_{\theta}(\sqrt{\bar\alpha_t}\bold{x}_0 + \sqrt{1-\bar\alpha_t}\bold{w},t) ||^2\big]$

并且， $\bold{s}$ 只包含二次项和一次项，所以在忽略常数的情况下，等价于优化

$L_{obj} = ||\frac{\sqrt{1-\bar\alpha_t}}{2\sqrt{1-\alpha_t}}\bold{w} - \epsilon_{\theta}(\sqrt{\bar\alpha_t}\bold{x}_0 + \sqrt{1-\bar\alpha_t}\bold{w},t) ||^2$

上述推导有点问题，越推导越感觉不对，先贴出目前比较认可的损失函数：

$L_{t}^{simple} = \mathbb{E}_{t\sim[1,T],\bold{x}_0,\epsilon_t\sim\mathcal N(\bold{0},\bold{I})}\big[ \| \epsilon_t - \epsilon_{\theta}(\sqrt{\bar\alpha_t}\bold{x}_0 + \sqrt{1-\bar \alpha_{t}} \epsilon_t, t)\|^2\big]$