对于经典的扩散模型DDPM,图像逆向生成推理过程中,我们需要一步步地将随机高斯噪声\(\mathbf{x}_T\)还原成有意义的图像\(\mathbf{x}_0\),总共需要\(T\)次推理,生成速度较慢。为此,DDIM(Denoising Diffusion Implicit Models)采用更小的采样次数来加速推理过程,同时从随机噪音生成图像的过程并没有引入其它噪音,即逆向过程是一个确定的过程。DDIM与DDPM的优化目标是一致的,但并不要求前向过程是一个马尔可夫链。
前提
DDPM的优化目标\(L_t\)以及采样过程\(q(\mathbf{x}_{t-1}\vert\mathbf{x}_t)\)(由\(p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t)\)去近似)如下:
\[ \begin{aligned} L_t &= \mathbb{E}_{t \sim [1, T], \mathbf{x}_0, \boldsymbol{\epsilon}_t} \Big[\|\boldsymbol{\epsilon}_t - \boldsymbol{\epsilon}_\theta(\sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon}_t, t)\|^2 \Big] \\ p_\theta(\mathbf{x}_{t-1} \vert \mathbf{x}_t) &= \mathcal{N}(\mathbf{x}_{t-1}; \frac{1}{\sqrt{\alpha_t}} \Big( \mathbf{x}_t - \frac{1 - \alpha_t}{\sqrt{1 - \bar{\alpha}_t}} \boldsymbol{\epsilon}_\theta(\mathbf{x}_t, t) \Big), \frac{1-\bar\alpha_{t-1}}{1-\bar\alpha_t}\beta_t) \\ \end{aligned} \]
从上面两个式子可以得到DDPM的损失函数只依赖于\(q(\mathbf{x}_t\vert\mathbf{x}_0)\),而采样过程只依赖于\(q(\mathbf{x}_{t-1}\vert\mathbf{x}_t)\)。DDIM中把DDPM公式推导用到的\(q(\mathbf{x}_{t}\vert\mathbf{x}_{t-1})\)给去掉了,这意味着DDIM并没有去定义前向过程的具体操作。
逆向过程
在DDPM中,我们借助贝叶斯定理将逆向过程转变为了前向过程:
\[ q(\mathbf{x}_{t-1}\vert\mathbf{x}_t,\mathbf{x}_0) = \frac{q(\mathbf{x}_t\vert\mathbf{x}_{t-1},\mathbf{x}_0)q(\mathbf{x}_{t-1}\vert\mathbf{x}_0)}{q(\mathbf{x}_t\vert\mathbf{x}_0)} \]
整理后两边同时对\(\mathbf{x}_t\)取积分可以把\(q(\mathbf{x}_t\vert\mathbf{x}_0)\)给消除掉,最后得到下式:
\[ \begin{aligned} \int q(\mathbf{\mathbf{x}_{t-1}\vert\mathbf{x}_t,\mathbf{x}_0})q(\mathbf{x}_t\vert\mathbf{x}_0)d\mathbf{x}_t = q(\mathbf{x}_{t-1}\vert\mathbf{x}_0) \end{aligned} \tag{1} \]
上式即是扩散模型需要满足的边际分布条件。
对于\(q(\mathbf{x}_{t-1}\vert\mathbf{x}_t,\mathbf{x}_0)\),它是一个高斯分布,所以我们直接假设它为下面的形式:
\[ \begin{aligned} q(\mathbf{x}_{t-1}\vert\mathbf{x}_t,\mathbf{x}_0) = \mathcal{N}(\mathbf{x}_{t-1}; \kappa_t\mathbf{x}_t+\lambda_t\mathbf{x}_0,\sigma^2_t\mathbf{I}) \end{aligned} \tag{2} \]
由于DDIM的优化目标与DDPM相同,所以我们可以复用DDPM的训练过程,即\(q(\mathbf{x}_{t-1}\vert\mathbf{x}_0)\)与\(q(\mathbf{x}_t\vert\mathbf{x}_0)\),将这两个变量与(1)和(2)结合,可以得到下面式子:
\[ \begin{aligned} \sqrt{\overline\alpha_{t-1}} &= \kappa_t\sqrt{\overline\alpha_t} + \lambda_t \\ 1-\overline\alpha_t &= \kappa_t^2(1-\overline\alpha_t)+\sigma^2_t \end{aligned} \]
将\(\sigma_t\)看作参数,则可以解出:
\[ \begin{aligned} \kappa_t &= \frac{\sqrt{1-\overline\alpha_t-\sigma^2_t}}{1-\overline\alpha_t} \\ \lambda_t &= \sqrt{\overline\alpha_{t-1}} - \frac{\sqrt{1-\overline\alpha_t-\sigma^2_t}\sqrt{\overline\alpha_t}}{\sqrt{1-\overline\alpha_t}} \end{aligned} \]
代入(2)式可得:
\[ \begin{aligned} &q(\mathbf{x}_{t-1}\vert\mathbf{x}_t,\mathbf{x}_0) \\ =& \mathcal{N}(\mathbf{x}_{t-1};\sqrt{\overline\alpha_{t-1}}\mathbf{x}_0+\sqrt{1-\overline\alpha_t-\sigma^2_t}\frac{\mathbf{x}_t-\sqrt{\overline\alpha_t}\mathbf{x}_0}{\sqrt{1-\overline\alpha_t}}, \sigma^2_t\mathbf{I}) \end{aligned} \tag{3} \]
关于(3)式更加详细的推导说明可以参考这篇博客,原论文中也有另外一种推导方法可以参考。
由于我们已知\(q(\mathbf{x}_t\vert\mathbf{x}_0)\),即已知\(\mathbf{x}_t\)和\(\mathbf{x}_0\)的转换关系,由重参数法可以进一步地得到\(\mathbf{x}_{t}\)到\(\mathbf{x}_{t-1}\)的转换如下:
\[ \mathbf{x}_{t-1} = \sqrt{\overline\alpha_{t-1}}\frac{\mathbf{x}_t-\sqrt{1-\overline\alpha_t}\epsilon_\theta(\mathbf{x}_t,t)}{\sqrt{\overline\alpha_t}} + \sqrt{1-\overline\alpha_{t-1}-\sigma^2_t}\cdot\epsilon_\theta(\mathbf{x}_t,t) + \sigma_t\epsilon_t \]
与DDPM一样,\(\epsilon_\theta(\mathbf{x}_t,t)\)是由神经网络预测的噪音。
DDIM作者将\(\sigma_t\)定义为了下面形式:
\[ \begin{aligned} \sigma_t &= \eta\cdot\sqrt{\tilde{\beta}_t} \\ &= \eta\cdot\sqrt{\frac{1-\overline\alpha_{t-1}}{1-\overline\alpha_t}}\sqrt{1-\alpha_t} \end{aligned} \tag{4} \]
当\(\eta\)为1时,上式有\(\sigma^2_t = \tilde\beta_t\),即与DDPM的逆向过程一致。而当\(\eta=0\),此时的逆向过程没有随机噪音(\(\sigma_t\epsilon_t\)项为0),即最开始输入的随机噪音\(\mathbf{x}_T\)确定时,逆向过程就是一个确定的过程,这种情况下的该模型就是DDIM。
总体上看,DDIM与DDPM的训练过程是一样的,所以我们可以利用训练好的DDPM模型,接上DDIM的逆向过程来做推理。
加速采样
DDIM依然是需要逐步地采样得到生成的图片,但它采样的序列\(\tau=[\tau_1,\tau_2,\dots,\tau_{len(\tau)}]\)是DDPM采样序列\([1,2,\dots,T]\)的任意子序列。假如我们用的是DDPM按\(T\)步采样训练好的模型,利用DDIM,我们可以在推理阶段将\(T\)步缩减为\(len(\tau)\)步,达到加速推理的作用。