Guided Diffusion/Diffusion Models Beat GANs on Image Synthesis (Paper reading)

Guided Diffusion/Diffusion Models Beat GANs on Image Synthesis (Paper reading)

Prafulla Dhariwal, OpenAI, NeurlPS2021, Cited: 555, Code, Paper.

1. 前言

对于条件图像合成，我们通过分类器指导进一步提高样本质量：一种简单、计算效率高的方法，使用分类器的梯度来权衡样本质量的多样性。我们在 ImageNet 128×128 上实现了 2.97 的 FID，在 ImageNet 256×256 上实现了 4.59，在 ImageNet 512×512 上实现了 7.72。即使每个样本只有 25 次前向传播，我们也能匹配 BigGAN-deep，同时保持更好的分布覆盖。最后，我们发现分类器引导与上采样扩散模型结合得很好，在 ImageNet 512×512 上进一步将 FID 提高到 3.85。

2. 整体思想

在扩散模型的梯度上加上额外的条件梯度使扩散模型可以有条件生成。

    def cond_fn(x, t, y=None):
        assert y is not None
        with th.enable_grad():
        	# x_in 就是 x_t
            x_in = x.detach().requires_grad_(True)
            # 分类器对x_t的预测结果
            logits = classifier(x_in, t)
            # 将预测的结果转换成概率
            log_probs = F.log_softmax(logits, dim=-1)
            # 对于GT的概率
            selected = log_probs[range(len(logits)), y.view(-1)]
            # 将所有概率求和得到梯度修正
            return th.autograd.grad(selected.sum(), x_in)[0] * args.classifier_scale

3. 方法

首先，定义一个有条件的扩散过程$\hat{q}$，并假设$\hat{q}(y|x_0)$是已知的：
$$\begin{array}{rl}\hat{q}(x_0)& :=q(x_0)\\ \hat{q}(y|x_0)& :=\text{Known labels per sample}\\ \hat{q}(x_{t+1}|x_t,y)& :=q(x_{t+1}|x_t)\\ \hat{q}(x_{1:T}|x_0,y)& :=\prod _T^{t=1}\hat{q}(x_t|x_{t-1},y)\end{array}\text{\hspace{1em}(1)(2)(3)(4)}$$

(1)可知训练样本是不变的，(2)可知每一个样本都有已知标签，(3)可知加噪过程不变，(4)可知有条件的联合分布也是一个马尔可夫性质的。那么通过以上定义的假设，接下来分别证明

$\hat{q}(x_{t+1}|x_t)$的边缘分布可以表示为${\int }_y\hat{q}(x_{t+1},y|x_t)dy$，我们把$x_t$当作一个条件时，可以将上式看作边缘概率与条件概率相乘（多变量贝叶斯公式）${\int }_y\hat{q}(x_{t+1}|x_t,y)\hat{q}(y|x_t)dy$，继续推导：
$$\begin{array}{rl}{\int }_y\hat{q}(x_{t+1}|x_t,y)\hat{q}(y|x_t)dy& ={\int }_y\hat{q}(x_{t+1}|x_t)\hat{q}(y|x_t)dy\\ & =\hat{q}(x_{t+1}|x_t){\int }_y\hat{q}(y|x_t)dy,{\int }_y\hat{q}(y|x_t)dy=1\\ & =\hat{q}(x_{t+1}|x_t)=\hat{q}(x_{t+1}|x_t,y)\end{array}$$
相同的逻辑可以推断出联合分布$\hat{q}(x_{1:T}|x_0)=q(x_{1:T}|x_0)$。根据这个联合分布得：
$$\begin{array}{rl}\hat{q}(x_t)& ={\int }_{x_0:t-1}\hat{q}(x_0,...,x_t)d_{x_0:t-1}\\ & ={\int }_{x_0:t-1}\hat{q}(x_0)\hat{q}(x_1,...,x_t|x_0)d_{x_0:t-1}\\ & ={\int }_{x_0:t-1}q(x_0)q(x_1,...,x_t|x_0)d_{x_0:t-1}\\ & ={\int }_{x_0:t-1}q(x_0,...,x_t)d_{x_0:t-1}\\ & =q(x_t)\end{array}$$
由于$\hat{q}(x_t)=q(x_t)$且$\hat{q}(x_{t+1}|x_t)=q(x_{t+1}|x_t)$，由贝叶斯公式可知任意时刻$x_t$在$\hat{q}$和$q$中的表现是一样的。那么在有条件下的逆过程$\hat{q}(x_t|x_{t+1},y)$是什么样的呢？
$$\begin{array}{rl}\hat{q}(x_t|x_{t+1},y)& =\frac{\hat{q}(x_t,x_{t+1},y)}{\hat{q}(x_{t+1},y)}\\ & =\frac{\hat{q}(x_t,x_{t+1},y)}{\hat{q}(y|x_{t+1})\hat{q}(x_{t+1})}\\ & =\frac{\hat{q}(x_t|x_{t+1})\hat{q}(y|x_t,x_{t+1})\hat{q}(x_{t+1})}{\hat{q}(y|x_{t+1})\hat{q}(x_{t+1})}\\ & =\frac{\hat{q}(x_t|x_{t+1})\hat{q}(y|x_t,x_{t+1})}{\hat{q}(y|x_{t+1})}\end{array}\text{\hspace{1em}(5)(6)(7)(8)}$$
公式8中的$\hat{q}(y|x_t,x_{t+1})$：
$$\begin{array}{rl}\hat{q}(y|x_t,x_{t+1})& =\hat{q}(x_{t+1}|x_t,y)\frac{\hat{q}(y|x_t)}{\hat{q}(x_{t+1}|x_t)}\\ & =\hat{q}(x_{t+1}|x_t)\frac{\hat{q}(y|x_t)}{\hat{q}(x_{t+1}|x_t)}\\ & =\hat{q}(y|x_t)\end{array}\text{\hspace{1em}(9)(10)(11)}$$
则公式8可以表示为：$\frac{q(x_t|x_{t+1})\hat{q}(y|x_t)}{\hat{q}(y|x_{t+1})}$，且$\hat{q}(y|x_{t+1})$是不依赖于$x_t$的，则$\hat{q}(y|x_t,x_{t+1})\propto Z\hat{q}(x_t|x_{t+1})\hat{q}(y|x_t,x_{t+1})$，$Z$是一个归一化常数。$q$是一个训练好的扩散过程，因此可以用$p_{\theta }(x_t|x_{t+1})$表示，那么该如何获得$\hat{q}(y|x_t)$？我们可以在每一个时刻$x_t$，训练一个分类器$p_{\varphi }(y|x_t)$。接下来如何通过采样得到$x_0$呢？
$$p_{\theta ,\varphi }(x_t|x_{t+1},y)=Z{p}_{\theta }(x_t|x_{t+1})p_{\varphi }(y|x_t)\text{\hspace{1em}(12)}$$
在DDPM中，假设了$p_{\theta }(x_t|x_{t+1})=N(\mu ,\Sigma )$， 该多元高斯分布的似然函数为$log{p}_{\theta }(x_t|x_{t+1})=-\frac{1}{2}(x_t-\mu {)}^T{\Sigma }^{-1}(x_t-\mu )+C$ 。同时，我们假设$log{p}_{\varphi }(y|x_t)$ 与${\Sigma }^{-1}$有更低的曲率，也就是说二阶导数更小。那么我们可以用Taylor公式将$log{p}_{\varphi }(y|x_t)$展开到一阶导数：
$$\begin{array}{rl}log{p}_{\varphi }(y|x_t)& \approx log{p}_{\varphi }(y|x_t){|}_{x_t=\mu }+(x_t-\mu ){\nabla }_{x_t}log{p}_{\varphi }(y|x_t){|}_{x_t=\mu }\\ & =(x_t-\mu )g+C_1\end{array}\text{\hspace{1em}(13)}$$
其中$g={\nabla }_{x_t}log{p}_{\varphi }(y|x_t){|}_{x_t=\mu }$且$C_1$是个常数。 则：
$$\begin{array}{rl}log(p_{\theta }(x_t|x_{t+1})p_{\varphi }(y|x_t))& \approx -\frac{1}{2}(x_t-\mu {)}^T{\Sigma }^{-1}(x_t-\mu )+(x_t-\mu )g+C_2\\ & =-\frac{1}{2}(x_t-\mu -\Sigma g{)}^T{\Sigma }^{-1}(x_t-\mu -\Sigma g)+\frac{1}{2}{g}^T\Sigma g+C_2\\ & =-\frac{1}{2}(x_t-\mu -\Sigma g{)}^T{\Sigma }^{-1}(x_t-\mu -\Sigma g)+C_3\\ & =logp(z)+C_4,z\sim N(\mu +\Sigma g,\Sigma )\end{array}\text{\hspace{1em}(14)(15)(16)(17)}$$
公式14中的$\Sigma ={\Sigma }^T$可以推出公式15，而公式15中的第二项与$x_t$无关，则可以推出公式16，公式17说明了这个新的似然为均值为$\mu +\Sigma g$，方差不变的新分布，也就是说新的逆扩散条件分布$p_{\theta ,\varphi }$与DDPM中的逆过程仅仅相差在均值上 。那么将该方法应用在DDPM中，如下：

首先通过扩散模型中的神经网络获得均值和方差，可学习的方差为${\Sigma }_{\theta }(x_t,t)=exp(vlog({\beta }_t)+(1-v)log{\tilde{\beta }}_t)$其中，${\tilde{\beta }}_t=\frac{1-{\bar{\alpha }}_{t-1}}{1-{\bar{\alpha }}_t}{\beta }_t$。然后我们从新的分布中进行采样得到$x_{t-1}$，这里的$g$就是$x_t$时刻分类器的梯度。

4. 总结

条件扩散模型是相对热点的研究方向，本文的添加条件的方法是一种简单高效的方法，应用于low-level任务中的文章目前也被研究探讨，与之前的条件扩散模型相比，guided diffusion models对于一些特殊的任务更友好，而且openai的代码质量也很高。本文还有很多值得讨论的内容，如神经网络的修改，DDIM的应用等。