对于经典的扩散模型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)$步,达到加速推理的作用。