ddpm (新)

太多博客或者学习资料喜欢一上来就从 ELBO 等角度去推导 loss ，但是对于掌握一个模型，如果我们连模型在做什么都没有一个宏观的把握，只是一味的陷入类似于收敛性证明的漩涡，这么做我认为是丢掉了重点。之所以大家喜欢这么做的原因，主要是原论文的出发点就是从 ELBO 角度去把东西一步步推出来的。我认为这样无可非议，因为人家是第一个发明这个算法的人，需要足够的严谨性和逻辑自洽。但是我们作为计算机兴趣研究者，在已知结论的前提下，是不需要走作者的老路。而应该用一种更宏观的视角，去化繁为简，把自己从繁重的证明中剥离出来，从上帝视角看看 ddpm 到底在做一个什么事情，怎么做的，原理是什么，就足够了。

我们毕竟不是数学专业，当然有更聪明的人能从公式推导见微知著，完全掌握。但我相信更多的人往往是陷入漩涡后就无法脱身，并且离场的时候还给 ddpm 打下了 “数学门槛高” 的标签，从而对后来者产生了不正确的观念影响。

好了，废话说了那么多，下面我将从一个宏观的视角解释 ddpm 。

‍

1 前置知识

1. 一些约定俗称的符号与规定：

   • 在 ddpm 中，$p$ 一般表示去噪过程，$q$ 一般表示加噪过程。一般课内的概率论常用公式推导，我都用的是 $p$ 。

   • $p, q$ 表示的是概率值，不是一个随机变量或者分布，所以不能写 $p/q \sim N(\mu, \sigma^2 I)$。$x_t$ 表示的是一个分布。以下是正确的写法：

     • $x_t \sim N(\mu, \sigma^2 I)$
     • $q(x_t | x_0) = N(x_t; \mu, \sigma^2 I) = \frac{1}{\sqrt{2\pi}\sigma}\exp(...)$
     • $x_t = \mu + \sigma\epsilon$

   • 当然这个错误我也可能经常犯，符号乱用，但是自己清楚就行。

2. $p(x_{0:T}) = p(x_0, x_1, \cdots, x_T)$​

3. $p(x_{0:T}) = ?$

   • 先从一个特例推起：

     $$
     \begin{aligned}
     &p(x_0, x_1, x_2) = p(x_0, x_1 | x_2)p(x_2) \
     \because & p(x_0, x_1) = p(x_0 | x_1) | p(x_1) \
     \therefore & p(x_0, x_1 | x_2) = p(x_0 | x_1, x_2)p(x_1 | x_2) \qquad \text{copy } x_2 \text{ directly after the conditon} \
     \because & \text{Markov chain}, \therefore p(x_0 | x_1, x_2) = p(x_0 | x_1) \
     \therefore & p(x_0, x_1, x_2) = p(x_0 | x_1)p(x_1 | x_2)p(x_2) \
     \end{aligned}
     $$

   • 所以有：

     $$
     \begin{aligned}
     p(x_{0:T})=p(x_{T})\prod_{t=1}^{T}p\left(x_{t-1}\left|\right|x_{t}\right) \
     p(x_{0:T})=p(x_{0})\prod_{t=1}^{T}p\left(x_{t}\left|\right|x_{t-1}\right) \
     \end{aligned}
     $$

4. $q(x_{1:T} | x_0) = ?$

   $$
   \begin{aligned}
   \because & q(x_{1:T}) = q(x_{1})\prod_{t=2}^{T}q\left(x_{t}\left|x_{t-1}\right.\right) \
   \therefore & q(x_{1:T} | x_0) = q(x_1 | x_0)\prod_{t=2}^T q(x_t | x_{t-1}, x_0) = q(x_1 | x_0) \prod_{t=2}^T q(x_t | x_{t-1}) = \prod_{t=1}^T q(x_t | x_{t-1})
   \end{aligned}
   $$

5. Jensen 不等式，对于上凸函数，有：$f(E(x)) \ge E(f(x))$

6. 最大似然不等式推导

   $$
   \begin{aligned}
   \log p(x_0) &= \log \int p(x_{0:T})d(x_{1:T}) \qquad \text{marginal distribution} \
   &= \log \int q(x_{1:T} | x_0)\frac{p(x_{0:T})}{q(x_{1:T} | x_0)}d(x_{1:T}) \qquad \text{Since the process of adding noise is a known condition, it must be used} \
   &= \log E_{q(x_{1:T} | x_0)}\left[ \frac{p(x_{0:T})}{q(x_{1:T} | x_0)} \right] \qquad x_0\text{ is actually an invariant, so we don't have to worry about it when we change expectations} \
   &\ge E_{q(x_{1:T}|x_0)}\log \frac{p(x_{0:T})}{q(x_{1:T} | x_0)} \qquad \text{Jensen inequality} \
   \end{aligned}
   $$

‍

2 加噪

我们直接从结论出发，ddpm 论文作者，对干净的样本 $x_0$ ，是做了一个加噪的过程，具体用下列公式表达：

$$
x_t = \sqrt{1 - \beta_t}x_{t-1} + \sqrt{\beta_t}\epsilon_t, \quad \epsilon_t \sim N(0, I)
$$

$\beta_t$ 是人为定义好的常量，随时间增大而增大，可以认为是噪声级别，时间越靠后扰动程度就越大。

下面这是个递推公式，所以我们可以一步到位：

令 $\alpha_t = 1 - \beta_t$，$\overline{\alpha_t} = \alpha_t\alpha_{t-1}\cdots\alpha_1$

先推一个特例：

$$
\begin{aligned}
&x_3 = \sqrt{1 - \beta_3}x_2 + \sqrt{\beta_3}\epsilon_3 \
&x_2 = \sqrt{1 - \beta_2}x_1 + \sqrt{\beta_2}\epsilon_2 \
&x_1 = \sqrt{1 - \beta_1}x_0 + \sqrt{beta_1}\epsilon_1 \
\therefore x_3 &= \sqrt{1 - \beta_3}(\sqrt{1 - \beta_2}x_1 + \sqrt{\beta_2}\epsilon_2) + \sqrt{\beta_3}\epsilon_3 \
&= \sqrt{1 - \beta_3}\sqrt{1 - \beta_2}x_1 + \sqrt{1 - \beta_3}\sqrt{\beta_2}
\epsilon_2 + \sqrt{\beta_3}\epsilon_3 \
&= \sqrt{1 - \beta_3}\sqrt{1 - \beta_2}(\sqrt{1 - \beta_1}x_0 + \sqrt{\beta_1}\epsilon_1) + \sqrt{1 - \beta_3}\sqrt{\beta_2}\epsilon_2 + \sqrt{\beta_3}\epsilon_3 \
&= \sqrt{1 - \beta_3}\sqrt{1 - \beta_2}\sqrt{1 - \beta_1}x_0 + \sqrt{1 - \beta_3}\sqrt{1 - \beta_2}\sqrt{\beta_1}\epsilon_1 + \sqrt{1 - \beta_3}\sqrt{\beta_2}\epsilon_2 + \sqrt{\beta_3}\epsilon_3 \
&= \sqrt{\alpha_3\alpha_2\alpha_1}x_0 + \sqrt{\alpha_3\alpha_2\beta_1}\epsilon_1 + \sqrt{\alpha_3\beta_2}\epsilon_2 + \sqrt{\beta_3}\epsilon_3
\end{aligned}
$$

观察到其中，上式可以写成一个正态分布的形式：

$$
\begin{aligned}
\because
&\sqrt{\alpha_3\alpha_2\beta_1}\epsilon_1 \sim N(0, \alpha_3\alpha_2\beta_1) \
&\sqrt{\alpha_3\beta_2}\epsilon_2 \sim N(0, \alpha_3\beta_2) \
&\sqrt{\beta_3}\epsilon_3 \sim N(0, \beta_3) \
\therefore
&\sqrt{\alpha_3\alpha_2\beta_1}\epsilon_1 + \sqrt{\alpha_3\beta_2}\epsilon_2 + \sqrt{\beta_3}\epsilon_3 \sim N(0, (\alpha_3\alpha_2\beta_1 + \alpha_3\beta_2 + \beta_3)I) \
\because
&\alpha_3\alpha_2\beta_1 + \alpha_3\beta_2 + \beta_3 = (\alpha_3\alpha_2\beta_1 + \alpha_3\beta_2 + \beta_3)

• \alpha_1\alpha_2\alpha_3 - \alpha_1\alpha_2\alpha_3 \
  &= \alpha_3\alpha_2(\beta_1 + \alpha_1) + \alpha_3\beta_2 + \beta_3 - \alpha_1\alpha_2\alpha_3 \
  &= \alpha_3\alpha_2 + \alpha_3\beta_2 + \beta_3 - \alpha_1\alpha_2\alpha_3 \
  &= \alpha_3(\alpha_2 + \beta_2) + \beta_3 - \alpha_1\alpha_2\alpha_3 \
  &= 1 - \alpha_1\alpha_2\alpha_3 \
  \therefore
  & \sqrt{\alpha_3\alpha_2\beta_1}\epsilon_1+\sqrt{\alpha_3\beta_2}\epsilon_2+\sqrt{\beta_3}\epsilon_3 \sim N(0, (1 - \overline{\alpha_3})I)
  \end{aligned}
  $$

所以可得：

$$
q(x_3 | x_0) = N(x_3; \sqrt{\overline{\alpha_3}}x_0, (1 - \overline{\alpha_3})I)
$$

所以：

$$
\begin{aligned}
& q(x_t | x_0) = N(x_t; \sqrt{\overline{\alpha_t}}x_0, (1 - \overline{\alpha_t})I) \
& x_t = \sqrt{\overline{\alpha_t}}x_0 + \sqrt{1 - \overline{\alpha_t}}\epsilon_t, \quad \epsilon_t \sim N(0, I)
\end{aligned}
$$

一步加噪的公式被我们成功推出来了。

现在，我们只需要有一个干净样本 $x_0$ ，然后从均匀分布中采样一个时间步 $t$ ，查表得到 $\overline{\alpha_t}$ ，再从标准高斯分布中采样一个 $\epsilon$ ，就可以用上述公式把 $x_0$ 加噪加到到第 $t$ 步的时候。

而且，我们发现，前向加噪的过程，无论是每一步时间，或者一步到位，都是 $x_r$ 是服从关于 $x_s$ 的正态分布的（$s < r$）。

‍

2 去噪

既然前面我们推出来前向过程 $x_r$ 服从 $x_s$ 的正态分布（$s < r$），那我们也希望反向去噪过程也可以满足这样的性质。即 $x_s$ 服从 $x_r$ 的正态分布（当我们已知 $x_r$ 时）。

让我们来验证一下。

先把一些已知条件列在这，后续会用到：

$$
\begin{aligned}
&q(x_t | x_{t-1}) = N(x_t; \sqrt{1 = \beta_t}x_{t-1}, \beta_t I) \
&q(x_t | x_0) = N(x_t; \sqrt{\overline{\alpha_t}}x_0, (1 - \overline{\alpha_t})I) \
&p(x) = \frac{1}{\sqrt{2\pi}\sigma}\exp(-\frac{(x - \mu)^2}{2\sigma^2}) = \frac{1}{\sqrt{2\pi}\sigma}\exp\left[-\frac12 (\frac{x^2}{\sigma^2} - \frac{2\mu x}{\sigma^2} + \frac{\mu^2}{\sigma^2})\right]
\end{aligned}
$$

我们的目标是想把 $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)}
$$

Well, is hard to calculate... 无法用上我们的已知条件，ok，那就加条件！不妨去求 $q(x_{t-1} | x_t, x_0)$ ，来吧：

$$
\begin{aligned}
q(x_{t-1} | x_t, x_0) &= \frac{q(x_t | x_{t-1}, x_0)q(x_{t-1} | x_0)}{q(x_t | x_0)} \
&= \frac{q(x_t | x_{t-1})q(x_{t-1} | x_0)}{q(x_t | x_0)} \
&= ... \
&= \exp\left[-\frac12 \left( \left(\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \overline{\alpha_{t-1}}}\right)^2 x_{t-1} - \left( \frac{2\sqrt{\alpha_t}}{\beta_t}x_t + \frac{2\sqrt{\overline{\alpha_{t-1}}}}{1 - \overline{\alpha_{t-1}}}x_0 \right)x_{t-1} + C(x_t, x_0) \right)\right]
\end{aligned}
$$

$\cdots$ 过程太繁琐，没必要推了，都是些配凑法之类的。

$C(x_t, x_0)$ 是关于 $x_t, x_0$ 的表达式，太长了，我就不写出来了。

似乎，这是个正态分布？我们可以用待定系数法求一下 $\sigma^2, \mu$ 。

$$
\begin{cases}
\frac{\alpha_t}{\beta_t} + \frac{1}{1 - \overline{\alpha_{t-1}}} = \frac{1}{\sigma^2} \
\frac{2\sqrt{\alpha_t}}{\beta_t}x_t + \frac{2\sqrt{\overline{\alpha_{t-1}}}}{1 - \overline{\alpha_{t-1}}}x_0 = \frac{2\mu}{\sigma^2} \
C(x_t, x_0) = \frac{\mu^2}{\sigma^2}
\end{cases}
\Rightarrow
\begin{cases}
\mu = (\frac{2\sqrt{\alpha_t}}{\beta_t}x_t + \frac{2\sqrt{\overline{\alpha_{t-1}}}}{1 - \overline{\alpha_{t-1}}}x_0)\frac{\sigma^2}{2} \
\sigma^2 = \frac{\beta_t(1 - \overline{\alpha_{t-1}})}{1 - \overline{\alpha_t}}
\end{cases}
$$

所以，确实可以推出来，反向过程也满足正态分布转移！可以写为（为了标识我用下标 $+$ 表示刚求出来的俩待定系数）：

$$
q(x_{t-1} | x_t, x_0) = N(x_{t-1}; \mu_{+}(x_t, x_0), \sigma_+^2I)
$$

看上面这个转移方程，也就是当我知道 $x_t, x_0$ 后，就可以把 $\mu_+, \sigma_+$ 算出来，从而通过正态分布采样转移到 $x_{t-1}$ 。然后一直做下去，最终就能得到 $x_0$ 。

其实这个形式还可以改写，因为我们知道：$x_t = \sqrt{\overline{\alpha_t}}x_0 + \sqrt{1 - \overline{\alpha_t}}\epsilon_t$ ，所以 $x_0 = \frac{x_t - \sqrt{1 - \overline{\alpha_t}}\epsilon}{\sqrt{\overline{\alpha_t}}}$ ，所以把这个反代进 $\mu_+$ 的表达式里，也就是 $\mu_+$ 可以用 $x_t, \epsilon$ 来表示：$\mu_+(x_t, \epsilon)$ 。

好了，在整个反向推理过程中，我们唯独不知道的是 $\epsilon$ 。不知道怎么办，用模型预测它就行了。

所以反向推理的过程是：

用模型先预测出 $\epsilon$ ，然后从 $x_t$ 开始，求出 $\mu_+(x_t, \epsilon), \sigma_+(t)$ ，然后正态分布采样得到 $x_{t-1}$ ，以此类推最终得到 $\hat{x}_0$ ，游戏结束。

训练就是有一个干净样本 $x_0$ ，然后从均匀分布中采样一个时间步 $t$ ，查表得到 $\overline{\alpha_t}$ ，再从标准高斯分布中采样一个 $\epsilon$ ，然后用加噪公式把 $x_0$ 加噪得到 $x_t$ 。然后通过 $f_\theta(x_t, t)$ 算出预测噪声。然后算一算 $|\epsilon - f_\theta(x_t, t)|_2$ 作为 loss 。

‍

3 灵魂发问

其实从宏观上，ddpm 在做什么事情已经很赤裸裸了。

本质就是加噪与去噪，而且前后向过程都是服从正态分布转移的。

但是，我很奇怪的一点就是，我们的模型是在预测初始噪声 $\epsilon$（或者说预测 $x_0$ 也行）。

如果我们连初始图像 $x_0$ 都预测出来了，那任务不就结束了吗？？你搞这么一堆前向反向的干啥，浪费时间？

很好，我给自己的解释是，$x_0$ 像素空间编码空间太大且分布不规律，很难预测，这样做不会 work 。

ok，那我进一步发问，既然 $x_0$ 不好预测，那上面的推导其实预测初始噪声 $\epsilon$ 也是可以的，有了它直接通过前向公式 $x_t = \sqrt{\overline{\alpha_t}}x_0 + \sqrt{1 - \overline{\alpha_t}}\epsilon_t$ 反算出 $x_0$ 不就行了吗？

这里我给出的解释是这样的：

我们知道，加噪过程中的初始噪声 $\epsilon$ 是采样得到的，这意味着它并不是一个固定的值，而是服从一个分布。

而我们模型做预测，每次只能预测出一个固定的噪声值，所以如果直接把这个固定的噪声值就当作当初我们前向加噪用的初始噪声，是不对的。

因为一个是值，一个是服从分布。

how to solve it ？虽然我们知道初始噪声不是一个固定值，但是我们知道它服从 $N(0, I)$ 啊，也就是说它的均值为 0！

而反向很多次的迭代去噪过程就很好了利用了这个性质，因为迭代过程中预测了很多次初始噪声，我们有理由相信，当去噪过程次数足够多的时候，模型预测的这么多个初始噪声的均值也为 0！

这就对齐了初始噪声不是固定值与模型预测的初始噪声是固定值的 gap。

（上述这个解释是我自己口胡的，我认为在直觉上能给我们一个很好的启发。严谨的数学分析需要后续用 SDE/ODE 数学公式去分析前向反向的公式，以后会见到的）