扩散模型医学图像生成 扩散效应理论模型

引言

在前面的博客中，我们讨论了生成模型VAE和GAN，近年来，新的生成模型——扩散模型受到越来越多的关注，因此值得好好去研究一番。扩散模型（Diffusion Models）最早由 [2] 于2015年提出，但直到2020年论文 [3] 发表之后才得到关注，本文详细梳理了 [3] 中的公式推导部分，帮助大家更好理解其中的数学原理。

---

数学模型

如下图所示（引自[1]），$x_0$ 是原始数据，$q$ 是扩散模型，每扩散一次，都会在前一期数据的基础上添加部分噪声，当 $t \to \infty$，$x_T$ 完全被噪声淹没，成为各向同性的高斯分布 $x_T \sim \mathcal{N}(0,I)$，$p_\theta$ 是生成模型，使用参数为 $\theta$ 的网络来近似，将噪声恢复成有效信息，整个模型满足马尔可夫链条件。

目标函数

本文采用自顶向下的形式进行讲解，首先说明最终关注的目标函数，然后针对其中的细节分别深入。直观来说，我们的目标是让近似分布 $p_\theta(x_0)$ 尽可能接近数据的真实分布 $q(x_0)$，所以目标函数可以用交叉熵来表示：$\begin{aligned} \mathcal{L}_{CE}&=-\mathbb{E}_{q(x_0)}log\ p_\theta(x_0)\\ &=-\mathbb{E}_{q(x_0)}log[\int p_\theta(x_{0:T})dx_{1:T}]\\ &=-\mathbb{E}_{q(x_0)}log[\int q(x_{1:T}|x_0)\frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)}dx_{1:T}]\\ &=-\mathbb{E}_{q(x_0)}log[\mathbb{E}_{q(x_{1:T}|x_0)}\frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)}]\\ &\le -\mathbb{E}_{q(x_{0:T})}log\frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)},\ Jensen\ inequality \end{aligned}$ 可以得到：$\mathbb{E}_{q(x_0)}log\ p_\theta(x_0) \ge \mathbb{E}_{q(x_{0:T})}log\frac{p_\theta(x_{0:T})}{q(x_{1:T}|x_0)}$ 不等式右边的项就是对数似然下界，记为 $\mathcal{L}_{LB}$，只要让其越大，不等式左边的项也就越大，交叉熵也就越小。

马克洛夫链假设

为了让 $\mathcal{L}_{LB}$ 便于优化，需要补充一些知识。前面提到，模型满足马克洛夫链条件（Markov Chain），即当前状态 $x_t$ 仅与上一状态 $x_{t-1}$ 有关，假设马克洛夫关系为 $A \to B \to C$，可以得到性质：$\begin{aligned} p(B,C|A)&=\frac{p(B,C,A)p(A,B)}{p(A)p(A,B)}=p(B|A)p(C|A,B)\\ &=p(B|A)p(C|B)\tag{1} \end{aligned}$ 利用公式（1），可以得到：$\tag{2} q(x_{1:T}|x_0)=\prod^T_{t=1}q(x_t|x_{t-1})\\ p_\theta(x_{0:T})=p_\theta(x_T)\prod^T_{t=1}p_\theta(x_{t-1}|x_t)$ 将公式（2）代入 $\mathcal{L}_{LB}$，可以得到：$\tag{3} \begin{aligned} \mathcal{L}_{LB}&=\mathbb{E}_{q(x_{0:T})}[log\frac{q(x_{1:T}|x_0)}{p_\theta(x_{0:T})}]\\ &=\mathbb{E}_{q}[log\frac{\prod^T_{t=1}q(x_t|x_{t-1})}{p_\theta(x_T)\prod^T_{t=1}p_\theta(x_{t-1}|x_t)}]\\ &=\mathbb{E}_{q}[-log\ p_\theta(x_T)+\sum^T_{t=1}log\frac{q(x_t|x_{t-1})}{p_\theta(x_{t-1}|x_t)}]\\ &=\mathbb{E}_{q}[-log\ p_\theta(x_T)+\sum^T_{t=2}log\frac{q(x_t|x_{t-1})}{p_\theta(x_{t-1}|x_t)}+log\frac{q(x_1|x_0)}{p_\theta(x_0|x_1)}]\\ &=\mathbb{E}_{q}[-log\ p_\theta(x_T)+\sum^T_{t=2}log\frac{q(x_{t-1}|x_t,x_0)}{p_\theta(x_{t-1}|x_t)} \cdot \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)}]\\ &=\mathbb{E}_{q}[-log\ p_\theta(x_T)+\sum^T_{t=2}log\frac{q(x_{t-1}|x_t,x_0)}{p_\theta(x_{t-1}|x_t)}+\sum^T_{t=2}log\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)}]\\ &=\mathbb{E}_{q}[-log\ p_\theta(x_T)+\sum^T_{t=2}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)}]\\ &=\mathbb{E}_{q}[log\frac{q(x_T|x_0)}{p_\theta(x_T)}+\sum^T_{t=2}log\frac{q(x_{t-1}|x_t,x_0)}{p_\theta(x_{t-1}|x_t)}-log\ p_\theta(x_0|x_1)] \end{aligned}$ 将公式（3）的最后一行化成 KL 散度的形式：$\tag{4} \mathbb{E}_{q(x_0)}D_{KL}[q(x_T|x_0)||p_\theta(x_T)]+\\ \mathbb{E}_{q(x_0,x_t)}\sum^T_{t=2}D_{KL}[q(x_{t-1}||x_0,x_t)||p_\theta(x_{t-1}|x_t)]-\\ \mathbb{E}_{q(x_0,x_1)}log\ p_\theta(x_0|x_1)$ 公式（4）第一项对应 VAE 中的正则化损失 $D_{KL}(q_\phi(z|x)||p_\theta(z))$，第三项对应于重建损失 $\mathbb{E}_{q_\phi(z|x)}[log(p_\theta(x|z))]$，第二项是多个 KL 散度的和，每个度量 $p$ 后验分布和 $q$ 已知 $x_0$

重参数化

利用重参数化技巧可以让公式（4）中的 $x_t$ 可解，进一步简化 $\mathcal{L}_{LB}$。给定真实数据 $x_0 \sim q(x)$，扩散过程的每一步可以表示为：$q(x_t|x_{t-1})=\mathcal{N}(x_t;\ \sqrt{1-\beta_t}x_{t-1},\beta_tI)\tag{5}$ 其中 $\beta_t$ 是一个超参数。利用重参数化技巧，可以使用 $x_0$ 直接计算任意时间点 t 上的 $x_t$，不需要一步步迭代。假设 $\alpha_t=1-\beta_t,\overline{\alpha_t}=\prod^t_{i=1}\alpha_i,\epsilon_t \sim \mathcal{N}(0,I)$：$\tag{6} \begin{aligned} x_t&=\sqrt{\alpha_t}x_{t-1}+\sqrt{1-\alpha_t}\epsilon_{t-1}\\ &=\sqrt{\alpha_t}[\sqrt{\alpha_{t-1}}x_{t-2}+\sqrt{1-\alpha_{t-1}}\epsilon_{t-2}]+\sqrt{1-\alpha_t}\epsilon_{t-1}\\ &=\sqrt{\alpha_t\alpha_{t-1}}x_{t-2}+\sqrt{\alpha_t-\alpha_t\alpha_{t-1}}\epsilon_{t-2}+\sqrt{1-\alpha_t}\epsilon_{t-1}\\ &=\sqrt{\alpha_t\alpha_{t-1}}x_{t-2}+\sqrt{1-\alpha_t\alpha_{t-1}}\overline{\epsilon}_{t-2}\\ &=...\\ &=\sqrt{\overline{\alpha}_t}x_0+\sqrt{1-\overline{\alpha}_t}\epsilon\\ &\therefore q(x_t|x_0)=\mathcal{N}(x_t;\ \sqrt{\overline{\alpha}_t}x_0,(1-\overline{\alpha}_t)I) \end{aligned}$ 其中 $\sqrt{\alpha_t-\alpha_t\alpha_{t-1}}\epsilon_{t-2}+\sqrt{1-\alpha_t}\epsilon_{t-1} \sim \mathcal{N}(0,1-\alpha_t\alpha_{t-1})$，从而得到 $\overline{\epsilon}_{t-2}$。进一步，可以计算 $q(x_{t-1}|x_t,x_0)$ 的解析式：$\tag{7} \begin{aligned} q(x_{t-1}|x_t,x_0)&=q(x_t|x_{t-1},x_0)\frac{q(x_{t-1}|x_0)}{q(x_t|x_0)}\\ &\propto exp[-\frac{1}{2}(\frac{(x_t-\sqrt{\alpha_t}x_{t-1})^2}{\beta_t}+\frac{(x_{t-1}-\sqrt{\overline{\alpha}_{t-1}}x_0)^2}{1-\overline{\alpha}_{t-1}}-\frac{(x_t-\sqrt{\overline{\alpha}_t}x_0)^2}{1-\overline{\alpha}_t})]\\ &=exp[-\frac{1}{2}((\frac{\alpha_t}{\beta_t}+\frac{1}{1-\overline{\alpha}_{t-1}})x^2_{t-1}-(\frac{2\sqrt{\alpha_t}}{\beta_t}x_t+\frac{2\sqrt{\overline{\alpha}_{t-1}}}{1-\overline{\alpha}_{t-1}}x_0)x_{t-1}+C(x_t,x_0))] \end{aligned}$ 我们知道，$ax^2+bx$ 可以化成 $a(x+\frac{b}{2a})^2+c$，那么对于指数项是 $ax^2+bx$ 这种格式的高斯分布，$\mu=-\frac{b}{2a},\sigma^2=\frac{1}{a}$。设 $q(x_{t-1}|x_t,x_0)=\mathcal{N}(x_{t-1};\ \tilde{\mu}(x_t,x_0),\tilde{\beta}_tI)$，代入公式（7）有：$\tag{8} \begin{aligned} \tilde{\beta}_t&=\frac{1-\overline{\alpha}_{t-1}}{1-\overline{\alpha}_t}\cdot\beta_t\\ \tilde{\mu}_t(x_t,x_0)&=\frac{\sqrt{\alpha_t}(1-\overline{\alpha}_{t-1})}{1-\overline{\alpha}_t}x_t+\frac{\sqrt{\overline{\alpha}_{t-1}}\beta_t}{1-\overline{\alpha}_t}x_0 \end{aligned}$ 根据公式（6），将 $x_0=\frac{1}{\sqrt{\overline{\alpha}_t}}(x_t-\sqrt{1-\overline{\alpha}_t}\epsilon_t)$ 代入公式（8），可以进一步化简 $\tilde{\mu}_t(x_t,x_0)$：$\tilde{\mu}_t(x_t,x_0)=\frac{1}{\sqrt{\alpha_t}}(x_t-\frac{1-\alpha_t}{\sqrt{1-\overline{\alpha}_t}}\epsilon_t)\tag{9}$

高斯分布的 KL 散度

重参数化技巧使能够解析计算 $x_t$ 和 $q(x_{t-1}|x_t,x_0)$，除此之外，当前分布为高斯分布时，KL 散度的计算也能得到简化，考虑一维高斯分布 $a \sim \mathcal{N}(\mu_1,\sigma^2_1),b \sim \mathcal{N}(\mu_2,\sigma^2_2)$，它们的 KL 散度为（同理可以扩展到多维高斯分布）：$\tag{10} \begin{aligned} D_{KL}(a||b)&=\mathbb{E}_alog\ \frac{a}{b}\\ &=\mathbb{E}_a[log\frac{\sigma_2}{\sigma_1}+\frac{(x-\mu_2)^2}{2\sigma^2_2}+\frac{(x-\mu_1)^2}{2\sigma^2_1}]\\ &=log\frac{\sigma_2}{\sigma_1}+\frac{1}{2}\mathbb{E}_a[\frac{(x-\mu_2)^2}{\sigma^2_2}+\frac{(x-\mu_1)^2}{\sigma^2_1}]\\ &=log\frac{\sigma_2}{\sigma_1}+\frac{1}{2}\mathbb{E}_a[(\frac{1}{\sigma^2_2}-\frac{1}{\sigma^2_1})x^2+(\frac{2\mu_1}{\sigma^2_1}-\frac{2\mu_2}{\sigma^2_2})x+\frac{\mu^2_2}{\sigma^2_2}-\frac{\mu^2_1}{\sigma^2_1}]\\ &=log\frac{\sigma_2}{\sigma_1}+\frac{\sigma^2_1+(\mu_1-\mu_2)^2}{2\sigma^2_2}-\frac{1}{2} \end{aligned}$

优化目标函数

公式（4）中的第一项是常数，第三项可以看作是第二项 $t=1$ 时的结果，所以我们主要考虑第二项。设 $p_\theta(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};\ \mu_\theta(x_t,t),\sigma^2_tI)$，根据公式（10），第二项可以化简为：$L_{t-1}=\mathbb{E}_q[\frac{1}{2\sigma^2_t}||\tilde{\mu}_t(x_t,x_0)-\mu_\theta(x_t,t)||^2]+C\tag{11}$ 将公式（9）代入公式（11）可以得到：$L_{t-1}=\mathbb{E}_{x_t,\epsilon}[\frac{1}{2\sigma^2_t}||\frac{1}{\sqrt{\alpha_t}}(x_t-\frac{\beta_t}{\sqrt{1-\overline{\alpha}_t}}\epsilon)-\mu_\theta(x_t,t)||^2]\tag{12}$ [3]作者进行了参数化 $\mu_\theta(x_t,t)=\frac{1}{\sqrt{\alpha_t}}(x_t-\frac{\beta_t}{\sqrt{1-\overline{\alpha}_t}}\epsilon_\theta(x_t,t))$，相当于网络拟合的是每个时间点的噪声（为什么选择这样参数化我还没明白），同时代入公式（6），公式（12）进一步化简为：$L_{t-1}=\mathbb{E}_{x_0,\epsilon}[\frac{\beta^2_t}{2\sigma^2_t\alpha_t(1-\overline{\alpha}_t)}||\epsilon-\epsilon_\theta(\sqrt{\overline{\alpha}_t}x_0+\sqrt{1-\overline{\alpha}_t}\epsilon,t)||^2]$ [3]作者发现，将前面的系数丢掉，训练更加稳定，因此得到最终的损失：$L_{simple}=\mathbb{E}_{t,x_0,\epsilon}[||\epsilon-\epsilon_\theta(\sqrt{\overline{\alpha}_t}x_0+\sqrt{1-\overline{\alpha}_t}\epsilon,t)||^2]$

参考

[1] Understanding Diffusion Models: A Unified Perspective [2] Deep Unsupervised Learning using Nonequilibrium Thermodynamics [3] Denoising Diffusion Probabilistic Models [4] What are Diffusion Models? [5] Probabilistic Diffusion Model概率扩散模型理论