扩散模型初探

扩散模型的推导

主要是根据以下网址学习diffusion的数学形式：网址

首先使用$q$表示前向扩散过程，使用$p$表示逆向过程。

$$q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;{\sqrt{1-\beta }}_t{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)=\prod _{t=1}^{T}q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})$$

当${\beta }_t$较小的时候，逆向过程也可以认为是一个高斯分布：

$$p_{\theta }({\mathbf{x}}_{0:T})=p({\mathbf{x}}_T)\prod _{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))$$

前向添加噪声的过程是一个马尔科夫链，$t$时刻添加噪声的图像只和$t-1$时刻有关。
令${\alpha }_t=1-{\beta }_t,{\overline{\alpha }}_t=\prod _{i=1}^t{\alpha }_i:$

在前向过程中添加高斯噪声：

$$\begin{array}{rl}{\mathbf{x}}_t& =\sqrt{{\alpha }_t}{\mathbf{x}}_{t-1}+\sqrt{1-{\alpha }_t}{\mathit{ϵ}}_{t-1}\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{\mathbf{x}}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{\overline{\mathit{ϵ}}}_{t-2}\\ & =\dots \\ & =\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\overline{\alpha }}_t}\mathit{ϵ}\\ \mathrm{从}\mathrm{而}\mathrm{得}\mathrm{到}：\\ q({\mathbf{x}}_t|{\mathbf{x}}_0)& =\mathcal{N}({\mathbf{x}}_t;\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0,(1-{\overline{\alpha }}_t)\mathbf{I})\end{array}$$

我们需要根据前向过程求得逆向过程的分布情况，即求得如下的分布情况(我们已经根据一些前置知识知道它是正态分布，所以只用求取如下的参数)：

$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_{t-1};\widetilde{\mathit{\mu }({\mathbf{x}}_t,{\mathbf{x}}_0)},{\tilde{\mathit{\beta }}}_t\mathbf{I})$$

根据贝叶斯公式：

$$\begin{array}{rl}q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)& =q({\mathbf{x}}_t|{\mathbf{x}}_{t-1},{\mathbf{x}}_0)\frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}{q({\mathbf{x}}_t|{\mathbf{x}}_0)}\\ & \propto \exp (-\frac{1}{2}(\frac{({\mathbf{x}}_t-\sqrt{{\alpha }_t}{\mathbf{x}}_{t-1}{)}^2}{{\beta }_t}+\frac{({\mathbf{x}}_{t-1}-\sqrt{{\overline{\alpha }}_{t-1}}{\mathbf{x}}_0{)}^2}{1-{\overline{\alpha }}_{t-1}}-\frac{({\mathbf{x}}_t-\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0{)}^2}{1-{\overline{\alpha }}_t}))\\ & =\exp (-\frac{1}{2}(\frac{{\mathbf{x}}_t^2-2\sqrt{{\alpha }_t}{\mathbf{x}}_t{\mathbf{x}}_{t-1}+{\alpha }_t{\mathbf{x}}_{t-1}^2}{{\beta }_t}+\frac{{\mathbf{x}}_{t-1}^2-2\sqrt{{\overline{\alpha }}_{t-1}}{\mathbf{x}}_0{\mathbf{x}}_{t-1}+{\overline{\alpha }}_{t-1}{\mathbf{x}}_0^2}{1-{\overline{\alpha }}_{t-1}}-\frac{({\mathbf{x}}_t-\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0{)}^2}{1-{\overline{\alpha }}_t}))\\ & =\exp (-\frac{1}{2}((\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\overline{\alpha }}_{t-1}})x_{t-1}^2-(\frac{2\sqrt{{\alpha }_t}}{{\beta }_t}{x}_t+\frac{2{\sqrt{\alpha }}_{t-1}}{1-{\overline{\alpha }}_{t-1}}{x}_0)x_{t-1}+C(x_t,x_0))\end{array}$$

根据正太分布的公式：

$$\begin{array}{rl}f(x)& =\frac{1}{\sigma \sqrt{2\pi }}{e}^{-\frac{1}{2}{\left(\frac{x-\mu }{\sigma }\right)}^2}\\ & =\frac{1}{\sigma \sqrt{2\pi }}\exp (-\frac{1}{2}\left(\frac{x^2-2\mu x+{\mu }^2}{{\sigma }^2}\right))\end{array}$$

从而可以求解出${\tilde{\mathit{\beta }}}_t$和$\tilde{\mathit{\mu }}({\mathbf{x}}_t,{\mathbf{x}}_0)$：

$$\begin{array}{rl}{\tilde{\beta }}_t& =1/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\overline{\alpha }}_{t-1}})=1/(\frac{{\alpha }_t-{\overline{\alpha }}_t+{\beta }_t}{{\beta }_t(1-{\overline{\alpha }}_{t-1})})=\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}\cdot {\beta }_t\\ {\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)& =(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{\mathbf{x}}_t+\frac{\sqrt{{\overline{\alpha }}_{t-1}}}{1-{\overline{\alpha }}_{t-1}}{\mathbf{x}}_0)/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\overline{\alpha }}_{t-1}})\\ & =(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{\mathbf{x}}_t+\frac{\sqrt{{\overline{\alpha }}_{t-1}}}{1-{\overline{\alpha }}_{t-1}}{\mathbf{x}}_0)\frac{1-{\overline{\alpha }}_{t-1}}{1-{\overline{\alpha }}_t}\cdot {\beta }_t\\ & =\frac{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{\mathbf{x}}_t+\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}{\mathbf{x}}_0\end{array}$$

根据上述$\text{bold}{x}_0,\text{bold}{x}_t$关系，将$\tilde{\mathit{\mu }}({\mathbf{x}}_t,{\mathbf{x}}_0)$表示：

$$\begin{array}{rl}{\tilde{\mathit{\mu }}}_t& =\frac{\sqrt{{\alpha }_t}(1-{\overline{\alpha }}_{t-1})}{1-{\overline{\alpha }}_t}{\mathbf{x}}_t+\frac{\sqrt{{\overline{\alpha }}_{t-1}}{\beta }_t}{1-{\overline{\alpha }}_t}\frac{1}{\sqrt{{\overline{\alpha }}_t}}({\mathbf{x}}_t-\sqrt{1-{\overline{\alpha }}_t}{\mathit{ϵ}}_t)\\ & =\frac{1}{\sqrt{{\alpha }_t}}({\mathrm{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{ϵ}_t)\end{array}$$

扩散模型现在需要的就是通过一个网络来预测$\tilde{\mathit{\mu }}({\mathbf{x}}_t,{\mathbf{x}}_0)$，如果我们得到了很好的$\tilde{\mathit{\mu }}({\mathbf{x}}_t,{\mathbf{x}}_0)$参数，那么我们可以能够得到逆向采样过程的分布，通过不断地进行采样，就可以生成逆向采样过程的图像了。
我们需要设计网络来预测，那么我们就需要设计loss函数来进行优化。

首先，我们需要生成地图像分布要和真实图像一致，采用最大似然估计，这个地方的推理和VAE是一致的。

$$\begin{array}{rl}-\log p_{\theta }({\mathbf{x}}_0)& \le -\log p_{\theta }({\mathbf{x}}_0)+D_{\mathrm{KL}}(q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)\Vert p_{\theta }({\mathbf{x}}_{1:T}|{\mathbf{x}}_0))\\ & =-\log p_{\theta }({\mathbf{x}}_0)+{\mathbb{E}}_{{\mathbf{x}}_{1:T}\sim q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}[\log \frac{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})/p_{\theta }({\mathbf{x}}_0)}]\\ & =-\log p_{\theta }({\mathbf{x}}_0)+{\mathbb{E}}_q[\log \frac{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})}+\log p_{\theta }({\mathbf{x}}_0)]\\ & ={\mathbb{E}}_q[\log \frac{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})}]\end{array}$$

那么我们就将对于图像的分布规约到了对于逆向过程去估计前向过程的分布。这一步就是VAE的$L_{VLB}$：

$$L_{\mathrm{VLB}}={\mathbb{E}}_{q({\mathbf{x}}_{0:T})}[\log \frac{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})}]\ge -{\mathbb{E}}_{q({\mathbf{x}}_0)}\log p_{\theta }({\mathbf{x}}_0)$$

但是上述虽然规约到了逆向过程，但是仍然不是方便计算的表达式，我们需要进一步的归约到可以算的KL散度中：

$$\begin{array}{rl}{L}_{\mathrm{VLB}}& ={\mathbb{E}}_{q({\mathbf{x}}_{0:T})}[\log \frac{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{0:T})}]\\ & ={\mathbb{E}}_q[\log \frac{\prod _{t=1}^{T}q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})}{p_{\theta }({\mathbf{x}}_T)\prod _{t=1}^T{p}_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}]\\ & ={\mathbb{E}}_q[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=1}^T\log \frac{q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}]\\ & ={\mathbb{E}}_q[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}+\log \frac{q({\mathbf{x}}_1|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}]\\ & ={\mathbb{E}}_q[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=2}^T\log (\frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}\cdot \frac{q({\mathbf{x}}_t|{\mathbf{x}}_0)}{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)})+\log \frac{q({\mathbf{x}}_1|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}]\\ & ={\mathbb{E}}_q[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_t|{\mathbf{x}}_0)}{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0)}+\log \frac{q({\mathbf{x}}_1|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}]\\ & ={\mathbb{E}}_q[-\log p_{\theta }({\mathbf{x}}_T)+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}+\log \frac{q({\mathbf{x}}_T|{\mathbf{x}}_0)}{q({\mathbf{x}}_1|{\mathbf{x}}_0)}+\log \frac{q({\mathbf{x}}_1|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}]\\ & ={\mathbb{E}}_q[\log \frac{q({\mathbf{x}}_T|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_T)}+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}-\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)]\\ & ={\mathbb{E}}_q[\underset{L_T}{\underbrace{D_{\mathrm{KL}}(q({\mathbf{x}}_T|{\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_T))}}+\sum _{t=2}^T\underset{L_{t-1}}{\underbrace{D_{\mathrm{KL}}(q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)\parallel p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t))}}\underset{L_0}{\underbrace{-\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)}}]\end{array}$$

只有中间一项是在训练过程中是变量，因此在优化过程中考虑中间一项即可，即优化这两个分布的KL散度。
我们设逆向分布为:

$$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))$$

具有两个参数${\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t)$.
我们需要预测这两个变量来拟合分布。在论文中采用固定的${\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t)$才能取得比较好的结果。
在对${\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t)$进行预测的时候仍然具有两种方法。

• 一种是直接预测${\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t)$,

• 另一种是根据之前推导出的相关公式：$$\begin{aligned}
  \boldsymbol{\mu}_\theta(\mathbf{x}_t,t)& =\color{red}{\frac1{\sqrt{\alpha_t}}\left(\mathrm{x}t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}t}}\boldsymbol{\epsilon}\theta(\mathrm{x}t,t)\right)} \
  \mathrm{Thus~}\mathbf{x}& =\mathcal{N}(\mathbf{x};\frac1{\sqrt{\alpha_t}}\Big(\mathbf{x}_t-\frac{1-\alpha_t}{\sqrt{1-\bar{\alpha}t}}\boldsymbol{\epsilon}\theta(\mathbf{x}t,t)\Big),\boldsymbol{\Sigma}\theta(\mathbf{x}_t,t))
  \end{aligned}$$
  可以直接预测${\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,t)$，然后根据上述公式计算${\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t)$。

如果对${ϵ}_t$进行预测，可以得到下述的更新函数：

$$\begin{array}{rl}{L}_t& ={\mathbb{E}}_{{\mathbf{x}}_0,ϵ}[\frac{1}{2\Vert {\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t){\Vert }_2^2}\Vert {\tilde{\mathit{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)-{\mathit{\mu }}_{\theta }({\mathbf{x}}_t,t){\Vert }^2]\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,ϵ}[\frac{1}{2\Vert {\mathbf{\Sigma }}_{\theta }{\Vert }_2^2}\Vert \frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{\mathit{ϵ}}_t)-\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\overline{\alpha }}_t}}{\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,t)){\Vert }^2]\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,\mathit{ϵ}}[\frac{(1-{\alpha }_t{)}^2}{2{\alpha }_t(1-{\overline{\alpha }}_t)\Vert {\mathbf{\Sigma }}_{\theta }{\Vert }_2^2}\Vert {\mathit{ϵ}}_t-{\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,t){\Vert }^2]\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,\mathit{ϵ}}[\frac{(1-{\alpha }_t{)}^2}{2{\alpha }_t(1-{\overline{\alpha }}_t)\Vert {\mathbf{\Sigma }}_{\theta }{\Vert }_2^2}\Vert {\mathit{ϵ}}_t-{\mathit{ϵ}}_{\theta }(\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\overline{\alpha }}_t}{\mathit{ϵ}}_t,t){\Vert }^2]\end{array}$$

如果忽略掉前面关于t的权重项，可以得到（忽略之后的实验效果更好）：

$$\begin{array}{rl}{L}_t^{\mathrm{simple}}& ={\mathbb{E}}_{t\sim [1,T],{\mathbf{x}}_0,{\mathit{ϵ}}_t}[\Vert {\mathit{ϵ}}_t-{\mathit{ϵ}}_{\theta }({\mathbf{x}}_t,t){\Vert }^2]\\ & ={\mathbb{E}}_{t\sim [1,T],{\mathbf{x}}_0,{\mathit{ϵ}}_t}[\Vert {\mathit{ϵ}}_t-{\mathit{ϵ}}_{\theta }(\sqrt{{\overline{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\overline{\alpha }}_t}{\mathit{ϵ}}_t,t){\Vert }^2]\end{array}$$