Denoising Diffusion Implicit Models
概
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)\) 的流程一步一步来. 我们可以:
- 先假定:
\[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}), \]- 假设:
\[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}), \]此时我们的推导依旧是成立的. 但是我们可以用更少的次数采样.
-
当然, 此时, 有一些需要注意的点:
- 用更少的次数采样, 容易造成图像产生更多的噪声;
- 此时用更小的 \(\sigma_t\) 可以缓解这一问题.
-
实际上, DDIM 所介意的是 \(\sigma_t=0\).
