Denoising Diffusion Implicit Models

• 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{\alpha_t \bm{x}_0}}{\sqrt{1 - \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{\alpha_t \bm{x}_{\theta}(\bm{x}_t, t)}}{\sqrt{1 - \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{\alpha_t \bm{x}_{\theta}(\bm{x}_t, t)}}{\sqrt{1 - \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{\alpha_{\tau_i}}\bm{x}_0, (1 - \alpha_{\tau_i})\bm{I}), \]

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

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

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

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

代码

official