深度学习（生成式模型）——Classifier Guidance Diffusion

前言

几乎所有的生成式模型，发展到后期都需要引入"控制"的概念，可控制的生成式模型才能更好应用于实际场景。本文将总结《Diffusion Models Beat GANs on Image Synthesis》中提出的Classifier Guidance Diffusion（即条件扩散模型），其往Diffusion Model中引入了控制的概念，可以控制DDPM、DDIM生成指定类别（条件）的图片。

问题建模

本章节所有符号定义与DDPM一致，在条件$y$下的Diffusion Model的前向与反向过程可以定义为
$$\begin{array}{r}\hat{q}(x_{t+1}|x_t,y)\\ \hat{q}(x_t|x_{t+1},y)\end{array}$$
只要求出上述两个概率密度函数，我们即可按条件生成图像。

我们利用$\hat{q}$表示条件扩散模型的概率密度函数，$q$表示扩散模型的概率密度函数。

条件扩散模型的前向过程

对于前向过程，作者定义了以下等式
$$\begin{array}{rl}\hat{q}(x_0)& =q(x_0)\\ \hat{q}(x_{t+1}|x_t,y)& =q(x_{t+1}|x_t)\\ \hat{q}(x_{1:T}|x_0,y)& =\prod _{t=1}^T\hat{q}(x_t|x_{t-1},y)\end{array}$$

基于上述第二行定义，可知基于条件$y$的diffusion model的前向过程与普通的diffusion model一致，即$\hat{q}(x_{t+1}|x_t)=q(x_{t+1}|x_t)$。即加噪过程与条件$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,y|x_{t+1})}{\hat{q}(y|x_{t+1})}\\ & =\frac{\hat{q}(y|x_t,x_{t+1})\hat{q}(x_t|x_{t+1})}{\hat{q}(y|x_{t+1})}\end{array}\text{\hspace{1em}(1.0)}$$

已知条件扩散模型的前向过程与扩散模型一致，则有

$$\hat{q}(x_{1:T}|x_0)=q(x_{1:T}|x_0)$$

进而有
$$\begin{array}{rl}\hat{q}(x_t)& =\int \hat{q}(x_0,...,x_t)d{x}_{0:t-1}\\ & =\int \hat{q}(x_0)\hat{q}(x_{1:t}|x_0)d{x}_{0:t-1}\\ & =\int q(x_0)q(x_{1:t}|x_0)d{x}_{0:t-1}\\ & =q(x_t)\end{array}$$

对于$\hat{q}(x_t|x_{t+1})$，则有
$$\begin{array}{rl}\hat{q}(x_t|x_{t+1})& =\frac{\hat{q}(x_t,x_{t+1})}{\hat{q}(x_{t+1})}\\ & =\frac{\hat{q}(x_{t+1}|x_t)\hat{q}(x_t)}{\hat{q}(x_{t+1})}\\ & =\frac{q(x_{t+1}|x_t)q(x_t)}{q(x_{t+1})}\\ & =q(x_t|x_{t+1})\end{array}$$

对于$\hat{q}(y|x_t,x_{x_{t+1}})$，我们有
$$\begin{array}{rl}\hat{q}(y|x_t,x_{x_{t+1}})& =\frac{\hat{q}(x_{t+1}|x_t,y)\hat{q}(y|x_t)}{\hat{q}(x_{t+1}|x_t)}\\ & =\frac{\hat{q}(x_{t+1}|x_t)\hat{q}(y|x_t)}{\hat{q}(x_{t+1}|x_t)}\\ & =\hat{q}(y|x_t)\end{array}$$

因此式1.0为

$$\begin{array}{rl}\hat{q}(x_t|x_{t+1},y)& =\frac{\hat{q}(y|x_t,x_{t+1})\hat{q}(x_t|x_{t+1})}{\hat{q}(y|x_{t+1})}\\ & =\frac{\hat{q}(y|x_t)q(x_t|x_{t+1})}{\hat{q}(y|x_{t+1})}\end{array}$$

由于在反向过程中，$x_{t+1}$是已知的，因此$\hat{q}(y|x_{t+1})$也可看成已知值，设其倒数为$Z$，则有

$$\begin{array}{r}\hat{q}(x_t|x_{t+1},y)=Z\hat{q}(y|x_t)q(x_t|x_{t+1})\end{array}$$

取log可得
$$\begin{array}{r}\log \hat{q}(x_t|x_{t+1},y)=\log Z+\log \hat{q}(y|x_t)+\log \hat{q}(x_t|x_{t+1})\end{array}\text{\hspace{1em}(1.1)}$$

设$\hat{q}(x_t|x_{t+1})=\mathcal{N}({\mu }_t,{\sum }_t^2)$，则有
$$\log \hat{q}(x_t|x_{t+1})=-\frac{1}{2}(x_t-{\mu }_t{)}^T({\sum }_t{)}^{-1}(x_t-{\mu }_t)+C\text{\hspace{1em}(1.2)}$$

对于$\log \hat{q}(y|x_t)$，在$x_t={\mu }_t$处做泰勒展开，则有

$$\begin{array}{rl}\log \hat{q}(y|x_t)& \approx \log \hat{q}(y|x_t){|}_{x_t={\mu }_t}+(x_t-{\mu }_t){\nabla }_{x_t}\log \hat{q}(y|x_t){|}_{x_t={\mu }_t}\\ & =C_1+(x_t-{\mu }_t)g\end{array}\text{\hspace{1em}(1.3)}$$
其中$g={\nabla }_{x_t}\log \hat{q}(y|x_t){|}_{x_t={\mu }_t}$，结合式1.1、1.2、1.3，有

$$\begin{array}{rl}\log \hat{q}(x_t|x_{t+1},y)& \approx C_1+(x_t-{\mu }_t)g+\log Z-\frac{1}{2}(x_t-{\mu }_t{)}^T(\sum {}_t{)}^{-1}(x_t-{\mu }_t)+C\\ & =(x_t-{\mu }_t)g-\frac{1}{2}(x_t-{\mu }_t{)}^T(\sum {}_t{)}^{-1}(x_t-{\mu }_t)+C_2\\ & =-\frac{1}{2}(x_t-{\mu }_t-\sum {}_{t}g{)}^T(\sum {}_t{)}^{-1}(x_t-{\mu }_t-\sum {}_{t}g)+C_3\end{array}$$

最终有

$$\begin{array}{r}\hat{q}(x_t|x_{t+1},y)\approx \mathcal{N}({\mu }_t+{\sum }_{t}g,({\sum }_t{)}^2)\\ g={\nabla }_{x_t}\log \hat{q}(y|x_t){|}_{x_t={\mu }_t}\end{array}\text{\hspace{1em}(1.4)}$$

为了获得${\nabla }_{x_t}\log \hat{q}(y|x_t)$，Classifier Guidance Diffusion在训练好的Diffusion model的基础上额外训练了一个分类头。

假设$x_t\approx {\mu }_t$，则Classifier Guidance Diffusion的反向过程为:

其中$p_{\varphi }(y|x_t)=\hat{q}(y|x_t)$，$s$为一个超参数。

式1.4有个问题，当方差$\sum$取值为0时，$\sum {\nabla }_{x_t}\log \hat{q}(y|x_t)$取值将为0，无法控制生成指定条件的图像。因此式1.4不适用于DDIM等确定性采样的扩散模型。

在推导DDIM的采样公式前，我们先了解一下用Tweedie方法做参数估计的流程。

Tweedie方法主要用于指数族概率分布的参数估计，而高斯分布属于指数族概率分布，自然也适用。假设有一批样本$z$，则利用样本$z$估计高斯分布$\mathcal{N}(Z;\mu ,{\sum }^2)$的均值$\mu$的公式为

$$E[\mu |z]=z+{\sum }^2{\nabla }_z\log p(z)\text{\hspace{1em}(1.5)}$$

已知DDPM、DDIM的前向过程有

$$q(x_t|x_0)=\mathcal{N}(x_t;\sqrt{{\bar{\alpha }}_t}{x}_0,(1-{\bar{\alpha }}_t)\mathcal{I})\text{\hspace{1em}(1.6)}$$

结合式1.5、1.6可得

$$\begin{array}{r}\sqrt{{\bar{\alpha }}_t}{x}_0=x_t+(1-{\bar{\alpha }}_t){\nabla }_{x_t}\log p(x_t)\end{array}$$
进而有
$$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0-(1-{\bar{\alpha }}_t){\nabla }_{x_t}\log p(x_t)\text{\hspace{1em}(1.7)}$$
设${ϵ}_t$服从标准正态分布，则从式1.6可知

$$x_t=\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_t\text{\hspace{1em}(1.8)}$$

结合式1.7、1.8，则有

$${\nabla }_{x_t}\log p(x_t)=-\frac{1}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_t\text{\hspace{1em}(1.9)}$$

已知DDIM的采样公式为

$$x_{t-1}=\sqrt{{\bar{\alpha }}_{t-1}}\frac{x_t-\sqrt{1-{\bar{\alpha }}_t}{ϵ}_{\theta }(x_t)}{\sqrt{{\bar{\alpha }}_t}}+\sqrt{1-{\bar{\alpha }}_t-{\delta }_t^2}{ϵ}_{\theta }(x_t)\text{\hspace{1em}(2.0)}$$

结合式1.9、2.0可将DDIM的采样公式转变为

$$x_{t-1}=\sqrt{{\bar{\alpha }}_{t-1}}\frac{x_t-\sqrt{1-{\bar{\alpha }}_t}(-\sqrt{1-{\bar{\alpha }}_t}{\nabla }_{x_t}\log p(x_t))}{\sqrt{{\bar{\alpha }}_t}}+\sqrt{1-{\bar{\alpha }}_t-{\delta }_t^2}(-\sqrt{1-{\bar{\alpha }}_t}{\nabla }_{x_t}\log p(x_t))\text{\hspace{1em}(2.1)}$$

我们只需要将其中的${\nabla }_{x_t}\log p(x_t)$替换为${\nabla }_{x_t}\log p(x_t|y)$，即可引入条件$y$来控制DDIM的生成过程，利用贝叶斯定理，我们有

$$\begin{array}{rl}\log p(x_t|y)& =\log p(y|x_t)+\log p(x_t)-\log p(y)\\ {\nabla }_{x_t}\log p(x_t|y)& ={\nabla }_{x_t}\log p(y|x_t)+{\nabla }_{x_t}\log p(x_t)-{\nabla }_{x_t}\log p(y)\\ & ={\nabla }_{x_t}\log p(y|x_t)+{\nabla }_{x_t}\log p(x_t)\\ & ={\nabla }_{x_t}\log p(y|x_t)-\frac{1}{\sqrt{1-{\bar{\alpha }}_t}}{ϵ}_t\end{array}\text{\hspace{1em}(2.2)}$$
则有

$$-\sqrt{1-{\bar{\alpha }}_t}{\nabla }_{x_t}\log p(x_t|y)={ϵ}_t-\sqrt{1-{\bar{\alpha }}_t}{\nabla }_{x_t}\log p(y|x_t)\text{\hspace{1em}(2.3)}$$

至此，我们可以得到DDIM的采样流程为

对于DDIM等确定性采样的扩散模型，其应在训练好的Diffusion model的基础上额外训练了一个分类头，从而转变为Classifier Guidance Diffusion。

条件扩散模型的训练目标

注意到$\hat{q}(x_t|x_{t+1})=q(x_t|x_{t+1})$，并且上述的推导过程并没有改变$q(x_t|x_{t+1})、q(x_{t+1}|x_t)$的形式，因此Classifier Guidance Diffusion的训练目标与DDPM、DDIM是一致的，都可以拟合训练数据。