Denoising Diffusion Implicit Models

目录

• 概
• Motivation
• 代码

Song J., Meng C. and Ermon S. Denoising diffusion implicit models. In International Conference on Learning Representations (ICLR), 2021.

概

DDIM 从另一种观点 理解 DDPM, 并以此推导出更加快速的采样方式.

Motivation

• DPM 的前向过程一般是:

  $$q(\bm{x}_{1:T}|\bm{x}_0) = \prod_{t=1}^T q(\bm{x}_t|\bm{x}_{t-1}), \: q(\bm{x}_t|\bm{x}_{t-1}) = \mathcal{N}(\sqrt{\beta_t}\bm{x}_{t-1}, (1 - \beta_t) \bm{I}).$$

• 通过 $q(\bm{x}_t|\bm{x}_{t-1})$ 我们可以得到

  $$q(\bm{x}_t|\bm{x}_0) = \mathcal{N}(\bm{x}_t; \sqrt{\bar{\alpha}_t} \bm{x}_0, (1 - \bar{\alpha}_t) \bm{I}), \: \bar{\alpha}_t = \prod_{\tau=1}^t \alpha_{\tau}, \alpha_{\tau}:= 1 - \beta_{\tau}.$$

• 然后结合二者, 我们可以通过贝叶斯公式得到:

  $$\tag{1} \begin{array}{ll} q(\bm{x}_{t-1}|\bm{x}_t, \bm{x}_0) &= \frac{q(\bm{x}_{t}|\bm{x}_{t-1}) q(\bm{x}_{t-1}|\bm{x}_0)}{q(\bm{x}_t|\bm{x}_0)} \\ &= \mathcal{N}(\bm{x}_t; \frac{\sqrt{\bar{\alpha}_{t-1}}(1 - \alpha_t)\bm{z}_0 + \sqrt{\alpha_t} (1 - \bar{\alpha}_{t-1})\bm{z}_t}{1 - \bar{\alpha}_t}, \frac{(1 - \alpha_t)(1 - \bar{\alpha_{t-1}})}{1 - \bar{\alpha_t}}\bm{I}). \end{array}$$

• 根据 这儿 我们知道, 实际上, 我们要做的就是建模 $p_{\theta}(\bm{x}_{t-1}|\bm{x}_t)$ 来近似 (1)：

  $$\tag{2} \begin{array}{ll} &= \underbrace{\mathbb{E}_{q(\bm{x}_{1}|\bm{x}_0)}[\log p(\bm{x}_0|\bm{x}_1; \theta)]}_{-\mathcal{L}_{1} \text{ reconstruction term}} \\ &\quad - \underbrace{\mathbf{KL}(q(\bm{x}_T|\bm{x}_{0}) \| p(\bm{x}_T))}_{\mathcal{L}_{0} \text{ prior matching term }} \\ &\quad \quad - \underbrace{\sum_{t=2}^{T} \mathbb{E}_{q(\bm{x}_{t}|\bm{x}_0)} \Big \{\mathbf{KL}(q(\bm{x}_{t-1}|\bm{x}_{t}, \bm{x}_0) \| p(\bm{x}_{t-1}|\bm{x}_{t}; \theta) \Big \}}_{\mathcal{L}_{2:T} \text{ consistency term } } \\ \end{array}$$

• 更特别地, 现在 DDPM 更习惯用一种更为简单的形式 (忽略系数):

  $$\mathbb{E}_{q(\bm{x}_t|\bm{x}_0)}\Big\{ \|\epsilon_{\theta}(\bm{x}_t, t) - \epsilon_t\|^2 \Big\},$$

  或者

  $$\mathbb{E}_{q(\bm{x}_t|\bm{x}_0)}\Big\{ \|\bm{x}_{\theta}(\bm{x}_t, t) - \bm{x}_t\|^2 \Big\}.$$

• 现在问题来了, 此时我们只要知道 $q(\bm{x}_t| \bm{x}_0)$ 就可以了, 该问题就可以求解了, 完全没必要按照 $q(\bm{x}_t|\bm{x}_{t-1}) \rightarrow q(\bm{x}_t|\bm{x}_0) \rightarrow q(\bm{x}_{t-1}|\bm{x}_t, \bm{x}_0)$ 的流程一步一步来. 我们可以:

  1. 先假定:
  $$q(\bm{x}_t|\bm{x}_0) = \mathcal{N}(\bm{x}_t; \sqrt{\bar{\alpha}_t} \bm{x}_0, (1 - \bar{\alpha}_t) \bm{I}),$$
  2. 假设:
  $$q(\bm{x}_{t-1}|\bm{x}_t, \bm{x}_0) = \mathcal{N}(\bm{x}_{t-1}; \kappa_t \bm{x}_t + \lambda_t \bm{x}_0, \sigma_t^2 \bm{I}),$$

  且待定系数 $\kappa, \lambda, \sigma$ 应当使得关系

  $$\tag{3} \int_{\bm{x}_t} q(\bm{x}_{t-1}|\bm{x}_t, \bm{x}_0) q(\bm{x}_t|\bm{x}_0) = q(\bm{x}_{t-1}|\bm{x}_0)$$

  成立.

• 首先, 容易观察到按照 $\bm{x}_t \sim q(\bm{x}_t|\bm{x}_0), \bm{x}_{t-1} \sim q(\bm{x}_{t-1}|\bm{x}, \bm{x}_0)$ 得到的 $\bm{x}_{t-1}$ 满足:

  $$\begin{array}{ll} \bm{x}_{t-1} &=\kappa_t \bm{x}_{t} + \lambda_t \bm{x}_0 + \sigma_t \epsilon \\ &=\kappa_t [\sqrt{\bar{\alpha}_t} \bm{x}_0 + \sqrt{(1 - \bar{\alpha}_t)} \bm{\epsilon}'] + \lambda_t \bm{x}_0 + \sigma_t \bm{\epsilon} \\ &=(\kappa_t \sqrt{\bar{\alpha}_t} + \lambda_t) \bm{x}_0 + \sqrt{\kappa_t^2(1 - \bar{\alpha}_t) + \sigma_t^2} \bm{\epsilon} \\ &\sim\mathcal{N}(\bm{x}_{t-1}; (\kappa_t \sqrt{\bar{\alpha}_t}, + \lambda_t) \bm{x}_0, \kappa_t^2(1 - \bar{\alpha}_t) + \sigma_t^2 \bm{I}). \\ \end{array}$$

• 所以, 为了保证 (3) 严格成立, 需要

  $$\kappa_t \sqrt{\bar{\alpha}_t} + \lambda_t = \sqrt{\bar{\alpha}_{t-1}}, \\ \kappa_t^2(1 - \bar{\alpha}_t) + \sigma_t^2 = 1 - \bar{\alpha}_{t-1}.$$

• 三个未知参数, 两个方差, 将 $\sigma_t$ 看成可变参数, 可得:

  $$\lambda_t = \sqrt{\bar{\alpha}_{t-1}} - \sqrt{\bar{\alpha}_t} \sqrt{\frac{1 - \bar{\alpha}_{t-1}- \sigma_t^2}{1 - \bar{\alpha}_{t}}}, \\ \kappa_t = \sqrt{\frac{1 - \bar{\alpha}_{t-1}- \sigma_t^2}{1 - \bar{\alpha}_{t}}}.$$

• 此时:

  $$q_{\sigma}(\bm{x}_{t-1}|\bm{x}_t, \bm{x}_0) = \mathcal{N}(\sqrt{\bar{\alpha}_{t-1}}\bm{x}_0 + \sqrt{1 - \bar{\alpha}_{t-1} - \sigma_t^2} \cdot \frac{\bm{x}_t - \sqrt{\bar{\alpha}_t} \bm{x}_0}{\sqrt{1 - \bar{\alpha}_t}}, \sigma_t^2 \bm{I}).$$

• 于是, 我们可以通过控制 $\sigma_t$ 来获得更多的可能性 (DDPM 可以看成是一个特例), 当然, 我们的近似 $p_{\theta}(\bm{x}_{t-1}|\bm{x}_t)$ 也要相应的改变:

  $$p_{\sigma}(\bm{x}_{t-1}|\bm{x}_t) = \mathcal{N}(\sqrt{\bar{\alpha}_{t-1}}\bm{x}_{\theta}(\bm{x}_t, t) + \sqrt{1 - \bar{\alpha}_{t-1} - \sigma_t^2} \cdot \frac{\bm{x}_t - \sqrt{\bar{\alpha}_t} \bm{x}_{\theta}(\bm{x}_t, t)}{\sqrt{1 - \bar{\alpha}_t}}, \sigma_t^2 \bm{I}).$$

• 不过, 因为我们在训练的时候一般用的是简化的损失, 所以实际上即使我们用了不同的 $\sigma_t$, 最后的简化的损失也依旧会退化为 DDPM 所用的简化损失的形式. 这意味着, 我们采用一般的 DDPM 的损失训练最优后, 对于其它 $\sigma_t$ 的损失也是最优的 (理想情况下). 于是, 这意味着我们才反向采样的时候可以依赖任意的 $\sigma_t$ !

• 我们可以发现, 反向采样过程为:

  $$\bm{x}_{t-1} \leftarrow \sqrt{\bar{\alpha}_{t-1}}\bm{x}_{\theta}(\bm{x}_t, t) + \sqrt{1 - \bar{\alpha}_{t-1} - \sigma_t^2} \cdot \frac{\bm{x}_t - \sqrt{\bar{\alpha}_t} \bm{x}_{\theta}(\bm{x}_t, t)}{\sqrt{1 - \bar{\alpha}_t}} + \sigma_t \epsilon_t.$$

• 原先采样需要严格按照 $[1, \ldots, T]$ 的倒序, 现在我们可以用一个 sub-sequence $[\tau_1, \ldots, \tau_{\tau_S}]$ 来替代. 此时我们实际上假设我们的前向是服从:

  $$q(\bm{x}_{\tau_i}|\bm{x}_0) = \mathcal{N}(\sqrt{\bar{\alpha}_{\tau_i}}\bm{x}_0, (1 - \bar{\alpha}_{\tau_i})\bm{I}),$$

  此时我们的推导依旧是成立的. 但是我们可以用更少的次数采样.

• 当然, 此时, 有一些需要注意的点:

  1. 用更少的次数采样, 容易造成图像产生更多的噪声;
  2. 此时用更小的 $\sigma_t$ 可以缓解这一问题.

• 实际上, DDIM 所介意的是 $\sigma_t=0$.

代码

official