Analytic-DPM与DPM-solver

本文参考苏神的博客网站。

【暂待完善】

Analytic-DPM

这篇文章提出，目前诸如DDPM和DDIM等模型在最优化ELBO的过程中，存在一个潜在的优化空间是，在给定的前向扩散过程的超参数下，其一般认为反向扩散过程的方差$\sigma_n^2$就是最优的方差取值，但是该文通过分析得到实际上并非如此，核心的原因就在于我们拟合${\bf x}_0$时存在潜在的误差。【这篇文章提出的数学框架实际上指出，我们原本的扩散模型训练流程是没有问题的，但是在推理的时候，我们可以使用训练好的模型通过蒙特卡洛方法来拟合出理论上最优的推理流程方差策略，从而在已有训练模型的基础上进一步提升推理效果，这样看起来那些用模型来拟合最优方差的工作或许被杀死了？当然也可能没有，因为这篇文章的理论推理实际上也引入了一些假设，这些假设会限制优化空间，其实我们可以进一步扩展分析空间来发一篇新的理论文章，或者在指出这篇文献存在的潜在假设局限后，指出我们可以用模型拟合得到潜在的更广的优化空间】

回顾DDPM和DDIM的推导，我们不难发现，最重要的优化项是

$$\mathbb{E}_{ {\bf x}_{0:T}\sim q({\bf x}_{0:T})}\left[\sum_{t>1}\log\frac{p_\theta({\bf x}_{t-1}|{\bf x}_t)}{q({\bf x}_{t-1}|{\bf x}_{t},{\bf x}_0)}\right]$$

我们希望这一项尽可能地大，该优化目标实际上可以拆分为多个如下的子目标

$$\mathbb{E}_{ {\bf x}_{0,(t-1),t}\sim q({\bf x}_{0,(t-1),t})}\left[\log\frac{p_\theta({\bf x}_{t-1}|{\bf x}_t)}{q({\bf x}_{t-1}|{\bf x}_{t},{\bf x}_0)}\right]$$

无论是DDPM还是DDIM，我们实际上都可以推导得到

$$q({\bf x}_{t-1}|{\bf x}_{t},{\bf x}_0)=\mathcal{N}({\bf x}_{t-1};a{\bf x}_t+b{\bf x}_0,\sigma_t^2{\bf I})$$ $$q({\bf x}_{t}|{\bf x}_0)=\mathcal{N}({\bf x}_{t};c{\bf x}_0,\lambda_t^2{\bf I})$$

假设

$$p_\theta({\bf x}_{t-1}|{\bf x}_t)=\mathcal{N}({\bf x}_{t-1};{\bf \mu}_t({\bf x}_t),\delta^2_t{\bf }I)$$

那么下面我们要求解，什么样的$\mu_t$和$\delta_t$是最优的解。

$$\begin{align} \mathbb{E}_{ {\bf x}_{0,(t-1),t}\sim q({\bf x}_{0,(t-1),t})}\left[\log\frac{p_\theta({\bf x}_{t-1}|{\bf x}_t)}{q({\bf x}_{t-1}|{\bf x}_{t},{\bf x}_0)}\right]\\ =\mathbb{E}_{ {\bf x}_{0,(t-1),t}\sim q({\bf x}_{0,(t-1),t})}\left[\frac{||{\bf x}_{t-1}-a{\bf x}_t-b{\bf x}_0||^2}{2\sigma_t^2}-\frac{||{\bf x}_{t-1}-\mu_t({\bf x}_t)||^2}{2\delta_t^2}-n\log\delta_t\right]+C\\ =\mathbb{E}_{ {\bf x}_{0,(t-1),t}\sim q({\bf x}_{0,(t-1),t})}\left[-\frac{||{\bf x}_{t-1}-\mu_t({\bf x}_t)||^2}{2\delta_t^2}-n\log\delta_t\right]+C'\\ =-\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0)}\left[\int\left[\frac{||{\bf x}_{t-1}-\mu_t({\bf x}_t)||^2}{2\delta_t^2}+n\log\delta_t\right]q({\bf x}_{t-1},{\bf x}_t|{\bf x}_0)d{\bf x}_{(t-1),t}\right]+C'\\ =-\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0)}\left[\int\left[\frac{||{\bf x}_{t-1}-\mu_t({\bf x}_t)||^2}{2\delta_t^2}+n\log\delta_t\right]q({\bf x}_{t-1}|{\bf x}_t,{\bf x}_0)q({\bf x}_t|{\bf x}_0)d{\bf x}_{(t-1),t}\right]+C'\\ =-\frac{1}{Z}\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0)}\Bigg[\int\bigg[\frac{||{\bf x}_{t-1}-\mu_t({\bf x}_t)||^2}{2\delta_t^2}+n\log\delta_t\bigg]\exp\bigg[-\frac{||{\bf x}_{t-1}-a{\bf x}_t-b{\bf x}_0||^2}{2\sigma_t^2}\\-\frac{||{\bf x}_{t}-c{\bf x}_0||^2}{2\lambda_t^2}\bigg]d{\bf x}_{(t-1),t}\Bigg]+C' \end{align}$$

最重要的项是

$$\mathcal{L}=\frac{1}{Z}\int\int\bigg[\frac{||{\bf x}_{t-1}-\mu_t||^2}{2\delta_t^2}+n\log\delta_t\bigg]\exp\bigg[-\frac{||{\bf x}_{t-1}-a{\bf x}_t-b{\bf x}_0||^2}{2\sigma_t^2}-\frac{||{\bf x}_{t}-c{\bf x}_0||^2}{2\lambda_t^2}\bigg]d{\bf x}_{t-1}d{\bf x}_{t}$$

我们希望求解的$\mu_t$和$\delta_t$都是${\bf x}_t$和${\bf x}_0$的函数(其中${\bf x}_0$在实际使用的时候是采用模型拟合出来的)。首先对${\bf x}_{t-1}$的积分是安全的，借由高斯分布的性质，我们可以求得

$$\begin{align} \mathcal{L}=\frac{1}{Z'}\int\bigg[\frac{n\sigma_t^2+(a{\bf x}_t+b{\bf x_0})^T(a{\bf x}_t+b{\bf x_0})-2(a{\bf x}_t+b{\bf x_0})^T\mu_t+||\mu_t||^2}{2\delta_t^2}+n\log\delta_t\bigg]\exp\bigg[-\frac{||{\bf x}_{t}-c{\bf x}_0||^2}{2\lambda_t^2}\bigg]d{\bf x}_{t}\\ =\frac{1}{Z'}\int\bigg[\frac{n\sigma_t^2+||\mu_t-(a{\bf x}_t+b{\bf x_0})||^2}{2\delta_t^2}+n\log\delta_t\bigg]\exp\bigg[-\frac{||{\bf x}_{t}-c{\bf x}_0||^2}{2\lambda_t^2}\bigg]d{\bf x}_{t} \end{align}$$

从这里可以看出为了让这一项尽可能小，最优的情形显然是

$$\mu_t=a{\bf x}_t+b{\bf x}_0$$

这样的话，我们就得到

$$\begin{align} \mathcal{L}=\frac{1}{Z'}\int\bigg[\frac{n\sigma_t^2}{2\delta_t^2}+n\log\delta_t\bigg]\exp\bigg[-\frac{||{\bf x}_{t}-c{\bf x}_0||^2}{2\lambda_t^2}\bigg]d{\bf x}_{t} \end{align}$$

接下来的任务就是求解最优的$\delta_t$，容易求得$\delta_t=\sigma_t$，这个结果就是我们常见的结果，我们一般用模型拟合${\bf x}_0$，即$\mu_t=a{\bf x}_t+b\hat{\bf x}_0({\bf x}_t)$，这里的

$$\hat{\bf x}_0({\bf x}_t)=\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0|{\bf x}_t)}[{\bf x}_0]$$

但是这个结果是suboptimal的，下面展开分析。

首先考虑一下，为什么我们在推理过程中会出现对${\bf x}_0$的预测需求呢，因为我们在建模的是$p_\theta({\bf x}_{t-1}|{\bf x}_t)$，而我们真实知道的是$q({\bf x}_{t-1}|{\bf x}_t,{\bf x}_0)$，要让二者尽可能接近，那么对于现有的已知条件${\bf x}_t$，我们需要尽可能地推理出其对应的真实${\bf x}_0$，也就是说我们实际上实现了

$$p_\theta({\bf x}_{t-1}|{\bf x}_t)=\int q({\bf x}_{t-1}|{\bf x}_t,{\bf x}_0')p_\theta({\bf x}_0'|{\bf x}_t)d{\bf x}_0'$$

在此前的文章的潜在理念是认为对于给定的${\bf x}_t$，我们总能找到其真实对应的${\bf x}_0$，即

$$p_\theta({\bf x}_0'|{\bf x}_t)=\delta({\bf x}_0,f_\theta({\bf x}_t))$$

这个式子的含义是

$${\bf x}_0'=f_\theta({\bf x}_t)={\bf x}_0$$

但是实际上我们的预测是有偏差的，该文章(潜在地)假设这种偏差是满足一个较为简单的高斯分布的，即

$$p_\theta({\bf x}_0'|{\bf x}_t)=\mathcal{N}({\bf x}_0';{\bf x}_0,\gamma_t^2{\bf I})$$

那么这样就又多出来一个需要分析的项$\gamma_t$，如何估计这个方差呢，实际上我们只需让其尽可能接近$q({\bf x}_0|{\bf x}_t)$的真实方差${\bf\Sigma}({\bf x}_t)$即可

$$\begin{align} {\bf\Sigma}({\bf x}_t)=\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0|{\bf x}_t)}\left[({\bf x}_0-\mathbb{E}_{q({\bf x}_0|{\bf x}_t)}[{\bf x}_0])^T({\bf x}_0-\mathbb{E}_{q({\bf x}_0|{\bf x}_t)}[{\bf x}_0])\right]\\ =\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0|{\bf x}_t)}\left[({\bf x}_0-\hat{\bf x}_0({\bf x}_t))^T({\bf x}_0-\hat{\bf x}_0({\bf x}_t)\right] \end{align}$$

我们的目的就是令

$$\begin{align} \gamma_t^2{\bf I}\approx\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0|{\bf x}_t)}\left[({\bf x}_0-\hat{\bf x}_0({\bf x}_t))^T({\bf x}_0-\hat{\bf x}_0({\bf x}_t)\right] \end{align}$$

考虑到$\gamma_t^2{\bf I}$的特殊性质(对角矩阵，且对角元素全部相同，并且与随机变量无关)，我们可以通过一些trick进行计算

$$\begin{align} \mathbb{E}_{ {\bf x}_t\sim q({\bf x}_t) }[\gamma_t^2{\bf I}]\approx\mathbb{E}_{ {\bf x}_t\sim q({\bf x}_t) }\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0|{\bf x}_t)}\left[({\bf x}_0-\hat{\bf x}_0({\bf x}_t))^T({\bf x}_0-\hat{\bf x}_0({\bf x}_t)\right]\\ \gamma_t^2{\bf I}\approx\mathbb{E}_{ {\bf x}_t\sim q({\bf x}_t) }\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0|{\bf x}_t)}\left[({\bf x}_0-\hat{\bf x}_0({\bf x}_t))^T({\bf x}_0-\hat{\bf x}_0({\bf x}_t)\right]\\ \gamma_t^2{\bf I}\approx\mathbb{E}_{ {\bf x}_t\sim q({\bf x}_t) }\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0|{\bf x}_t)}\left[({\bf x}_0-{\bf x}_t+{\bf x}_t-\hat{\bf x}_0({\bf x}_t))^T({\bf x}_0-{\bf x}_t+{\bf x}_t-\hat{\bf x}_0({\bf x}_t)\right]\\ n\gamma_t^2\approx\mathbb{E}_{ {\bf x}_t\sim q({\bf x}_t) }\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0|{\bf x}_t)}\left[||{\bf x}_0-{\bf x}_t||^2+||{\bf x}_t-\hat{\bf x}_0({\bf x}_t)||^2+({\bf x}_0-{\bf x}_t)^T({\bf x}_t-\hat{\bf x}_0({\bf x}_t))+({\bf x}_t-\hat{\bf x}_0({\bf x}_t))^T({\bf x}_0-{\bf x}_t)\right]\\ n\gamma_t^2\approx\mathbb{E}_{ {\bf x}_t\sim q({\bf x}_t) }\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0|{\bf x}_t)}[||{\bf x}_0-{\bf x}_t||^2+||{\bf x}_t-\hat{\bf x}_0({\bf x}_t)||^2]-2\mathbb{E}_{ {\bf x}_t\sim q({\bf x}_t) }[||{\bf x}_t-\hat{\bf x}_0({\bf x}_t)||^2]\\ n\gamma_t^2\approx\mathbb{E}_{ {\bf x}_t\sim q({\bf x}_t) }\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0|{\bf x}_t)}[||{\bf x}_0-{\bf x}_t||^2]-\mathbb{E}_{ {\bf x}_t\sim q({\bf x}_t) }[||{\bf x}_t-\hat{\bf x}_0({\bf x}_t)||^2]\\ n\gamma_t^2\approx\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0) }[||{\bf x}_0||^2-c{\bf x}_0^T{\bf x}_0+n\lambda_t^2]-\mathbb{E}_{ {\bf x}_t\sim q({\bf x}_t) }[||{\bf x}_t-\hat{\bf x}_0({\bf x}_t)||^2]\\ n\gamma_t^2\approx(1-c)\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0) }[||{\bf x}_0||^2]-\mathbb{E}_{ {\bf x}_t\sim q({\bf x}_t) }[||{\bf x}_t-\hat{\bf x}_0({\bf x}_t)||^2]+n\lambda_t^2\\ n\gamma_t^2\approx\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0) }[||{\bf x}_0||^2-c{\bf x}_0^T{\bf x}_0+n\lambda_t^2]-\mathbb{E}_{ {\bf x}_t\sim q({\bf x}_t) }[||{\bf x}_t-\hat{\bf x}_0({\bf x}_t)||^2]\\ \gamma_t^2\approx\frac{1-c}{n}\mathbb{E}_{ {\bf x}_0\sim q({\bf x}_0) }[||{\bf x}_0||^2]-\frac{1}{n}\mathbb{E}_{ {\bf x}_t\sim q({\bf x}_t) }[||{\bf x}_t-\hat{\bf x}_0({\bf x}_t)||^2]+\lambda_t^2 \end{align}$$

中间进行了一次矩阵的迹的计算。可以看到，我们可以在模型$\hat{\bf x}_0({\bf x}_t)$训练完成后，通过简单的蒙特卡洛方法来估计$\gamma_t$，在估计得到$\gamma_t$之后，再接着推导$\mathcal{L}$，我们可以得到

$$\begin{align} \mathcal{L}=\frac{1}{Z'}\int\bigg[\frac{n\sigma_t^2+nb^2\gamma_t^2}{2\delta_t^2}+n\log\delta_t\bigg]\exp\bigg[-\frac{||{\bf x}_{t}-c{\bf x}_0||^2}{2\lambda_t^2}\bigg]d{\bf x}_{t} \end{align}$$

容易求解得到最优的$\delta_t=\sqrt{\sigma_t^2+b^2\gamma_t^2}$，代入$\gamma_t$即可。

对于DDIM来说

$$\begin{align} a=\frac{\sqrt{(1-\bar\alpha_{t-1})(1-\bar\alpha_t-\sigma^2_t)}}{1-\bar\alpha_t}\\ b=\sqrt{\bar\alpha_{t-1}}-\frac{\sqrt{\bar\alpha_t(1-\bar\alpha_{t-1})(1-\bar\alpha_t-\sigma^2_t)}}{1-\bar\alpha_t}\\ \sigma_t^2=\frac{1-\bar\alpha_{t-1}}{1-\bar\alpha_t}\sigma_t^2\\ c=\sqrt{\bar\alpha_t}\\ \lambda_t^2=1-\bar\alpha_t \end{align}$$

而在这篇文章中

$$\begin{align} a=\frac{\sqrt{\bar\beta_{t-1}-\lambda_t^2} }{\sqrt{\bar\beta_t} }\\ b=\sqrt{\bar\alpha_{t-1}}-\frac{\sqrt{\bar\alpha_t(\bar\beta_{t-1}-\lambda_t^2)} }{\sqrt{\bar\beta_t} }\\ \sigma_t^2=\lambda_t^2\\ c=\sqrt{\bar\alpha_t}\\ \lambda_t^2=\bar\beta_t \end{align}$$

代入求得的式子，就可以得到原本文章中的结果。

在附录中的证明过程中，作者写了很多的引理，但是感觉没必要写这么多引理，这些引理本质上就是推理过程中的中间步骤，可能是为了让推理部分尽可能简洁吧。

DPM-solver

作者的PPT讲解已经比较清晰易懂，暂时贴上来，具体的解析之后补充。

本质上是采用了宋飏论文中提出的正向SDE过程对应的ODE过程，通过泰勒展开来近似求解该ODE过程的积分，从而显著提升推理效率。【Runge-Kutta Method也是一种数值求解ODE的算法】

原作者讲解PPT