Diffusion Model & Conditional Diffusion Model
UPD 4.24.2022
- 最近朱军老师实验室Fan Bao等人进一步推进了理论结果,在优化时考虑了方差并给出了分析的结果
Background
其实diffusion model也在很久之前就出现了,我在Generative Model Roadmap这篇文章中也有简单的介绍,最近又被大家翻出来,这篇文章我将介绍一下Diffusion Model的数学形式以及最近出现的Conditional Diffusion Model。其中的相当一部分内容抄自Lilian Weng的一篇博客,其中的一些计算也加入了我自己的理解。
Unconditional Diffusion Model
Diffusion Model 事实上是一个正向的Markov过程,我们将从严格定义的正向过程出发,展示我们如何将一个数据集映射到标准高斯分布,以及推导出从标准高斯分布映射回原始数据集的逆向过程,当然其中有不可计算的东西,当然由万能的神经网络来完成这最后一块拼图。
Forward process
我们有一个随机变量$\mathbf{x}_{0} \sim q(\mathbf{x_o})$来表征一个数据集所代表的连续化的分布,我们将x考虑为离散的Markov过程,并且有这样的条件分布,这些条件分布的参数将由$\beta_t$完全表示
$$
q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}\right)=\mathcal{N}\left(\mathbf{x}_{t} ; \sqrt{1-\beta_{t}} \mathbf{x}_{t-1}, \beta_{t} \mathbf{I}\right) \quad q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)=\prod_{t=1}^{T} q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}\right)
$$
这么明显的服从高斯分布的独立增量可以让我们轻易地构造出一个布朗运动,并将离散的随机过程$\frac{x_t}{\sqrt{\bar{\alpha}_{t}}}$嵌入到该布朗运动中。
并且由于高斯分布的特性,我们可以直接从$x_0$得到$x_t$的分布,其中$\alpha_{t}:=1-\beta_{t}$,$\bar{\alpha}_{t}:=\prod_{s=0}^{t} \alpha_{s}$
\begin{equation}
q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)=\mathcal{N}\left(\mathbf{x}_{t} ; \sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0},\left(1-\bar{\alpha}_{t}\right) \mathbf{I}\right)
\end{equation}
Reverse Progress
首先我们观察正向过程的最后一个公式,可以看出其均值是递减的,方差是递增的,很容易选择合适的条件让$q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)$变成标准高斯分布,进而对$x_0$的积分变得Trivial,也就是$q(x_t)$也是标准高斯分布。注意到当$\beta_t$足够小时,$q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)$也是可以被当做高斯分布的,但是我们很难轻松地得到显示表达式,因为这需要我们对$x_0$做积分,因此我们考虑去学习一个模型$p_\theta$来去近似这样的Reverse过程的条件分布,进而我们可以从标准高斯分布来生成原始数据啦。
\begin{equation}
p_{\theta}\left(\mathbf{x}_{0: T}\right)=p\left(\mathbf{x}_{T}\right) \prod_{t=1}^{T} p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right) \quad p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)=\mathcal{N}\left(\mathbf{x}_{t-1} ; \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right), \mathbf{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)\right)
\end{equation}
注意到我们之前说了$x_t$是可以嵌入到一个布朗运动的,熟悉布朗运动的读者可以轻松地根据布朗运动关于两端中间某一时刻的条件分布仍然是高斯分布可以轻松得到(当然记不住结论的话也可以在明确了是高斯分布之后根据Lilian Weng的方法配方得到均值和方差)
\begin{equation}
\begin{aligned}
\tilde{\mu}_{t}\left(x_{t}, x_{0}\right) &:=\frac{\sqrt{\bar{\alpha}_{t-1}} \beta_{t}}{1-\bar{\alpha}_{t}} x_{0}+\frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right)}{1-\bar{\alpha}_{t}} x_{t} \\
\tilde{\beta}_{t} &:=\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}} \beta_{t} \\
q\left(x_{t-1} \mid x_{t}, x_{0}\right) &=\mathcal{N}\left(x_{t-1} ; \tilde{\mu}\left(x_{t}, x_{0}\right), \tilde{\beta}_{t} \mathbf{I}\right)
\end{aligned}
\end{equation}
又注意到我们可以设从$x_o$生成$x_t$的增量对应的随机变量为$z_t$,进而消掉$x_o$,我们的均值变为
\begin{equation}
\begin{aligned}
\tilde{\boldsymbol{\mu}}_{t} \left(x_{t}, x_{0}\right) &=\frac{\sqrt{\alpha_{t}}\left(1-\bar{\alpha}_{t-1}\right)}{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}}}\left(\mathbf{x}_{t}-\sqrt{1-\bar{\alpha}_{t}} \mathbf{z}_{t}\right) \\
&=\frac{1}{\sqrt{\alpha_{t}}}\left(\mathrm{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \mathbb{z}_{t}\right)
\end{aligned}
\end{equation}
之后就是关于训练的Loss了,由于我们无法直接优化log likelihood,我们可以抄一下VAE的变分下界,
\begin{equation}
\begin{aligned}
-\log p_{\theta}\left(\mathbf{x}_{0}\right) & \leq-\log p_{\theta}\left(\mathbf{x}_{0}\right)+D_{\mathrm{KL}}\left(q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right) \| p_{\theta}\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)\right) \\
&=-\log p_{\theta}\left(\mathbf{x}_{0}\right)+\mathbb{E}_{\mathbf{x}_{1: T} \sim q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)}\left[\log \frac{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0: T}\right) / p_{\theta}\left(\mathbf{x}_{0}\right)}\right] \\
&=-\log p_{\theta}\left(\mathbf{x}_{0}\right)+\mathbb{E}_{q}\left[\log \frac{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0: T}\right)}+\log p_{\theta}\left(\mathbf{x}_{0}\right)\right] \\
&=\mathbb{E}_{q}\left[\log \frac{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0: T}\right)}\right] \\
\text { Let } L_{\mathrm{VLB}} &=\mathbb{E}_{q\left(\mathbf{x}_{0: T}\right)}\left[\log \frac{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0: T}\right)}\right] \geq-\mathbb{E}_{q\left(\mathbf{x}_{0}\right)} \log p_{\theta}\left(\mathbf{x}_{0}\right)
\end{aligned}
\end{equation}
该下界也可以用JS不等式表出:
\begin{equation}
\begin{aligned}
L_{\mathrm{CE}} &=-\mathbb{E}_{q\left(\mathbf{x}_{0}\right)} \log p_{\theta}\left(\mathbf{x}_{0}\right) \\
&=-\mathbb{E}_{q\left(\mathbf{x}_{0}\right)} \log \left(\int p_{\theta}\left(\mathbf{x}_{0: T}\right) d \mathbf{x}_{1: T}\right) \\
&=-\mathbb{E}_{q\left(\mathbf{x}_{0}\right)} \log \left(\int q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right) \frac{p_{\theta}\left(\mathbf{x}_{0: T}\right)}{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)} d \mathbf{x}_{1: T}\right) \\
&=-\mathbb{E}_{q\left(\mathbf{x}_{0}\right)} \log \left(\mathbb{E}_{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)} \frac{p_{\theta}\left(\mathbf{x}_{0: T}\right)}{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)}\right) \\
& \leq-\mathbb{E}_{q\left(\mathbf{x}_{0: T}\right)} \log \frac{p_{\theta}\left(\mathbf{x}_{0: T}\right)}{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)} \\
&=\mathbb{E}_{q\left(\mathbf{x}_{0: T}\right)}\left[\log \frac{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0: T}\right)}\right]=L_{\mathrm{VLB}}
\end{aligned}
\end{equation}
这样的表达依旧不方便计算,为了令其可计算,也有如下的推导变到KL散度:
\begin{equation}
\begin{aligned}
L_{\mathrm{VLB}} &=\mathbb{E}_{q\left(\mathbf{x}_{0: T}\right)}\left[\log \frac{q\left(\mathbf{x}_{1: T} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0: T}\right)}\right] \\
&=\mathbb{E}_{q}\left[\log \frac{\prod_{t=1}^{T} q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}\right)}{p_{\theta}\left(\mathbf{x}_{T}\right) \prod_{t=1}^{T} p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)}\right] \\
&=\mathbb{E}_{q}\left[-\log p_{\theta}\left(\mathbf{x}_{T}\right)+\sum_{t=1}^{T} \log \frac{q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}\right)}{p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)}\right] \\
&=\mathbb{E}_{q}\left[-\log p_{\theta}\left(\mathbf{x}_{T}\right)+\sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}\right)}{p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)}+\log \frac{q\left(\mathbf{x}_{1} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)}\right] \\
&=\mathbb{E}_{q}\left[-\log p_{\theta}\left(\mathbf{x}_{T}\right)+\sum_{t=2}^{T} \log \left(\frac{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)} \cdot \frac{q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{0}\right)}\right)+\log \frac{q\left(\mathbf{x}_{1} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)}\right] \\
&=\mathbb{E}_{q}\left[-\log p_{\theta}\left(\mathbf{x}_{T}\right)+\sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)}+\sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t} \mid \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{0}\right)}+\log \frac{q\left(\mathbf{x}_{1} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)}\right] \\
&=\mathbb{E}_{q}\left[-\log p_{\theta}\left(\mathbf{x}_{T}\right)+\sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)}+\log \frac{q\left(\mathbf{x}_{T} \mid \mathbf{x}_{0}\right)}{q\left(\mathbf{x}_{1} \mid \mathbf{x}_{0}\right)}+\log \frac{q\left(\mathbf{x}_{1} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)}\right] \\
&=\mathbb{E}_{q}\left[\log \frac{q\left(\mathbf{x}_{T} \mid \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{T}\right)}+\sum_{t=2}^{T} \log \frac{q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right)}{p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)}-\log p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)\right]
\end{aligned}
\end{equation}
也就是
\begin{equation}
=\mathbb{E}_{q}[\underbrace{D_{\mathrm{KL}}\left(q\left(\mathbf{x}_{T} \mid \mathbf{x}_{0}\right) \| p_{\theta}\left(\mathbf{x}_{T}\right)\right)}_{L_{T}}+\sum_{t=2}^{T} \underbrace{D_{\mathrm{KL}}\left(q\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}, \mathbf{x}_{0}\right) \| p_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)\right)}_{L_{t-1}}\underbrace{-\log p_{\theta}\left(\mathbf{x}_{0} \mid \mathbf{x}_{1}\right)}_{L_{0}}]
\end{equation}
每个KL散度都是高斯分布之间的KL散度,我们有如下的闭形式
\begin{equation}
D_{\mathrm{KL}}\left(\mathcal{N}_{0} \| \mathcal{N}_{1}\right)=\frac{1}{2}\left(\operatorname{tr}\left(\Sigma_{1}^{-1} \Sigma_{0}\right)-k+\left(\mu_{1}-\mu_{0}\right)^{\top} \Sigma_{1}^{-1}\left(\mu_{1}-\mu_{0}\right)+\ln \left(\frac{\operatorname{det} \Sigma_{1}}{\operatorname{det} \Sigma_{0}}\right)\right)
\end{equation}
而Ho等人将$L_0$使用独立的离散解码器去建模,之后会详细介绍。
How to Train?
我们有了Loss的表达式,我们将如何训练呢?我们使用如下的符号去表征一步的Reverse Progress:
\begin{equation}
\boldsymbol{p}_{\theta}\left(\mathbf{x}_{t-1} \mid \mathbf{x}_{t}\right)=\mathcal{N}\left(\mathbf{x}_{t-1} ; \boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right), \mathbf{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)\right)
\end{equation}
那么我们的神经网络的最终目标是去学习 $\boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right), \mathbf{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)$
我们考虑$L_t$第一项的高斯分布的均值形式,我们不妨去fit它的形式,令
\begin{equation}
\boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right)=\frac{1}{\sqrt{\alpha_{t}}}\left(\mathbf{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \mathbb{z}_{\theta}\left(\mathbf{x}_{t}, t\right)\right)
\end{equation}
于是我们的$L_t$这一项KL散度可以使用如下的方式表示
\begin{equation}
\begin{aligned}
L_{t} &=\mathbb{E}_{\mathbf{x}_{0}, \mathbf{z}_t}\left[\frac{1}{2\left\|\boldsymbol{\Sigma}_{\theta}\left(\mathbf{x}_{t}, t\right)\right\|_{2}^{2}}\left\|\tilde{\boldsymbol{\mu}}_{t}\left(\mathbf{x}_{t}, \mathbf{x}_{0}\right)-\boldsymbol{\mu}_{\theta}\left(\mathbf{x}_{t}, t\right)\right\|^{2}\right] \\
&=\mathbb{E}_{\mathbf{x}_{0}, \mathbf{z}_t}\left[\frac{1}{2\left\|\boldsymbol{\Sigma}_{\theta}\right\|_{2}^{2}}\left\|\frac{1}{\sqrt{\alpha_{t}}}\left(\mathbf{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \mathbf{z}_{t}\right)-\frac{1}{\sqrt{\alpha_{t}}}\left(\mathbf{x}_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}} \mathbf{z}_{\theta}\left(\mathbf{x}_{t}, t\right)\right)\right\|^{2}\right] \\
&=\mathbb{E}_{\mathbf{x}_{0}, \mathbf{z}_t}\left[\frac{\beta_{t}^{2}}{2 \alpha_{t}\left(1-\bar{\alpha}_{t}\right)\left\|\mathbf{\Sigma}_{\theta}\right\|_{2}^{2}}\left\|\mathbf{z}_{t}-\mathbf{z}_{\theta}\left(\mathbf{x}_{t}, t\right)\right\|^{2}\right] \\
&=\mathbb{E}_{\mathbf{x}_{0}, \mathbf{z}_t}\left[\frac{\beta_{t}^{2}}{2 \alpha_{t}\left(1-\bar{\alpha}_{t}\right)\left\|\mathbf{\Sigma}_{\theta}\right\|_{2}^{2}}\left\|\mathbf{z}_{t}-\mathbf{z}_{\theta}\left(\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \mathbf{z}_{t}, t\right)\right\|^{2}\right]
\end{aligned}
\end{equation}
后来实践的时候,Ho等人发现抛弃系数反而有更好的效果:
\begin{equation}
L_{t}^{\text {simple }}=\mathbb{E}_{\mathbf{x}_{0}, \mathbf{z}_{t}}\left[\left\|\mathbf{z}_{t}-\mathbf{z}_{\theta}\left(\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}} \mathbf{z}_{t}, t\right)\right\|^{2}\right]
\end{equation}
这样我们统一了$L_0$到$L_{T-1}$,我们可以采样去训练,有了如下漂亮的算法
DDIM
事实上我们注意到,正向的Markov过程这个条件有点过强了,一方面在实际操作时也不是一步一步叠加噪声的,另一方面最后的Loss的推导也是无关于Markov的, 正向Markov唯一的作用是让我们有了$q\left(x_{t-1} \mid x_{t}, x_{0}\right) $这样的一个解析表达式. 那么我们为什么不放弃Markov的假设, 保留$q(x_t, x_0)$的边缘分布, 构造出合适的$q\left(x_{t-1} \mid x_{t}, x_{0}\right) $呢? Denoising Diffusion Implicit Models 这篇文章就是这么干的. 并且给出了一个可以自由选择$q\left(x_{t-1} \mid x_{t}, x_{0}\right) $方差的一个通用的表达式,并且保持$q(x_t, x_0)$的边缘分布和之前一致.
我们保持边缘分布为$q_{\sigma}\left(\boldsymbol{x}_{T} \mid \boldsymbol{x}_{0}\right)=\mathcal{N}\left(\sqrt{\bar{\alpha}_{T}} \boldsymbol{x}_{0},\left(1-\bar{\alpha}_{T}\right) \boldsymbol{I}\right)$, 设计出了这样一个更加自由控制的条件分布:
$$
q_{\sigma}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_{t}, \boldsymbol{x}_{0}\right)=\mathcal{N}\left(\sqrt{\bar{\alpha}_{t-1}} \boldsymbol{x}_{0}+\sqrt{1-\bar{\alpha}_{t-1}-\sigma_{t}^{2}} \cdot \frac{\boldsymbol{x}_{t}-\sqrt{\bar{\alpha}_{t}} \boldsymbol{x}_{0}}{\sqrt{1-\bar{\alpha}_{t}}}, \sigma_{t}^{2} \boldsymbol{I}\right)
$$
一个很直觉的构造想法是保持方差像这样
$$
\begin{aligned}
\mathbf{x}_{t-1} &=\sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t-1}} \mathbf{z}_{t-1} \\
&=\sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t-1}-\sigma_{t}^{2}} \mathbf{z}_{t}+\sigma_{t} \mathbf{z} \\
&=\sqrt{\bar{\alpha}_{t-1}} \mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t-1}-\sigma_{t}^{2}} \frac{\mathbf{x}_{t}-\sqrt{\bar{\alpha}_{t}} \mathbf{x}_{0}}{\sqrt{1-\bar{\alpha}_{t}}}+\sigma_{t} \mathbf{z}
\end{aligned}
$$
另外虽然这个有点像reverse progress, 但正如我们在DDPM中由正向公式推导出该公式一样, 我们同样可以推导出在这种情况下的forward progress, 尽管其不再是Markov的. 因此我们依旧用这个表征正向过程, 或者说指导我们如何从数据集构造用于训练的数据(事实上也并不需要, 我们本身用的也是$q\left(x_{t-1} \mid x_{t}, x_{0}\right) $这个分布).
在训练时我们将和之前一样, 我们的神经网络是从加噪图片中预测出噪声, 这篇文章给了一个全新的视角,尽管跟之前是一样的算法.
我们先根据预测出的噪声,直接根据$q(x_t \mid x_0)$预测出$x_0$:
$$
f_{\theta}^{(t)}\left(\boldsymbol{x}_{t}\right):=\left(\boldsymbol{x}_{t}-\sqrt{1-\bar{\alpha}_{t}} \cdot \epsilon_{\theta}^{(t)}\left(\boldsymbol{x}_{t}\right)\right) / \sqrt{\bar{\alpha}_{t}}
$$
之后再使用和$q_{\sigma}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_{t}, \boldsymbol{x}_{0}\right)$一样的形式去迭代
$$
p_{\theta}^{(t)}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_{t}\right)= \begin{cases}\mathcal{N}\left(f_{\theta}^{(1)}\left(\boldsymbol{x}_{1}\right), \sigma_{1}^{2} \boldsymbol{I}\right) & \text { if } t=1 \\ q_{\sigma}\left(\boldsymbol{x}_{t-1} \mid \boldsymbol{x}_{t}, f_{\theta}^{(t)}\left(\boldsymbol{x}_{t}\right)\right) & \text { otherwise }\end{cases}
$$
同时作者发现DDIM的一个好处是可以令$\sigma = 0$从而构造出一个确定的映射关系, 并且这样的模型可以在inference时有更少的数目.
Conditional Diffusion Model
之后来介绍我比较关心的如何让这样一个无条件的生成器中加入条件生成的功能。
Classifier Guidance
Ho等人提出了一种基于单独训练的分类器指导的条件生成的方法,该算法比较直观。我们先给出一个比较关键的公式,其中$\propto$右端的定义和之前提出的定义一致,$p_{\phi}\left(y \mid x_{t}\right)$给出了给定$x_t$后其标签的分布, $C_0$是一个归一化的常数,而左侧的严格定义将在下面给出。
\begin{equation}
p_{\theta, \phi}\left(x_{t} \mid x_{t+1}, y\right) = C_0 p_{\theta}\left(x_{t} \mid x_{t+1}\right) p_{\phi}\left(y \mid x_{t}\right)
\end{equation}
下面我们将构建带有conditional的Diffusion Model 以及上述公式的证明(参考自Ho的附录H)
我们给出一个很强的假设
\begin{equation}
q\left(x_{t+1} \mid x_{t}, y\right):=q\left(x_{t+1} \mid x_{t}\right)
\end{equation}
于是当我们给出了一个很强的全局lable $y$时,整个生成模型将保持和以前一致。
我们在这样的强假设的条件下,可以推导出label对随机过程的conditional分布也具有Markov性:
\begin{equation}
\begin{aligned}
q\left(y \mid x_{t}, x_{t+1}\right) &=q\left(x_{t+1} \mid x_{t}, y\right) \frac{q\left(y \mid x_{t}\right)}{q\left(x_{t+1} \mid x_{t}\right)} \\
&=q\left(x_{t+1} \mid x_{t}\right) \frac{q\left(y \mid x_{t}\right)}{q\left(x_{t+1} \mid x_{t}\right)} \\
&=q\left(y \mid x_{t}\right)
\end{aligned}
\end{equation}
同时我们有如下的推导:
\begin{equation}
\begin{aligned}
q\left(x_{t} \mid x_{t+1}, y\right) &=\frac{q\left(x_{t}, x_{t+1}, y\right)}{q\left(x_{t+1}, y\right)} \\
&=\frac{q\left(x_{t} \mid x_{t+1}\right) q\left(y \mid x_{t}, x_{t+1}\right) q\left(x_{t+1}\right)}{q\left(y \mid x_{t+1}\right) q\left(x_{t+1}\right)} \\
&=\frac{q\left(x_{t} \mid x_{t+1}\right) q\left(y \mid x_{t}\right)}{q\left(y \mid x_{t+1}\right)}
\end{aligned}
\end{equation}
因此我们需要用神经网络$p_{\phi}\left(y \mid x_{t}\right)$去拟合$q\left(y \mid x_{t}\right)$,可以看出这只是一个加噪图像的分类器,可以直接根据加噪的图片和标签进行训练。
下面我们将继续推导有了如上的分类器之后我们将如何使用DDPM进行条件生成。
我们考虑一步的Reverse Progress:
\begin{equation}
\begin{aligned}
p_{\theta}\left(x_{t} \mid x_{t+1}\right) &=\mathcal{N}(\mu, \Sigma) \\
\log p_{\theta}\left(x_{t} \mid x_{t+1}\right) &=-\frac{1}{2}\left(x_{t}-\mu\right)^{T} \Sigma^{-1}\left(x_{t}-\mu\right)+C
\end{aligned}
\end{equation}
其中的假设$\Sigma$是相当小的,因此我们可以考虑对$\log p_{\phi}(y \mid\ x_t)$在$x_t = \mu$进行Taylor展开保留一阶小量:
\begin{equation}
\begin{aligned}
\log p_{\phi}\left(y \mid x_{t}\right) &\left.\approx \log p_{\phi}\left(y \mid x_{t}\right)\right|_{x_{t}=\mu}+\left.\left(x_{t}-\mu\right) \nabla_{x_{t}} \log p_{\phi}\left(y \mid x_{t}\right)\right|_{x_{t}=\mu} \\
&=\left(x_{t}-\mu\right) g+C_{1}
\end{aligned}
\end{equation}
于是我们在生成时考虑最大化log likelihood $p_{\theta, \phi}\left(x_{t} \mid x_{t+1}, y\right)$
在计算时忽略$\Sigma$的一阶小量,我们有如下的推导
\begin{equation}
\begin{aligned}
\log \left(p_{\theta}\left(x_{t} \mid x_{t+1}\right) p_{\phi}\left(y \mid x_{t}\right)\right) & \approx-\frac{1}{2}\left(x_{t}-\mu\right)^{T} \Sigma^{-1}\left(x_{t}-\mu\right)+\left(x_{t}-\mu\right) g+C_{2} \\
&=-\frac{1}{2}\left(x_{t}-\mu-\Sigma g\right)^{T} \Sigma^{-1}\left(x_{t}-\mu-\Sigma g\right)+\frac{1}{2} g^{T} \Sigma g+C_{2} \\
&=-\frac{1}{2}\left(x_{t}-\mu-\Sigma g\right)^{T} \Sigma^{-1}\left(x_{t}-\mu-\Sigma g\right)+C_{3} \\
&=\log p(z)+C_{4}, z \sim \mathcal{N}(\mu+\Sigma g, \Sigma)
\end{aligned}
\end{equation}
也就是我们最终有了如下漂亮的条件生成的算法
Langevin Dynamics & Score-based Model: Another Pointview
除了上面一个的角度之外, 还有另外一个完全不同的视角来看待这个问题, 这里我们基本上是在抄另外一篇博客. 一个分布$p(x)$如果我们无法得到,但是我们能够得到score function$\nabla_x \log (p(x))$的话, Langevin动力学告诉我们我们依然可以根据$p(x)$来进行采样, 只不过需要通过一个迭代过程.
$$
\mathbf{x}_{i+1} \leftarrow \mathbf{x}_{i}+\epsilon \nabla_{\mathbf{x}} \log p(\mathbf{x})+\sqrt{2 \epsilon} \mathbf{z}_{i}, \quad i=0,1, \cdots, K
$$
但是呢, 如果直接去进行对$p(x)$的学习, 会遇到各种各样的问题, 一方面在数据特别稀疏的地方我们理想的$p(x)$和给出的数据点的gap比较大, 另外一方面在高维的数据中数据显然是几乎必然稀疏的, 或者说使其dense的Hausdorff维数是远小于数据所表现出来的维度的. 那么我们希望加入高斯噪声来让其在其所表现的维度上dense起来, 于是一个很天然的想法就是加入不同level的噪声扰动. 而如果加入的噪声是独立的高斯噪声的话, 我们完全可以通过加入另外单独采样自高斯的噪声来实现这一点. 我们记我们加入了方差为$\sigma_i$的噪声为$\boldsymbol{p}_{\sigma_{i}}(\mathbf{x})$, 那么我们需要神经网络通过Fisher divergence来学习对应的score function:
$$
\sum_{i=1}^{L} \lambda(i) \mathbb{E}_{p_{\sigma_{i}}(\mathbf{x})}\left[\left\|\nabla_{\mathbf{x}} \log p_{\sigma_{i}}(\mathbf{x})-\mathbf{s}_{\theta}(\mathbf{x}, i)\right\|_{2}^{2}\right]
$$
但是呢我们的目标梯度也不是可以直接拿过来的, 可是一些计算告诉我们, 在某些条件下优化上面的东西就等价于优化下面的这个
$$
\mathbb{E}_{p_{\sigma_i}(\mathbf{x})}\left[\operatorname{tr}\left(\nabla_{\mathbf{x}} \mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}, i)\right)+\frac{1}{2}\left\|\mathbf{s}_{\boldsymbol{\theta}}(\mathbf{x}, i)\right\|_{2}^{2}\right]
$$
又又又但是呢,这个东西的第一项计算成本依旧高昂, 我们又有方法规避掉这个, 比如所谓的Denoising score matching方法, 我们给数据加上已知的噪声$q_{\sigma}(\tilde{\mathbf{x}} \mid \mathbf{x})$, 然后呢就可以按这个算
$$
\frac{1}{2} \mathbb{E}_{q_{\sigma}(\tilde{\mathbf{x}} \mid \mathbf{x}) p_{\mathrm{data}}(\mathbf{x})}\left[\left\|\mathbf{s}_{\boldsymbol{\theta}}(\tilde{\mathbf{x}})-\nabla_{\tilde{\mathbf{x}}} \log q_{\sigma}(\tilde{\mathbf{x}} \mid \mathbf{x})\right\|_{2}^{2}\right]
$$
其中的$ p_{\mathrm{data}}(\mathbf{x})$代表真实的数据分布,
我们给每个噪声选了高斯分布之后, 我们的Loss 就长成了这样
$$
\ell(\boldsymbol{\theta} ; \sigma) \triangleq \frac{1}{2} \mathbb{E}_{p_{\text {data }}(\mathbf{x})} \mathbb{E}_{\tilde{\mathbf{x}} \sim \mathcal{N}\left(\mathbf{x}, \sigma^{2} I\right)}\left[\left\|\mathbf{s}_{\boldsymbol{\theta}}(\tilde{\mathbf{x}}, \sigma)+\frac{\tilde{\mathbf{x}}-\mathbf{x}}{\sigma^{2}}\right\|_{2}^{2}\right]
$$
整体上是
$$
\mathcal{L}\left(\boldsymbol{\theta} ;\left\{\sigma_{i}\right\}_{i=1}^{L}\right) \triangleq \frac{1}{L} \sum_{i=1}^{L} \lambda\left(\sigma_{i}\right) \ell\left(\boldsymbol{\theta} ; \sigma_{i}\right)
$$
然后我们有这样的算法去采样, 随机步伐越来越小, 每个噪声多次采样
后来这个作者用SDE的角度将他这种方差变到很大的方法(SMLD)和主流的映射到标准高斯分布的方法(DDPM)统一了起来, 分别叫做Variance Exploding (VE) SDE 和 Variance Preserving (VP) SDE.
原文链接:Diffusion Model & Conditional Diffusion Model
Skirrey's Blog 版权所有,转载请注明出处。
还没有任何评论,你来说两句吧!