条件控制扩散模型

条件控制扩散模型

背景

扩散模型在生成图片时，都是 输入一个随机高斯噪声数据， 然后 逐步的产出一张有意义的真实图片。这个过程中每一步都是一个随机过程，所以每次执行产出的图片都不一样， 生成的图像多样性非常好。

但这也是一个缺点：生成的图像不可控，无法控制这个生成过程并令其生成我们想要的图像内容。

直觉的讲，我们可以在扩散过程中 引入额外的信息来指导整个扩散模型， 假设这个额外的信息为 $y$，它可以是一段文本、一张图片或者图像的类别标签，引入 $y$ 之后的模型就变成了一个以 $y$ 为条件的条件概率分布。
$$p(x_{1:T}|y,x_0)$$

自然而然地，接下来就需要探讨，引入 $y$ 之后对前向扩散过程和逆向采样过程分别有什么影响，需要做出什么调整。

首先看下对前向扩散过程的影响，先说结论：引入 $y$ 之后，对前向扩散过程没有任何影响。

其实，从直觉上讲，前向扩散过程是对原始图片加噪声，直至变成纯噪声，这个过程显然与 $y$ 没有任何关系。但做研究要严谨，还是需要给出数学证明的。证明过程在论文 Dhariwal et al. (2021) 中已经给出，这里直接搬过来。

• 首先定义带有小帽子 $\hat{q}$ 的符号表示引入 $y$ 之后的条件扩散模型中的概率分布，不带小帽子的符号 $q$ 表示原始（无条件）扩散模型的概率分布。在这里可以先把 $y$ 理解成原始图像 $x_0$ 的类别标签（label），它是已知的，即 训练输入是 $(x_0,y)$ 这样的数据对，类似于图像分类任务的训练样本。

• 先考虑，单一时刻 t 的变化，扩散模型本身是一个马尔科夫链，其定义中已经明确了：任意时刻的 $x_t$ 只和它上一个时刻的状态 $x_{t-1}$ 有关。因此任意单一时刻的前向扩散条件概率（加噪声）过程一定是与 $y$ 无关的，即下面的式子是一定成立的。
  $$\hat{q}(x_t|x_{t-1},y):=q(x_t|x_{t-1})$$

• 基于这个事实，可以推导出另一个等式，即
  $$\begin{array}{rl}\hat{q}(x_t|x_{t-1})& ={\int }_y\hat{q}(x_t,y|x_{t-1})dy\\ & ={\int }_y\hat{q}(x_t|y,x_{t-1})\hat{q}(y|x_{t-1})dy\\ & ={\int }_{y}q(x_t|x_{t-1})\hat{q}(y|x_{t-1})dy\\ & =q(x_t|x_{t-1})\underset{\text{这里积分为}1}{\underbrace{{\int }_y\hat{q}(y|x_{t-1})dy}}\\ & =q(x_t|x_{t-1})\\ & =\hat{q}(x_t|x_{t-1},y)\end{array}$$

• 基于相同的逻辑，可以推导出条件扩散模型中的前向扩散过程的联合概率 $\hat{q}(x_{1:T}|x_0)$ 和非条件（原始）扩散模型的前向扩散过程联合概率 $q(x_{1:T}|x_0)$ 是等价的。推导过程如下：
  $$\begin{array}{rl}\hat{q}(x_{1:T}|x_0)& ={\int }_y\hat{q}(x_{1:T},y|x_0)dy\\ & ={\int }_y\hat{q}(y|x_0)\hat{q}(x_{1:T}|x_0,y)dy\\ & ={\int }_y\hat{q}(y|x_0)\prod _{t=1}^T\hat{q}(x_t|x_{t-1},y)dy\\ & ={\int }_y\hat{q}(y|x_0)\prod _{t=1}^T\hat{q}(x_t|x_{t-1})dy\\ & =\prod _{t=1}^T\hat{q}(x_t|x_{t-1})\underset{\text{积分为}1}{\underbrace{{\int }_y\hat{q}(y|x_0)dy}}\\ & =\prod _{t=1}^T\hat{q}(x_t|x_{t-1})\\ & =q(x_{1:T}|x_0)\end{array}$$

到这里，证明了加不加条件 $y$ 对前向扩散过程没有任何影响，即：条件扩散模型的前向过程与非条件扩散模型的前向过程完全一样。

接下来就该讨论条件 $y$ 对逆向采样过程的影响，显然我们的目的就是要令 $y$ 对采样过程产生影响。

• 原始的无条件的扩散模型的逆向过程联合概率为：
  $$p(x_{0:T})=p(x_T)\prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t)$$
  添加额外信息 $y$ 之后，条件扩散模型的逆向过程联合概率为：
  $$p(x_{0:T}|y)=p(x_T)\prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t,y)$$

扩散模型有三种解释，即参数化模型有三种预测方案，分别是：

1. 直接预测原始的 $x_0$， ${\hat{x}}_{\theta }(x_t,t)\approx x_0$
2. 预测噪声数据， ${\widehat{ϵ}}_{\theta }(x_t,t)\approx {ϵ}_t$
3. 预测得分（梯度）， ${\hat{s}}_{\theta }(x_t,t)\approx {\nabla }_{x_t}\log p(x_t)$

最直觉（简单）的添加条件的方式，就是和添加 $t$ 一样，直接把 $y$ 作为模型的输入，添加到模型中即可，

上面三种方式分别改成：

1. 直接预测原始的 $x_0$， ${\hat{x}}_{\theta }(x_t,t,y)\approx x_0$
2. 预测噪声数据 ， ${\widehat{ϵ}}_{\theta }(x_t,t,y)\approx {ϵ}_t$
3. 预测得分（梯度）， ${\hat{s}}_{\theta }(x_t,t,y)\approx {\nabla }_{x_t}\log p(x_t|y)$

首先这么做是可以的，但效果可能并不好，以这种方式训练的条件扩散模型可能会学会忽略或淡化任何给定的条件信息 Luo et al. (2022)。

接下来看看不同的 guidance 方案具体是如何做的。

参考：条件控制扩散模型

Classifier Guidance

原理

OpenAI 的团队在 2021 年发表一篇论文 Dhariwal et al. (2021)：“A. Diffusion models beat GANs on image synthesis”，在这篇论文中，提出一种 利用图片类别标签指导图像生成 的方案，称为 classifier guidance，通过这种改进使扩散模型生成图像的质量大幅提升，并在 IS 和 FID 评分上超过了 GAN 模型，所以看论文的名字，简单直接。
论文的源码在： https://github.com/openai/guided-diffusion。

实际上这篇论文做了很多改进，比如对 UNET 也做了改进。但这里只关注 guidance 部分。

原论文的推导过程比较繁杂，这里我们采用另一篇文章 Luo et al. (2022) 的推导方案，直接从 score function 的角度去理解。

如上文所述，扩散模型的降噪过程的预测模型可以从三个角度去理解，其中一种角度是参数化神经网络模型预测的是 $x_t$ 的对数梯度，也叫得分（score），公式表示为
$${\hat{s}}_{\theta }(x_t,t)\approx {\nabla }_{x_t}\log p(x_t)$$

引入条件信息 $y$ 之后，相当于原来的 ${\nabla }_{x_t}\log p(x_t)$ 变成了 ${\nabla }_{x_t}\log p(x_t|y)$，所以只要推导出 ${\nabla }_{x_t}\log p(x_t|y)$ 是什么即可，推导过程如下：
$$\begin{array}{rl}{\nabla }_{x_t}\log p(x_t|y)& ={\nabla }_{x_t}\log \left(\frac{p(x_t)p(y|x_t)}{p(y)}\right)\\ & ={\nabla }_{x_t}\log p(x_t)+{\nabla }_{x_t}\log p(y|x_t)-\underset{\text{与 }{x}_t\text{ 无关，为 0}}{\underbrace{{\nabla }_{x_t}\log p(y)}}\\ & =\underset{\text{unconditional score}}{\underbrace{{\nabla }_{x_t}\log p(x_t)}}+\underset{\text{adversarial gradient}}{\underbrace{{\nabla }_{x_t}\log p(y|x_t)}}\end{array}$$

展开后有两项，

• 第一项 ${\nabla }_{x_t}\log p(x_t)$ 就是无条件扩散模型的梯度本身，称为 无条件分数（unconditional score）；
• 第二项 ${\nabla }_{x_t}\log p(y|x_t)$ 称为 对抗梯度（adversarial gradient）。

其中 $p(y|x_t)$ 显然代表一个分类器， $x_t$ 是分类器的输入， $y$ 是分类器预测的标签。

显然，条件扩散模型的分数项 ${\nabla }_{x_t}\log p(x_t|y)$ 就是在非条件扩散模型分数项 ${\nabla }_{x_t}\log p(x_t)$ 的基础上加了一项分类器的对抗梯度。
$$\text{condition score}=\text{unconditional score}+\text{adversarial gradient}$$

分类器 $p(y|x_t)$ 的输入不是初始的图像数据 $x_0$，而是对应时刻加了噪声的 $x_t$。

这个分类器需要在训练条件扩散模型之前，独立的训练好。当然此分类训练过程中，需要用扩散模型的前向加噪过程得到 $x_t$ 作为分类器的输入。

在实现时，

1. 预训练好一个噪声分类器 $p(y|x_t)$。
2. 在原来非条件扩散模型每一个时刻 $t$ 的采样过程中，
   1. 把模型的预测输出 ${\hat{s}}_{\theta }(x_t,t)$
   2. 同时把 $x_t$ 输入到预训练好的噪声分类器 $p(y|x_t)$，并计算出梯度 ${\nabla }_{x_t}\log p(y|x_t)$。
   3. 计算 ${\hat{s}}_{\theta }(x_t,t,y)={\hat{s}}_{\theta }(x_t,t)+{\nabla }_{x_t}\log p(y|x_t)$。
   4. 继续原来处理。

噪声分类器的对抗梯度 ${\nabla }_{x_t}\log p(y|x_t)$ 影响着采样降噪过程，它可以是采样过程向着类别标签 $y$ 的方向前进，进而得到贴近标签 $y$ 的图像，比如当 $y=\text{"小狗"}$ 时， 就会生成有小狗的图像，并且生成的图像更加逼真。

然而引入 classifier guidance 后，虽然生成图像的质量提高了，更加逼真，但是显然它会降低多样性。

为了能更灵活地控制生成内容的方向，论文中增加了一个超参数 $\lambda$，作为一个权重项。显然这个权重可以调节类别引导的作用大小。
$${\nabla }_{x_t}\log p(x_t|y)={\nabla }_{x_t}\log p(x_t)+\lambda {\nabla }_{x_t}\log p(y|x_t)$$

分类指导允许在样本多样性和保真度之间进行权衡。一个关于分类器引导的美丽的二维插图，

虽然引入 classifier guidance 效果很明显，但缺点也很明显：

1. 需要额外一个分类器模型，极大增加了成本，包括训练成本和采样成本。
2. 分类器的类别毕竟是有限集，不能涵盖全部情况，对于没有覆盖的标签类别会很不友好。

后来《More Control for Free! Image Synthesis with Semantic Diffusion Guidance》推广了“Classifier”的概念，使得它也可以按图、按文来生成。

Classifier-Guidance 方案的训练成本比较低（熟悉 NLP 的读者可能还会想起与之很相似的 PPLM 模型），但是推断成本会高些，而且控制细节上通常没那么到位。

Classifier-Free Guidance

无分类器指导方法通过 联合训练条件扩散模型和无条件扩散模型，消除了训练时需要额外分类器的需求。

原理

紧跟 OpenAI 之后，谷歌大脑团队发表了一篇论文 Classifier-Free Diffusion Guidance，双方战况激烈啊。Classifier-free 只需要在 classifier guidance 基础上稍微变换一下即可。

将公式 ${\nabla }_{x_t}\log p(x_t|y)={\nabla }_{x_t}\log p(x_t)+{\nabla }_{x_t}\log p(y|x_t)$ 移项，得到：
$${\nabla }_{x_t}\log p(y|x_t)={\nabla }_{x_t}\log p(x_t|y)-{\nabla }_{x_t}\log p(x_t)$$

然后将其带入${\nabla }_{x_t}\log p(x_t|y)={\nabla }_{x_t}\log p(x_t)+\lambda {\nabla }_{x_t}\log p(y|x_t)$ 中，可以推导得到：
$$\begin{array}{rl}{\nabla }_{x_t}\log p(x_t|y)& ={\nabla }_{x_t}\log p(x_t)+\gamma \left({\nabla }_{x_t}\log p(x_t|y)-{\nabla }_{x_t}\log p(x_t)\right)\\ & ={\nabla }_{x_t}\log p(x_t)+\gamma {\nabla }_{x_t}\log p(x_t|y)-\gamma {\nabla }_{x_t}\log p(x_t)\\ & =\underset{\text{conditional score}}{\underbrace{\gamma {\nabla }_{x_t}\log p(x_t|y)}}+\underset{\text{unconditional score}}{\underbrace{(1-\gamma ){\nabla }_{x_t}\log p(x_t)}}\end{array}$$

Classifier-free 的训练成本更大了一些，但效果更好一些。财大气粗的 OpenAI’s GLIDE、OpenAI’s DALL·E 2，以及 Google 的 Imagen 都采用了这个机制。

实战

Classifier-Free Guidance 是一种 控制扩散模型输出“更符合某个条件” 的方法，通过 同时训练有条件和无条件模型，再 在采样时进行插值 实现引导。

在条件扩散模型中，我们通常希望模型生成“满足条件”的样本，比如：

• 输入文字“猫”，生成一只猫的图像；
• 输入初始状态 p0，生成满足物理约束的状态序列 p，其中初始状态就是 p0；

原始的条件扩散方法需要一个 额外的分类器，比如 DDPM 的 Classifier Guidance 就是这么做的，但训练分类器困难、开销大。

下面是 Classifier-Free Guidance（CFG） 的核心思想，

训练阶段，让模型随机“忘记条件”，用概率 $p$ 将条件设为 None：

if random() < p:
    cond = None  # 无条件训练
else:
    cond = y      # 条件输入（如 label, prompt 等）

也就是说，模型学到的是：

• 条件模型 ${ϵ}_{\theta }(x_t,t,y)$
• 无条件模型 ${ϵ}_{\theta }(x_t,t)$

但它们是 同一个模型共享参数训练 的！只靠有无条件区分路径。

采样阶段，模型预测的噪声通过“插值”进行调整：

$${\widehat{ϵ}}_{\text{guided}}={ϵ}_{\theta }(x_t,t)+w\cdot \left({ϵ}_{\theta }(x_t,t,y)-{ϵ}_{\theta }(x_t,t)\right)$$

• ${ϵ}_{\theta }(x_t,t,y)$ 是有条件预测的噪声；
• ${ϵ}_{\theta }(x_t,t)$ 是无条件预测的噪声；
• $w$ 是 guidance scale（比如 3～10），控制引导强度。

类似于 “把输出拉向有条件结果方向”。