DDPM详解

DDPM详解

参考 https://www.bilibili.com/video/BV1pa411u7G3/ 系列

DDPM 可以分为 Diffusion 和 Reverse 两个阶段。其中 Diffusion 阶段通过不断地对真实图片添加噪声，最终得到一张噪声图片。而 Reverse 阶段，模型需要学习预测出一张噪声图片中的噪声部分，然后减掉该噪声部分，即：去噪。随机采样一张完全噪声图片，通过不断地去噪，最终得到一张符合现实世界图片分布的真实图片。以下分别介绍两个阶段的具体原理与公式推导。

Diffusion 阶段

Diffusion 阶段就是不断地给真实图片加噪声，经过 $T$ 步加噪之后，得到一张完全噪声图片。要理解 diffusion 阶段，首先要理解每一步噪声是如何加的。即：单步扩散的过程，即从 $x_{t-1}$ 到 $x_t$ 的加噪声过程。这里，我们记真实图片为 $x_0$，中间过程的噪声图片为 $x_1,...x_{T-1}$ ，最终的完全噪声图片为 $x_T$ 。则有，单步扩散过程公式：
$$x_t=\sqrt{1-{\beta }_t}{x}_{t-1}-\sqrt{{\beta }_t}{z}_t{z}_t\sim \mathcal{N}(0,I)$$
可以看到，就是在输入图像的基础上加一个高斯噪声 $z_t$。其中 ${\beta }_{1\dots T}$ 是超参数，在 DDPM 原文 中， ${\beta }_{1\dots T}$ 从 $10^{-4}$ 到 $2\times 10^{-2}$ 线性增加。

训练时，这样单步单步地加噪声效率太低了。想要提高训练效率。需要实现从原图 $x_0$ 一步到位得到 $t$ 步的扩散结果 $x_t$，即我们现在已知的是 $x_t=f(x_{t-1})$ ，想要的是 $x_t=f(x_0)$

首先，记 ${\alpha }_t=1-{\beta }_t$ ，则有：
$$x_t=\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{z}_t$$
将 $x_{t-1}$ 代换为 $x_{t-2}$，将 $x_{t-2}$ 代换为 $x_{t-3}$ ，直到 $x_0$ ，推导如下：
$$\begin{array}{rl}{x}_t& =\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{z}_t\\ & =\sqrt{{\alpha }_t}(\sqrt{{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{z}_{t-1})+\sqrt{1-{\alpha }_t}{z}_t\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}{z}_{t-1}+\sqrt{1-{\alpha }_t}{z}_t\end{array}$$
由于 $z_t$、$z_{t-1}$ 是两次独立的对 $\mathcal{N}(0,I)$ 的采样，因此有：
$$\begin{gathered} \sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}{z}_{t-1}\sim \mathcal{N}(0,({\alpha }_t-{\alpha }_t{\alpha }_{t-1})I) \\ \sqrt{1-{\alpha }_t}{z}_t\sim \mathcal{N}(0,(1-{\alpha }_t)I) \\ \sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}{z}_{t-1}+\sqrt{1-{\alpha }_t}{z}_t\sim \mathcal{N}(0,(1-{\alpha }_t{\alpha }_{t-1})I) \end{gathered}$$
再结合重参数化技巧，有：
$$x_t=\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}z,z\sim \mathcal{N}(0,I)$$
推到 $x_0$，有：
$$\begin{array}{rl}{x}_t& =\sqrt{{\alpha }_t{\alpha }_{t-1}\dots {\alpha }_1}{x}_0+\sqrt{1-{\alpha }_t{\alpha }_{t-1}\dots {\alpha }_1}z,z\sim \mathcal{N}(0,I)\\ & =\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}z,记{\bar{\alpha }}_t=\prod _{i=1}^T{\alpha }_i\end{array}$$
diffusion 阶段总结：

• 单步扩散 ($x_{t-1}\to x_t$)：$x_t=\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{z}_t$

• 一步到位 ($x_0\to x_t$)：$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}z$

reverse阶段

如下图所示，我们先简单看下 diffusion 训练的过程。假设一个 batch size 为 4，最大时间步 $T=1000$ 。首先从均匀分布 $U(1,1000)$ 中采样出 4 个对应的时间步 $t$。然后使用我们的一步到位加噪公式得到噪声图。训练一个神经网络（通常是 UNet）得到预测噪声 $\tilde{z}$，计算其与实际添加噪声 $z$ 的损失 $loss(z,\tilde{z})$ ，反传梯度，更新网络权重。

在训练完成后，采样推理生图时，该怎么做呢？直觉上，我们已经有噪声的估计值 $\tilde{z}$ ，那么根据一步到位公式（式 (6)），可以直接算出真实图片 $x_0$ ：
$$x_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}\tilde{z})$$
但实际上，这种一步到位 reverse 的生图效果很差。在生图时，我们最好还是按照 diffusion 过程原始的定义，单步去噪。即由 $x_T$ 估计出 $x_{T-1}$，再到 $x_{T-2}$，一直单步到 $x_0$ 。为了完成这种单步 reverse 的过程，我们需要推导出根据 $x_t$ 和预测出的噪声 $\tilde{z}$ 如何计算 $x_{t-1}$，即 $x_{t-1}=f(x_t,\tilde{z})$ 。即下图红线所示

我们要推的是 $x_t\to x_{t-1}$ 的过程，即条件概率 $q(x_{t-1}|x_t)$，根据贝叶斯公式，有：
$$q(x_{t-1}|x_t)=\frac{q(x_t|x_{t-1})q(x_{t-1})}{q(x_t)}$$
注意 $q(x_t|x_{t-1})$ 和 $q(x_t|x_0)$ 分别就是我们前向 diffusion 过程中的单步扩散和一步到位的公式，它们都可以写成高斯分布的形式，即分别有：
$$\begin{gathered} x_t=\sqrt{{\alpha }_t}{x}_{t-1}-\sqrt{1-{\alpha }_t}{z}_t\sim \mathcal{N}(\sqrt{{\alpha }_t}{x}_{t-1},(1-{\alpha }_t)I) \\ {x}_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}z\sim \mathcal{N}(\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)I) \end{gathered}$$
然后，由于我们假设 diffusion 过程是一个马尔可夫过程，则 $x_t$ 只与 $x_{t-1}$ 有关，与 $x_{t-\mathrm{2...0}}$ 都无关，都是独立的。那么我们可以将 $x_0$ 加到条件中，而不影响结果，即 $q(x_t)=q(x_t|x_0)$。则 $q(x_{t-1}|x_t)$ 式中的三项可分别写为：
$$\begin{gathered} q(x_t|x_{t-1})\sim \mathcal{N}(\sqrt{{\alpha }_t}{x}_{t-1},(1-{\alpha }_t)I),式9(1)单步扩散 \\ q(x_t)=q(x_t|x_0)\sim \mathcal{N}(\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)I),式9(2)一步到位 \\ q(x_{t-1})=q(x_{t-1}|x_0)\sim \mathcal{N}(\sqrt{{\bar{\alpha }}_{t-1}}{x}_0,(1-{\bar{\alpha }}_{t-1})I),式9(2)一步到位 \end{gathered}$$

而在已知高斯分布的均值和方差时，有正比关系：$N(\mu ,{\sigma }^2)\propto \text{exp}(-\frac{1}{2}\frac{(x-\mu {)}^2}{{\sigma }^2})$， 将上面几个高斯分布的均值和方差分别代入，有：
$$\begin{array}{rl}q(x_{t-1}|x_t)& =\frac{q(x_t|x_{t-1})q(x_{t-1})}{q(x_t)}\\ & \propto \text{exp}(-\frac{1}{2}(\frac{(x_t-\sqrt{{\alpha }_t}{x}_{t-1}{)}^2}{1-{\alpha }_t}+\frac{(x_{t-1}-\sqrt{{\alpha }_{t-1}}{x}_0{)}^2}{1-{\bar{\alpha }}_{t-1}}-\frac{(x_t-\sqrt{{\alpha }_t}{x}_0{)}^2}{1-{\bar{\alpha }}_t}))\end{array}$$
注意，我们的目标是求分布 $q(x_{t-1}|x_t)$ ，我们也将其形式化为一个高斯分布。可以观察到，目前我们推导的结果是一个 $x_{t-1}$ 的二次式，将其配方，找到我们关心的 $q(x_{t-1}|x_t)$ 的均值和方差。 配方有：
$$\begin{array}{rl}q(x_{t-1}|x_t)& \propto \text{exp}(-\frac{1}{2}[\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}}]x_{t-1}^2-2[\frac{\sqrt{{\alpha }_t}{x}_t}{{\beta }_t}+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{x}_0}{1-{\bar{\alpha }}_{t-1}}]x_{t-1}+?)\\ & \propto \text{exp}(-\frac{1}{2}A{x}_{t-1}^2+B{x}_{t-1}+C)\\ & \propto \text{exp}(-\frac{1}{2}\cdot A(x_{t-1}+\frac{B}{2A}{)}^2+?)\end{array}$$
别忘了 ${\alpha }_t=1-{\beta }_t$ 。

$C$ 的具体值是多少我们不关心，这样我们就得到均值和方差：
$$\mu =-\frac{B}{2A},{\sigma }^2=\frac{1}{A}$$
其中，
$$A=\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}},B=-2(\frac{\sqrt{{\alpha }_t}{x}_t}{{\beta }_t}+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{x}_0}{1-{\bar{\alpha }}_{t-1}})$$
整理得方差：
$$\begin{array}{rl}{\sigma }^2& =\frac{1}{A}=1/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})\\ & =1/(\frac{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}+{\beta }_t}{{\beta }_t(1-{\bar{\alpha }}_{t-1})})\\ & =\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t\end{array}$$
均值：
$$\begin{array}{rl}\mu & =-\frac{B}{2A}=(\frac{\sqrt{{\alpha }_t}{x}_t}{{\beta }_t}+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{x}_0}{1-{\bar{\alpha }}_t})\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t\\ & =\sqrt{{\alpha }_t}\frac{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\end{array}$$

我们发现，$q(x_{t-1}|x_t)$ 的方差完全由 $\alpha ,\beta$ 计算得出，这是设定好的超参数。而 $q(x_{t-1}|x_t)$ 的均值是 $x_t$ 和 $x_0$ 的线性组合。目前其中 $x_t$ 是已知的，而 $x_0$ 也是未知的，我们还是把一步到位的 diffusion 过程的公式转换而来的式 (7) 代换 $x_0$，则有：
$$\mu =\sqrt{{\alpha }_t}\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{x}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}\cdot \frac{1}{\sqrt{{\bar{\alpha }}_t}}(x_t-\sqrt{1-{\bar{\alpha }}_t}\tilde{z})$$
到这里，我们发现 $q(x_{t-1}|x_t)$ 的均值由已知的 $x_t$ 和模型预测出的噪声 $\tilde{z}$ 组成，而 $q(x_{t-1}|x_t)$ 的方差则由超参数计算得出。从而我们就得到了想要的 $q(x_{t-1}|x_t)$。

继续整理一下：
$$\begin{array}{rl}\mu & =\frac{x_t}{\sqrt{{\alpha }_t}}\cdot \frac{{\alpha }_t-{\bar{\alpha }}_t+{\beta }_t}{1-{\bar{\alpha }}_t}-\frac{\tilde{z}}{\sqrt{{\alpha }_t}}\cdot \frac{{\beta }_t}{1-\sqrt{1-{\bar{\alpha }}_t}}\\ & =\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}\tilde{z})\end{array}$$
即最终有：
$$q(x_{t-1}|x_t)\sim \mathcal{N}(\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}\tilde{z}),\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_{t}I)$$
最终，我们得到 reverse 过程的公式 $x_{t-1}=f(x_t,\tilde{z})$：
$$\begin{gathered} x_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}\tilde{z})+\sqrt{\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t}z \\ \tilde{z}=\text{UNet}(x_t,t),z\sim \mathcal{N}(0,I) \end{gathered}$$
这里有个问题：为什么每次计算上一步 $x_{t-1}$ 时，还要额外采样一个高斯噪声 $z$ 加上去呢？其实，这是为了给 reverse 的过程带来一定的随机性，就像 LLM 每次采样也不一定是采概率最高的 token 一样，这也是对布朗运动的一种模拟。

总结

diffusion 过程

单步扩散：
$$\begin{gathered} x_t=\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{z}_t \\ {x}_t\sim \mathcal{N}(\sqrt{{\alpha }_t}{x}_{t-1},(1-{\alpha }_t)I) \end{gathered}$$
一步到位：
$$\begin{gathered} x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}z \\ {x}_t\sim \mathcal{N}(\sqrt{{\bar{\alpha }}_t}{x}_0,\sqrt{1-{\bar{\alpha }}_t}z) \end{gathered}$$
diffusion 的过程就是一个退化的过程，不断地给真实图片加噪声，直到得到完全的高斯噪声图片。其中每一步都是在上一步的基础上加了一定量的噪声，这个 “量” 由一系列超参数 ${\beta }_{\mathrm{1...}T}$ 控制。

在训练时，每次都单步扩散得到下一步的噪声图效率较低。结合高斯分布的性质，我们可以通过一步到位的方法获取第 $t$ 步的噪声图。

由于这两种表示中，$x_t$ 都中都包含一个高斯分布项 $z/z_t$ ，所以 $x_t$ 也可以写为高斯分布的形式。

reverse过程

$$x_{t-1}\sim \mathcal{N}(\mu ,{\sigma }^{2}I)$$

其方差和均值分别为：
$$\begin{gathered} {\sigma }^2=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t \\ \begin{array}{rl}\mu & =\sqrt{{\alpha }_t}\frac{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\\ & =\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{{\sqrt{1-\bar{\alpha }}}_t}\tilde{z})\end{array} \end{gathered}$$
从而有：
$$x_{t-1}=\frac{1}{\sqrt{{\alpha }_t}}(x_t-\frac{{\beta }_t}{\sqrt{1-{\bar{\alpha }}_t}}\tilde{z})+\sqrt{\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t}z$$
reverse 的过程，虽然我们可以直接根据 diffusion 过程中一步到位的公式，从 $x_T$ 一步到 $x_0$ 。但从一个完全高斯分布一步到真实图片无疑是很困难的，这样得到的生图效果并不好。因此我们这里一步一步 $x_T,x_{T-1},\dots ,x_0$ ，从 $x_T$ 到 $x_0$。

这就需要我们推导出从后到前相邻两步的关系，即 $q(x_t|x_{t-1})$ ，或者表达成 $x_{t-1}=f(x_t,\tilde{z}),\tilde{z}=\text{UNet}(x_t,t)$ 。这里我们将 $q(x_{t-1}|x_t)$ 描述为一个高斯分布，然后分别推导出其方差和均值。其方差就是完全由超参数 $\alpha ,\beta$ 组成的一个式子。

而其均值可以写作两种形式，第一种是写成 $x_t$ 和 $x_0$ 的线性组合的形式，注意这里的 $x_0$ 当然不是真实的 $x_0$ （真实的 $x_0$ 是我们的终极目标），而是我们通过倒腾一步到位 diffusion 公式，根据 $x_t$ 和 UNet 的预测值 $\tilde{z}$ 计算得到。也就是说，这里的 $x_0$ 是当前步的估计值。从这个角度来看，每一步都会估计一个 $x_0$ ，然后乘上权重，到最后我们估计的 $x_0$ 就是每一步估计 $x_{0(t)}$ 的加权和。当然，越靠近最终结果（$t$ 越小）的时候，噪声越小，我们估计得越准，其权重也越大。

如果我们对该形式进一步整理，就可以得到第二种形式，为 $x_t$ 和噪声估计值 $\tilde{z}$ 的加权和的形式。从这个角度来看，符合我们对去噪过程的直观理解，从 $x_t$ 中减去当前步的预测噪声 $\tilde{z}$ 。

当然，我们将 $q(x_{t-1}|x_t)$ 作为高斯分布的均值方差求出之后，也能将它写成一个高斯分布的式子，即式 (25)。

最后，我们对照一下 DDPM 原文中给出的训练和采样公式。可以看到，和我们上面推导的结果是一致的。

这里有个点提一下，在采样过程中，每一步都会加一个额外的噪声 ${\sigma }_{t}z$ ，这里的 ${\sigma }_t$ 就是我们上面推导出的 $q(x_{t-1}|x_t)$ 的标准差。我们之前提到，这是为了给采样过程带来一点随机性，所以注意在采样算法中，最后一步 $x_0$ 是不加噪声的，因为 $x_0$ 就已经是我们最终要的生图结果。另外，在实际采样中，这个 ${\sigma }_t$ 是噪声方差的一个上界，这个值也不必完全与推导结果加的一样，比他小甚至为 0，也是可以的。