Diffusion Model 的原理推演

一直以来科研人员对生成式模型进行不断探索，试图建立从高斯噪音到样本数据分布的映射。最近出现的 Diffusion Model 提出一个新的角度：将生成问题看作一个扩散过程的逆问题。定义一个特殊的扩散过程 $q(x_t|x_{t-1}) = \mathcal{N}(\sqrt{1-\beta_t} x_{t-1}, \beta_t I), x_0 \sim q(x_0)$ ，我们可以通过采样得到 $\{x_0, …, x_T\}$。假设扩散过程足够充分，可以认为 $x_T \sim \mathcal{N}(0, I)$。于是我们通过扩散过程建立起$q(x_0)$分布到高斯噪音的一个联系。可以看到扩散过程的逆过程正好就是一个生成过程，因此问题便转化为了试图求解扩散过程的逆$q(x_{t-1}|x_t)$。

定义参数为$\theta$的模型 $p_\theta(x_{t-1} | x_t)$ 为目标模型，便可通过$x_{t-1} \sim p_\theta(x_{t-1} | x_t)$的生成过程得到 $\{x_T, …, x_0\}$，记以此得到的$x_0$分布为 $p_\theta(x_0)$。通过最大化原始数据 $q(x_0)$ 在反向生成分布中的似然 $\mathbb{E}_{q(x_0)}\log p_\theta(x_0)$，我们便可以到一个适当的参数$\theta$用于生成接近$q(x_0)$。

1. 推导主干

1. 从最大似然$L_{CE}$到变分下界$L_{VLB}$：
   $L_{CE} = \mathbb{E}_{q(x_0)} \left[-\log p_{\theta}(x_0)\right]$，根据Lemma 1可知$L_{CE} \leq L_{VLB} = \mathbb{E}_{q(x_{0:T})}\left[\log \frac{q(x_{1:T} | x_0)}{p_\theta(x_{0:T})}\right].$

2. 将$L_{VLB}$展开为不同时间步Loss：
   根据 Lemma 2 将 $L_{VLB}$ 展开为 $L_0 + L_{1:T-1} + L_T$，其中

   • $L_T = D_{KL} (q(x_T|x_0)|p_\theta(x_T))$
   • $L_0 = - \log p_\theta(x_0|x_1)$
   • $L_t = D_{KL} (q(x_t|x_{t+1}, x_0) | p_\theta(x_t|x_{t+1}))$
     其中$L_T$为常数，真正可以被学习的是$L_t$和$L_0$。

3. 分时间步Loss的闭式解：
   将$p_{\theta}(x_{t-1}|x_t)$ 表征为 $\mathcal{N}(\mu_{\theta}(x_t, t), \Sigma_{\theta}(x_t, t))$，由 Lamma 3 可知 $q(x_{t-1}|x_t, x_0)$=$\mathcal{N}(\mu(x_t, x_0), \sigma_t^2 I)$ 也是高斯分布，$L_t$ 即为两个高斯分布KL Divergence的期望。$L_{t-1} = \mathbb E_{x_0, x_t} \left[\frac{1}{2\|\Sigma_\theta\|^2}\|\mu_{\theta}(x_t, t) - \mu(x_t, x_0)\|^2\right].$

4. 表达为噪音预测 $\epsilon_\theta$ 的Loss：
   重写$\mu_{\theta}(x_t, t)=\frac{1}{\sqrt{\alpha_t}} \left(x_t - \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon_{\theta}(x_t, t)\right)$，最终优化目标$L_i$成为 $\mathbb E_{x_0, \epsilon_t} \left[ \frac{\beta_t^2}{2\alpha_t(1-\bar\alpha_t)\|\Sigma_\theta\|^2}\| \epsilon_t - \epsilon_\theta(x_t, t)\|^2 \right].$

2. 训练与推理

$L = \mathbb E_{x_0, \epsilon_t} \left[ \frac{\beta_t^2}{2\alpha_t(1-\bar\alpha_t)|\Sigma_\theta|^2}| \epsilon_t - \epsilon_\theta(x_t, t)|^2 \right]$ 作为loss稍显复杂，而且$\Sigma_\theta$会带来训练的不稳定。Nichol & Dhariwal (2021) 提出使用简化loss $L_{\mathtt{simp}} = \mathbb E_{x_0, \epsilon_t} \left[ \| \epsilon_t - \epsilon_\theta(x_t, t)\|^2 \right].$ 该优化目标更简洁，也让训练变得更稳定，是现在主流Diffusion模型的实现。

与此同时，Nichol & Dhariwal (2021) 还提出用余弦$\{\bar\alpha_t\}$序列替代线性$\{\beta_t\}$序列以提高NLL。这样的参数选择使得扩散过程在开端和结尾处变动小，中间段变动大。需要注意的是要同时保证 $\beta_t < 1-\delta$，因此具体实现时通常先以$\{\bar\alpha_t\}$反算$\{\beta_t\}$，对$\{\beta_t\}$作范围限制，再得到最终$\{\bar\alpha_t\}$。

反向生成时，$x_{t-1} \sim p_\theta(x_t, t)$，依赖 $\mu_\theta(x_t, t)$ 和 $\Sigma_\theta$。采样的噪声强度$\Sigma_\theta$有很重要的影响，如果简单设为$0$则生成的样本会坍缩。但如果采用$L_{\mathtt{simp}}$训练，则无法获知 $\Sigma_\theta$。在Ho et al. (2020) 中作者使用 $\Sigma_\theta = \sigma_t^2 I$来进行生成得到了较好的效果，报告称 $\sigma_t^2 = \beta_t$ 或者 $\sigma_t^2 = \frac{1-\bar \alpha_{t-1}}{1-\bar \alpha_t} \beta_t$ 都能得到相似结果。然而笔者在toy example上复现后发现后者的结果还是要好不少。

3. 若干引理

Lemma 1. $L_{CE} \leq L_{VLB}$

$\mathbb{E}_{q(x_0)} \left[-\log p_{\theta}(x_0)\right] = L_{CE} \leq L_{VLB} = \mathbb{E}_{q(x_{0:T})}\left[\log \frac{q(x_{1:T} | x_0)}{p_\theta(x_{0:T})}\right]$

由Jense不等式可以自然推出该关系

$$\begin{align*} L_{CE} &= \mathbb{E}_{q(x_0)} [-\log p_{\theta}(x_0)] \\ &= -\mathbb{E}_{q(x_0)} \log \smallint p_\theta(x_{0:T}) dx_{1:T} \\ &= -\mathbb{E}_{q(x_0)} \log\left( \int q(x_{1:T}|x_0) \frac{p_\theta(x_{0:T}) }{q(x_{1:T}|x_0) } dx_{1:T}\right) \\ &= -\mathbb{E}_{q(x_0)} \log \left( \mathbb{E}_{q(x_{1:T}|x_0)} \left[ \frac{p_\theta(x_{0:T}) }{q(x_{1:T}|x_0) } \right] \right) \\ &\leq - \mathbb{E}_{q(x_{0:T})} \left[\log \frac{p_\theta(x_{0:T}) }{q(x_{1:T}|x_0) } \right] \\ &= L_{VLB} \end{align*}$$

Lemma 2. $L_{VLB}= L_0 + L_{1:T-1} + L_T$

由贝叶斯定理 $q(x_t|x_{t-1}, x_0)=\frac{q(x_t, x_{t-1}, x_0)}{q(x_{t-1}, x_0)} = q(x_{t-1}|x_t, x_0) \frac{q(x_t, x_0)}{q(x_{t-1}, x_0)}$

因此

$$\begin{align*} L_{VLB} &= \mathbb{E}_{q(x_{0:T})}\left[\log \frac{q(x_{1:T} | x_0)}{p_\theta(x_{0:T})}\right] \\ &= \mathbb{E}_{q} \bigg[ \log \frac{\prod_{t=1}^T q(x_t|x_{t-1}, x_0)}{p_\theta(x_T)\prod_{t=1}^T p_\theta(x_{t-1}|x_t)}\bigg] \\ &= \mathbb{E}_{q}\bigg[-\log p_\theta(x_T) + \sum_{t=1}^T \log \frac{q(x_t|x_{t-1}, x_0)}{ p_\theta(x_{t-1} | x_t)}\bigg] \\ &= \mathbb{E}_{q} \bigg[-\log p_\theta(x_T) + \sum_{t=2}^T \log \frac{q(x_t|x_{t-1}, x_0)}{ p_\theta(x_{t-1} | x_t)} + \log \frac{q(x_1| x_0)}{ p_\theta(x_0 | x_1)} \bigg] \\ &= \mathbb{E}_{q} \bigg[-\log p_\theta(x_T) + \sum_{t=2}^T \log \frac{q(x_{t-1}|x_t, x_0)}{ p_\theta(x_{t-1} | x_t)} \frac{q(x_t, x_0)}{q(x_{t-1}, x_0)} + \log \frac{q(x_1| x_0)}{ p_\theta(x_0 | x_1)} \bigg] \\ &= \mathbb{E}_{q} \bigg[-\log p_\theta(x_T) + \sum_{t=2}^T \log \frac{q(x_{t-1}|x_t, x_0)}{ p_\theta(x_{t-1} | x_t)} + \log \frac{q(x_T, x_0)}{q(x_1, x_0)} + \log \frac{q(x_1| x_0)}{ p_\theta(x_0 | x_1)} \bigg] \\ &= \mathbb{E}_{q} \bigg[ \log \frac{q(x_T| x_0)}{p_\theta(x_T)} + \sum_{t=2}^T \log \frac{q(x_{t-1}|x_t, x_0)}{ p_\theta(x_{t-1} | x_t)} - \log p_\theta(x_0 | x_1) \bigg] \\ &=\mathbb{E}_{q} \bigg[ \underbrace{D_{KL} (q(x_T|x_0)\|p_\theta(x_T))}_{L_T} + \sum_{t=2}^T \underbrace{D_{KL} (q(x_{t-1}|x_t, x_0) \|p_\theta(x_{t-1}|x_t))}_{L_{t-1}} + \underbrace{(- \log p_\theta(x_0 | x_1))}_{L_0} \bigg] \\ &= L_0 + L_{1:T-1} + L_T \end{align*}$$

Lemma 3. $q(x_{t-1}|x_t, x_0)$ 是高斯分布

$q(x_{t-1}|x_t, x_0) = \frac{q(x_t, x_{t-1}, x_0)}{q(x_t, x_0)} = q(x_t|x_{t-1}) \frac{q(x_{t-1}|x_0)}{q(x_t | x_0)}$

已知 $q(x_t|x_{t-1}) = \mathcal{N}(\sqrt{1-\beta_t} x_{t-1}, \beta_t I)$，令$\alpha_t = 1-\beta_t, \bar\alpha_t = \prod_{i=1}^{t} \alpha_i$，不难证明：$q(x_t|x_0) = \mathcal{N}(\sqrt{\bar\alpha_t} x_{t-1}, (1-\bar\alpha_t) I)$

因此

$$\begin{align*} q(x_{t-1}|x_t, x_0) \propto& \exp \bigg(-\frac{(x_t - \sqrt{1-\beta_t}x_{t-1})^2}{2\beta_t} - \frac{(x_{t-1} - \sqrt{\bar\alpha_{t-1}}x_0)^2}{2(1-\bar\alpha_{t-1})} \\ &+ \frac{(x_t - \sqrt{\bar\alpha_t}x_0)^2}{2(1-\bar\alpha_t)} \bigg) \\ =& \exp \left(- \frac{(x_{t-1} - \mu(x_t, x_0))^2}{2\sigma_t^2} + C(x_t, x_0) \right) \end{align*}$$

其中

$$\begin{align*} \mu(x_t, x_0) &= \frac{\sqrt{\alpha_t}(1-\bar\alpha_{t-1})}{1-\bar\alpha_t} x_t + \frac{\sqrt{\bar\alpha_{t-1}}\beta_t}{1-\bar\alpha_t} x_0\\ \sigma_t^2 &= \frac{1-\bar \alpha_{t-1}}{1-\bar \alpha_t} \beta_t \end{align*}$$

结合 $x_t = \sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t} \epsilon_t$ 可写为 $\mu(x_t, x_0) = \frac{1}{\sqrt{\alpha_t}}\left(x_t - \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\epsilon_t\right)$

4. 后记

DDPM的核心是构造了一个扩散过程，这个过程既可以分解为 $q(x_t|x_{t-1})$ 的每一小步，使得反向过程 $q(x_{t-1}|x_t)$ 接近于高斯分布，能够被模型学习，又具有解析的 $q(x_t|x_0)$ 形式，使得训练数据可以直接从原始分布一步采样得到，从而让训练不受扩散采样过程影响。

然而DDPM的缺陷在于反向过程迭代次数太多影响推理速度，因此Song et al. 2022 提出了DDIM，一个基于非马尔可夫扩散过程的Diffusion Model。通过巧妙设置非马尔可夫过程 $q(x_t|x_{t-1}, x_0)$，使得 $q(x_t|x_0)$ 具有和DDPM相同的表达式，从而有相同的训练loss，但在推理时可以使用不同的alpha schedule，从而加速生成。

Reference

[1] https://lilianweng.github.io/posts/2021-07-11-diffusion-models/
[2] Denoising Diffusion Probabilistic Models
[3] Improved Denoising Diffusion Probabilistic Models
[4] Denoising Diffusion Implicit Models
[5] https://github.com/milkpku/toy_diffusion