Diffusion Model 3 DDPM 逆扩散过程推导

回顾前向过程

从前向过程中，得到了扩散过程的分布：
$$\begin{array}{c}{x}_t=\sqrt{1-{\alpha }_t}\times {ϵ}_t+\sqrt{{\alpha }_t}\times x_{t-1}\\ x_t=\sqrt{1-\overline{\alpha }}\times ϵ+\sqrt{{\overline{\alpha }}_t}\times x_0\end{array}$$

重参数化：
$$q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=N(\sqrt{{\alpha }_t}{x}_{t-1},1-{\alpha }_t)，q({\mathbf{x}}_t|{\mathbf{x}}_0)=N(\sqrt{\overline{\alpha }}{x}_0,1-{\overline{\alpha }}_t)$$

逆扩散过程

随着步长$t$变大，数据样本${\mathbf{x}}_0$逐渐失去其可区分的特征。最终当 $T\to \infty$时，${\mathbf{x}}_T$相当于各向同性高斯分布。

各向同性的高斯分布是指在多元空间中，不同维度上的随机变量服从独立且具有相同的高斯分布。

如果我们可以反转上述过程并从$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$进行采样，我们将能够从高斯噪声输入${\mathbf{x}}_T\sim \mathcal{N}(0,\mathbf{I})$重新创建真实样本。请注意，如果${\beta }_t$足够小， $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$也将是高斯分布。

因此我们需要学习模型$p_{\theta }$来近似条件概率$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)$，以便运行反向扩散过程。

当以${\mathbf{x}}_0$为条件时，反向条件概率是容易处理的：
$$q({\mathbf{x}}_{t-1}|{\mathbf{x}}_t,{\mathbf{x}}_0)=\mathcal{N}({\mathbf{x}}_{t-1};\tilde{\mathbf{\mu }}({\mathbf{x}}_t,{\mathbf{x}}_0),{\tilde{\beta }}_t\mathbf{I})$$

对于正态分布，有如下正比关系：
$$N(\mu ,{\sigma }^2)\propto \exp \left[-\frac{{\left(x-\mu \right)}^2}{2{\sigma }^2}\right]$$

由贝叶斯可得：
$$\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{{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0{)}^2}{1-{\bar{\alpha }}_{t-1}}-\frac{({\mathbf{x}}_t-\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0{)}^2}{1-{\bar{\alpha }}_t}\left)\right)\\ & =\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{{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0{\mathbf{x}}_{t-1}+{\bar{\alpha }}_{t-1}{\mathbf{x}}_0^2}{1-{\bar{\alpha }}_{t-1}}-\frac{({\mathbf{x}}_t-\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0{)}^2}{1-{\bar{\alpha }}_t}\left)\right)\\ & =\exp (-\frac{1}{2}((\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}}){\mathbf{x}}_{t-1}^2-(\frac{2\sqrt{{\alpha }_t}}{{\beta }_t}{\mathbf{x}}_t+\frac{2\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0){\mathbf{x}}_{t-1}+C({\mathbf{x}}_t,{\mathbf{x}}_0)\left)\right)\end{array}$$

其中$C({\mathbf{x}}_t,{\mathbf{x}}_0)$是一些不涉及${\mathbf{x}}_{t-1}$的确定常数函数（不影响分布），详细信息被省略。$\propto \exp$为等价符号。

$a{x}^2+bx=a(x+\frac{b}{2a}{)}^2+C$，得到逆扩散过程的均值和方差分别为：
$$\begin{array}{rl}{\tilde{\beta }}_t& =1/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})=1/(\frac{{\alpha }_t-{\bar{\alpha }}_t+{\beta }_t}{{\beta }_t(1-{\bar{\alpha }}_{t-1})})=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}\cdot {\beta }_t\\ {\tilde{\mathbf{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)& =(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{\mathbf{x}}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0)/(\frac{{\alpha }_t}{{\beta }_t}+\frac{1}{1-{\bar{\alpha }}_{t-1}})\\ & =(\frac{\sqrt{{\alpha }_t}}{{\beta }_t}{\mathbf{x}}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}}{1-{\bar{\alpha }}_{t-1}}{\mathbf{x}}_0)\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}\cdot {\beta }_t\\ & =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{\mathbf{x}}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}{\mathbf{x}}_0\end{array}$$

在前向过程中我们推导得到:
$${\mathbf{x}}_0=\frac{1}{\sqrt{{\bar{\alpha }}_t}}({\mathbf{x}}_t-\sqrt{1-{\bar{\alpha }}_t}{\mathbf{ϵ}}_t)$$

所以均值改写为：
$$\begin{array}{rl}{\tilde{\mathbf{\mu }}}_t& =\frac{\sqrt{{\alpha }_t}(1-{\bar{\alpha }}_{t-1})}{1-{\bar{\alpha }}_t}{\mathbf{x}}_t+\frac{\sqrt{{\bar{\alpha }}_{t-1}}{\beta }_t}{1-{\bar{\alpha }}_t}\frac{1}{\sqrt{{\bar{\alpha }}_t}}({\mathbf{x}}_t-\sqrt{1-{\bar{\alpha }}_t}{\mathbf{ϵ}}_t)\\ & =\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_t)\end{array}$$

训练及损失

我们需要学习一个神经网络来近似反向扩散过程中的条件概率分布$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))$, 因为方差是由预先定义的加噪比例确定的，所以我们想要训练${\mathbf{\mu }}_{\theta }$来预测${\tilde{\mathbf{\mu }}}_t=\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_t)$，因为${\mathbf{x}}_t$是训练过程的中的输入项目，因此我们可以重新参数化高斯噪声项${\mathbf{ϵ}}_t$为${\mathbf{ϵ}}_{\theta }$，以时间步作为输入进行预测：

$$\begin{array}{rl}{\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t)& =\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t))\\ \text{Thus }{\mathbf{x}}_{t-1}& =\mathcal{N}({\mathbf{x}}_{t-1};\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)),{\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t))\end{array}$$

需要我们通过神经网络去预测，即这个网络输入$({\mathbf{x}}_t,t)$，预测出${\mathbf{x}}_{t-1}$到${\mathbf{x}}_t$加上的噪音。

损失项$L_t$被参数化以最小化与$\tilde{\mathbf{\mu }}$的差异:

$$\begin{array}{rl}{L}_t& ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}[\frac{1}{2\Vert {\mathbf{\Sigma }}_{\theta }({\mathbf{x}}_t,t){\Vert }_2^2}\Vert {\tilde{\mathbf{\mu }}}_t({\mathbf{x}}_t,{\mathbf{x}}_0)-{\mathbf{\mu }}_{\theta }({\mathbf{x}}_t,t){\Vert }^2]\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}[\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-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_t)-\frac{1}{\sqrt{{\alpha }_t}}({\mathbf{x}}_t-\frac{1-{\alpha }_t}{\sqrt{1-{\bar{\alpha }}_t}}{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t)\left){\Vert }^2\right]\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}[\frac{(1-{\alpha }_t{)}^2}{2{\alpha }_t(1-{\bar{\alpha }}_t)\Vert {\mathbf{\Sigma }}_{\theta }{\Vert }_2^2}\Vert {\mathbf{ϵ}}_t-{\mathbf{ϵ}}_{\theta }({\mathbf{x}}_t,t){\Vert }^2]\\ & ={\mathbb{E}}_{{\mathbf{x}}_0,\mathbf{ϵ}}[\frac{(1-{\alpha }_t{)}^2}{2{\alpha }_t(1-{\bar{\alpha }}_t)\Vert {\mathbf{\Sigma }}_{\theta }{\Vert }_2^2}\Vert {\mathbf{ϵ}}_t-{\mathbf{ϵ}}_{\theta }(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}{\mathbf{ϵ}}_t,t){\Vert }^2]\end{array}$$

根据经验论文经验发现，通过忽略权重项的简化目标来训练扩散模型效果更好：

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

最终的简单目标是：

$$L_{\text{simple}}=L_t^{\text{simple}}+C$$

其中$C$是不依赖于$\theta$的常量。

训练及采样算法

训练流程

采样流程