【Remix】拆解 DDIM 论文【扩散模型加速采样】

原文链接：[2010.02502] Denoising Diffusion Implicit Models (arxiv.org)

技术上，DDIM 提出

1. 直接拿到训好的 DDPM
2. 从 DDPM 的生成过程中删掉随机项
3. 在此基础上大幅减少采样步数

理论上，DDIM 指出

1. DDPM 训练中使用的不同权重相当于不同参数的 DDIM
2. 跳步采样的 loss function 与 DDPM 一致

实验上，DDIM 证明

1. DDIM 不论是普通采样还是快速采样，都有基本超越 DDPM 的效果
2. 对于不同的随机项选取，DDIM （即去掉随机项）基本上是最优的
3. DDIM 能保持 sample 的连续性

---

记号约定

省略向量符号。$x_0$ 表示真实数据，$\mu$ 为其分布。

前作回顾：NCSN

论文链接：[1907.05600] Generative Modeling by Estimating Gradients of the Data Distribution (arxiv.org)

1. 通过加噪，把低维流形上的分布扩散到全空间，并填充不同的高密度区域之间的空隙。
   • 取 $\tilde{x}\sim\mathcal{N}(x_0, \sigma^2\bm{I})$ ，则 $\tilde{x}$ 有概率密度函数 $q_\sigma$ 。
2. Denoising Score Matching 揭示了 $x_0$ 的 MSE 最优估计和 Stein score 的联系。
   • 根据正态分布的性质，$\newcommand{\mean}{\mathbb{E}}\mean[x_0 \mid \tilde{x}]=\tilde{x}+\sigma^2\nabla\log q_\sigma(\tilde{x})$ 。据此有 loss ：
     $$\frac{1}{2}\mean_{x_0\sim\mu,\ \tilde{x}\sim\mathcal{N}(x_0, \sigma^2\bm{I})}\left\|s_\sigma(x_0; \theta)+\frac{\tilde{x}-x_0}{\sigma^2}\right\|^2$$

前作回顾：DDPM

论文链接：[2006.11239] Denoising Diffusion Probabilistic Models (arxiv.org)

为了和 NCSN 一致，对论文里的记号做一些代换：令 $x_t=\color{red}{x}_t/\sqrt{\bar{\alpha}_t}, \sigma_t=\sqrt{1-\bar{\alpha}_t^{-1}}$ 。其中 $\color{red}{x}$ 为原文的记号，$x$ 为新记号。部分代换的中间结果放在下面以供读者验算：

$$\bar{\alpha}_t=\frac{1}{1+\sigma_t^2},\ \alpha_t=\frac{\bar{\alpha}_t}{\bar{\alpha}_{t-1}}=\frac{1+\sigma_{t-1}^2}{1+\sigma_t^2},\ \beta_t=1-\alpha_t=\frac{\sigma_t^2-\sigma_{t-1}^2}{1+\sigma_t^2},\ \tilde{\beta}_t=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_t}\beta_t=\frac{\sigma_{t-1}^2(\sigma_t^2-\sigma_{t-1}^2)}{\sigma_t^2}\bar{\alpha}_{t-1}$$

1. 继承了扩散过程 $x_t\sim\mathcal{N}(x_{t-1}, (\sigma_t^2-\sigma_{t-1}^2)\bm{I}), t \in [T]$ ，在此基础上改进了反向过程的训练目标。
   • 逆向推断的难度全在于初始分布的影响。
   • 如果已知 $x_0$ ，反向过程有封闭形式 $x_{t-1}\sim\mathcal{N}(x_0+(\sigma_{t-1}^2/\sigma_t^2)(x_t-x_0), \sigma_{t-1}^2(1-\sigma_{t-1}^2/\sigma_t^2)\bm{I})$ 。
   • 取 $\sigma_T \gg \mean\|x_0\|$ ，采样时可近似 $x_T \sim \mathcal{N}(0, \sigma_T^2\bm{I})$ 。
   • 使用神经网络估计 $\epsilon_t(x_t; \theta) \approx (x_t-x_0)/\sigma_t$ ，取 $\hat{x}_0=x_t-\sigma_t\epsilon_t(x_t; \theta)$ 代回原式。
   • 目标函数是 $\|\epsilon_t(x_t; \theta)-(x_t-x_0)/\sigma_t\|^2$ 的加权和。
2. 理论上，证明了一种对 $t$ 加权的训练目标是上述过程定义的 ELBO 。实验上采用不加权的训练目标，结果也不错。

---

仍然将论文里的记号与 NCSN 对齐：$x_t=\color{red}{x}_t/\sqrt{\alpha_t}, \sigma_t^2=1/\alpha_t-1, s_t^2=\color{red}{\sigma}_t^2/\alpha_{t-1}$ ，其中 $\color{red}{x}, \color{red}{\sigma}$ 为原文的记号，$x, \sigma$ 为新记号。

反向过程的随机项：从删除到调整

当 DDPM 训练得很好时，$\epsilon_t(x_t) \approx \mean[(x_t-x_0)/\sigma_t \mid x_t]=-\sigma_t\nabla\ln q_t(x_t)$ ，其中 $q_t$ 是 $x_t$ 的概率密度函数，这只与 $x_0 \sim \mu$ 以及 $x_t\sim\mathcal{N}(x_0, \sigma_t^2\bm{I})$ 有关，而与具体的扩散过程无关。

最简版本

• 使用 DDPM 预训练的 $\epsilon_t$ ，只修改生成算法。
• 假如已知 $x_0$ ，在保持 $x_t\sim\mathcal{N}(x_0, \sigma_t^2\bm{I})$ 的前提下直接删掉随机项，由此可以构造：$x_{t-1}=x_0+(\sigma_{t-1}/\sigma_t)(x_t-x_0)$ 。
• 依然估计 $x_T \sim \mathcal{N}(x_0, \sigma_T^2\bm{I})$ ，并用 $\hat{x}_0=x_t-\sigma_t\epsilon_t(x_t)$ 估计未知的 $x_0$ 。
• 效果比 DDPM 好。

理论版本

• 修改 DDPM 对 $x_t$ 的定义，使用一个非马尔可夫的过程来产生 $x_t$ 。
• 训练时，已知 $x_0$ ，从 $x_T \sim \mathcal{N}(x_0, \sigma_T^2\bm{I})$ 开始反向构造 $x_t$ ，保持 $x_t\sim\mathcal{N}(x_0, \sigma_t^2\bm{I})$ 不变：$x_{t-1} \sim \mathcal{N}(x_0+(\sqrt{\sigma_{t-1}^2-s_t^2}/\sigma_t)(x_t-x_0), s_t^2\bm{I})$ 。
• 采样时，在上述的式子上用 $\hat{x}_0=x_t-\sigma_t\epsilon_t(x_t; \theta)$ 估计 $x_0$ ，最后一步取 $x_0 \sim \mathcal{N}(x_1-\sigma_1\epsilon_1(x_1; \theta), s_1^2\bm{I})$ 。
• 对其计算 ELBO 说明不同的 $s_t$ 对应不同的 $\|\epsilon_t(x_t; \theta)-(x_t-x_0)/\sigma_t\|^2$ 权重，这验证了 DDPM 原来的不加权训练目标的有效性。
• 对 $s_t^2=\eta\sigma_{t-1}^2(1-\sigma_{t-1}^2/\sigma_t^2)$ 做了实验，其中 $\eta$ 从 0 (DDIM) 到 1 (DDPM) ，证明 $\eta=0$ 是最佳的。

加速采样和它的无聊理论

最简版本

• 采样过程为 $x_{t-1}=x_t-(\sigma_t-\sigma_{t-1})\epsilon_t(x_t)$ ，写成微分方程就是 $\mathrm{d}x_t=\epsilon_t(x_t)\mathrm{d}\sigma_t$ 。

• 所以可以取子序列 $0<\tau_1<\tau_2<\cdots<\tau_S=T$ 模拟：$x_{\tau_{i-1}}=x_{\tau_i}-(\sigma_{\tau_i}-\sigma_{\tau_{i-1}})\epsilon_{\tau_i}(x_{\tau_i})$ 。

理论版本

$$x_{\tau_{i-1}} \sim \mathcal{N}(x_0+(\sqrt{\sigma_{\tau_{i-1}}^2-s_{\tau_i}^2}/\sigma_{\tau_i})(x_{\tau_i}-x_0), s_{\tau_i}^2\bm{I})$$

• 为了和 DDPM 挂钩，需要保持 ELBO 的每一项都有一定权重。
• 所以用了个小技巧：直接把不在 $\tau$ 里的项和 $x_0$ 挂钩起来。
• 对于 $t \notin \tau$ 构造前向 $x_t \sim \mathcal{N}(x_0, \sigma_t^2\bm{I})$ 以及不参与采样的“反向”$p(x_0 \mid x_t)=\mathcal{N}(x_0; x_t-\sigma_t\epsilon_t(x_t; \theta), s_t^2\bm{I})$ 。