Denoising Diffusion Probabilistic Models论文解读

Denoising Diffusion Probabilistic Models^[1]

《Denoising Diffusion Probabilistic Models》简称ddpm。扩散模型的里程碑式文章，简单来说，扩散模型包含两个过程：

1. 前向过程（从右向左）：原始分布\(x_0\)不断加入噪声，经过\(T\)次加噪，最终得到一个标准正太分布\(x_T\sim\mathcal N(0,1)\)
2. 逆向过程（从左向右）：从一个标准正态分布\(x_T\sim\mathcal N(0,1)\)，通过不断降噪，学习一个\(\epsilon_\theta\)最终得到图像\(x_0\)

伪代码如下，具体内容会在损失函数小结详细展开：

在三篇前文中我们介绍了ddpm的背景知识、前向过程推导、逆向过程推导和扩散模型损失函数。我们将三篇文章的内容在此总结，并且进一步介绍ddpm的改进。

前向过程

在前向\(x_{t-1}\rightarrow x_{t}\)加噪的数据分布\(q\)人为的定义如下：

\[q(x_t|x_{t-1}):=\mathcal N(x_t;\sqrt{1-\beta_t}x_{t-1},\beta_tI) \tag{Eq.1} \]

我们通过推导得到了\(x_0\rightarrow x_T\)一步到位的数据分布，如下。其中\(\alpha_t=1-\beta_t, \bar \alpha_t=\alpha_t\alpha_{t-1}...\alpha_1\)：

\[q(x_t|x_0) = \mathcal N(x_t;\sqrt{\bar\alpha_t}x_0,(1-\bar\alpha_t)I) \tag{Eq.2} \]

根据重参数化采样的技巧，Eq.2可以得到如下公式：

\[x_t=\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon, \epsilon\in\mathcal N(0,1) \tag{Eq.3} \]

Eq.3可以看出前向过程就是通过对原始数据\(x_0\)按照\(\bar\alpha\)进行采样后，添加标准正态分布的噪声\(\epsilon\)，最后得到\(x_T\sim\mathcal N(0,1)\) 。我们可以观察到Eq.3展示了\(x_t\)和\(x_0\)是存在一种数据关系可以互相转换的：

\[x_0= \frac{1}{\sqrt{\bar\alpha_t}}(x_t-\sqrt{1-\bar\alpha_t}\epsilon) \tag{Eq.4} \]

我们可以将\(x_t\)和\(x_0\)看成是一个函数关系，即\(x_t(x_0,\epsilon)\)

逆向过程

所谓逆向过程，就是去除前向过程的噪声。逆向分布我们还是用\(q\)表示，因为我们想要知道逆向该如何进行，才能训练一个神经网络\(p_\theta\)去学习逆向分布。我们在前文中得到了逆向过程\(x_t\rightarrow x_{t-1}\)的分布如下：

\[q(x_{t-1}|x_t,x_0)=\mathcal N(x_{t-1};\tilde\mu_t(x_t,x_0), \tilde\beta_tI) \tag{Eq.5} \]

其中\(\tilde \mu (x_t,x_0)\)和\(\tilde \beta\)定义如下：

\[\tilde \mu_t(x_t, x_0)=\frac{\sqrt{\alpha_t}{(1-\bar\alpha_{t-1})}}{1-\bar{\alpha}_{t}}x_t+\frac{\beta_t\sqrt{\bar\alpha_{t-1}}}{1-\bar{\alpha}_{t}}x_0,\tilde \beta=\frac{{(1-\bar\alpha_{t-1})}}{1-\bar{\alpha}_{t}}\beta_t \tag{Eq.6} \]

由于在Eq.4发现\(x_t\)和\(x_0\)互相转换的关系，所以我们可以重新表示\(\tilde \mu\):

\[\begin{aligned} \tilde \mu_t(x_t, x_0)&=\tilde\mu\left(x_t\left(x_0,\epsilon\right),\frac{1}{\sqrt{\bar\alpha_t}}(x_t(x_0,\epsilon)-\sqrt{1-\bar\alpha_t}\epsilon)\right)\\ &=\frac{\sqrt{\alpha_t}{(1-\bar\alpha_{t-1})}}{1-\bar{\alpha}_{t}}x_t\left(x_0,\epsilon\right)+\frac{\beta_t\sqrt{\bar\alpha_{t-1}}}{1-\bar{\alpha}_{t}}\frac{1}{\sqrt{\bar\alpha_t}}(x_t\left(x_0,\epsilon\right)-\sqrt{1-\bar\alpha_t}\epsilon) \\ &=\frac{\alpha_t(1-\bar\alpha_{t-1})}{\sqrt{\alpha_t}(1-\bar\alpha_t)}x_t\left(x_0,\epsilon\right)+\frac{\beta_t(x_t\left(x_0,\epsilon\right)-\sqrt{1-\bar\alpha_t}\epsilon)}{\sqrt{\alpha_t}(1-\bar\alpha_t)}\\ &=\frac{1}{\sqrt{\alpha_t}}\left[\frac{(\alpha_t-\bar\alpha_t+\beta_t)x_t\left(x_0,\epsilon\right)}{1-\bar\alpha_t}-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon\right] \\ &=\frac{1}{\sqrt{\alpha_t}}\left[x_t\left(x_0,\epsilon\right)-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon\right] \end{aligned} \tag{Eq.7} \]

这样我们的\(\tilde \mu\)的表达式就可以只包含\(x_t\)。我们之所以煞费苦心地把式子把\(x_0\)转换掉，是因为在真正推理阶段，我们不可能又原图\(x_0\)，我们只会有\(x_t\)的状态。

同样通过重参数化采样的技巧Eq.5可以得到如下关系

\[\begin{aligned} x_{t-1}&=\tilde\mu_t(x_t, x_0)+\tilde\beta\epsilon \\ &=\frac{\sqrt{\alpha_t}{(1-\bar\alpha_{t-1})}}{1-\bar{\alpha}_{t}}x_t+\frac{\beta_t\sqrt{\bar\alpha_{t-1}}}{1-\bar{\alpha}_{t}}x_0+\frac{{(1-\bar\alpha_{t-1})}}{1-\bar{\alpha}_{t}}\beta_t\epsilon \\ &=\frac{1}{\sqrt{\alpha_t}}\left[x_t\left(x_0,\epsilon\right)-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon\right]+\frac{{(1-\bar\alpha_{t-1})}}{1-\bar{\alpha}_{t}}\beta_t\epsilon \end{aligned} \tag{Eq.8} \]

损失函数

至此我们表示出了我们需要的前向过程和逆向过程，我们通过训练一个\(p_\theta(x_{t-1}|x_t):=\mathcal N(x_{t-1};\mu_\theta(x_t,t),\Sigma_\theta(x_t, t)I)\) 来拟合逆向过程，其中\(\Sigma_\theta(x_t, t)\)我们定义为一个常量\(\sigma_t^2 = \beta_t\)，这种定义是为了计算简单，如果引入了方差计算，理论会变得更加复杂。最终的损失函数如下：

\[L = \underbrace{D_{KL}[q(x_T|x_0)\|p(x_T)]}_{L_T}+\underbrace{\sum_{t>1}^TD_{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} \tag{Eq.9} \]

我们观察损失函数中三项的含义：

• \(L_T\)：加噪的结果\(q(x_T|x_0)\)和神经网络的起始分布\(p(x_T)\)尽可能相似，很符合逻辑。由于\(q(x_T|x_0)\)中的\(\beta_t\)是人为定义常量，所以
• \(L_{t-1}\): 用\(q\)得到的正确逆向过程的数据分布和神经网络训练的逆向过程尽可能相似，很符合逻辑
• \(L_0\): 属于\(t=1\)的特殊项，和\(L_{t-1}\)的含义一致

我们主要关注的就是\(L_{t-1}\)这个逆向过程。我们知道\(q(x_{t-1}|x_t,x_0)\)的表示，通过高斯分布的KL散度公式可以将\(L_{t-1}\)写成如下形式：

\[D_{KL}[q(x_{t-1}|x_t,x_0)\|p_\theta(x_{t-1}|x_t)]=\log\frac{\sigma_p}{\sigma_q}+\frac{\sigma_q^2+(\tilde \mu_q-\tilde \mu_p)^2}{2\sigma_p^2} -\frac{1}{2} \tag{Eq.10} \]

其中\(\sigma_p=\beta_t\)是一个常量，\(\sigma_q=\tilde \beta\)在Eq.6中也是常量，\(\frac{1}{2}\)也是常量，在训练的时候都对梯度起不到作用。我们的关注点只有\(\tilde \mu\)，如此\(L_{t-1}\)的损失函数如下(对应原论文(8))：

\[L_{t-1}=\frac{1}{2\sigma_t^2}\|\tilde \mu_t(x_t,x_0)-\mu_\theta(x_t,t)\|^2+C \tag{Eq.11} \]

我们将Eq.7引入Eq.11中(对应原论文(10))：

\[\begin{aligned} L_{t-1} - C &=\frac{1}{2\sigma_t^2}\left\|\tilde \mu_t(x_t,x_0)-\mu_\theta(x_t,t)\right\|^2 \\ &=\frac{1}{2\sigma_t^2}\left\|\frac{1}{\sqrt{\alpha_t}}\left[x_t\left(x_0,\epsilon\right)-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon\right]-\mu_\theta(x_t,t)\right\|^2 \end{aligned} \tag{Eq.12} \]

在Eq.12中我们可以看到我们\(\mu_\theta\)训练目标就是预测\(\frac{1}{\sqrt{\alpha_t}}\left[x_t\left(x_0,\epsilon\right)-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon\right]\)。我们使用Eq.7方法同样表示\(\mu_\theta\)如下：

\[\mu_\theta(x_t,t)=\frac{1}{\sqrt{\alpha_t}}\left[x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon_\theta(x_t,t)\right] \tag{Eq.13} \]

在这里\(\epsilon_\theta\)函数用于预测\(x_t\)的\(\epsilon\)。则我们神经网络训练的数据分布\(p_\theta(x_{t-1}|x_t):=\mathcal N(x_{t-1};\mu_\theta(x_t,t),\sigma_t^2I)\) 就是计算如下公式：

\[x_{t-1}=\frac{1}{\sqrt{\alpha_t}}\left[x_t-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon_\theta(x_t,t)\right]+\sigma_t^2z, z\in\mathcal N(0,1) \tag{Eq.14} \]

和算法2的伪代码中间的采样函数一致。Eq.14是用于具体采样阶段的计算公式。

简化损失函数\(L_{t-1}\)

我们的损失函数目前是Eq.12的形式，我们将Eq.13带入Eq.12，并且扩写\(x_t\)函数(对应原论文(12))：

\[\begin{aligned} L_{t-1} - C &=\frac{1}{2\sigma_t^2}\left\|\frac{1}{\sqrt{\alpha_t}}\left[x_t\left(x_0,\epsilon\right)-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon\right]-\mu_\theta(x_t,t)\right\|^2 \\ &=\frac{1}{2\sigma_t^2}\left\|\frac{1}{\sqrt{\alpha_t}}\left[x_t\left(x_0,\epsilon\right)-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon \right] -\frac{1}{\sqrt{\alpha_t}}\left[x_t\left(x_0,\epsilon\right)-\frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon_\theta(x_t,t)\right]\right\|^2 \\ &=\frac{\beta_t^2}{2\sigma_t^2\alpha_t(1-\bar\alpha_t)}\left\|\epsilon-\epsilon_\theta(\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon, t)\right\|^2 \end{aligned} \tag{Eq.15} \]

我们可以省略系数，得到\(L_{t-1}\)的最终最简化版本如下：

\[L_{simple}(\theta)=\left\|\epsilon-\epsilon_\theta(\sqrt{\bar\alpha_t}x_0+\sqrt{1-\bar\alpha_t}\epsilon, t)\right\|^2 \]

我们可以看到损失函数从最初的Eq.9的最小化KL散度，变成了 Eq.11的预测\(\mu\)，再进一步变成了Eq.15的预测噪声\(\epsilon\)，几乎达到了最简化的情况。综上ddpm针对图像生成的训练目标，从复原数据分布，变成了复原分布的均值\(\mu\)，再进一步变成了预测噪声。其实在结束了所有推导后，重新看这个结论就非常直觉，前向过程通过加噪声将一张图变成高斯分布，那么训练的目标就是去还原噪声就可以恢复一张图。ddpm主要是从理论角度证实了两者的关系，并且确定了扩散模型预测噪声就足够了。

---

1. Ho J, Jain A, Abbeel P. Denoising diffusion probabilistic models[J]. Advances in neural information processing systems, 2020, 33: 6840-6851. ↩︎