Classifier Guidance 与 Classifier-Free Guidance

Classifier Guidance 与 Classifier-Free Guidance

DDPM 终于把 diffusion 模型做 work 了，但无条件的生成在现实中应用场景不多，我们终归还是要可控的图像生成。本文简要介绍两篇关于 diffusion 模型可控生成的工作。其中 Classifier-Free Guidance 的方法还是现在多数条件生成 diffusion 模型的主流思路。

Classifier Guidance: Diffusion Models Beat GANs on Image Synthesis

Classifier-Free Guidance: Classifier-Free Diffusion Guidance

Classifier Guidance

要做可控生成，即条件生成，首先想到我们可以拿类别来作为条件，比如要指定类别猫，就生成猫的图片。也就是说要给定类别 $y$，生成图片 $x$，即 $P(x|y)$ 。而一般分类器做的事情正好是反过来，给定图片，预测类别，即 $P(y|x)$ 。这刚好是一对逆条件概率，应该敏锐地想到贝叶斯公式就是处理这类逆概率问题的。推导如下：
$$\begin{array}{rl}\nabla \log P(x|y)& =\nabla \log \frac{P(x)P(y|x)}{P(y)}\\ & =\nabla \log P(x)+\nabla \log P(y|x)-\nabla \log P(y)\\ & =\nabla \log P(x)+\nabla \log P(y|x)\end{array}$$
其中 $P(y)$ 是某个类别的先验概率，是一个常数，其梯度为 0，故直接丢掉。这里的 $\nabla P(x)$ 实际就是 score-base model 中所谓的 score，score-based model 实际可以看作是 diffusion model 的另一种形式，这里不展开。

在结果中，第一项 $\nabla \log P(x)$ 就是原本无条件生成的梯度，而第二项 $\nabla P(y|x)$ 则相当于是分类器进行图形分类的梯度。也就是说，我们只要在无条件生成的基础上，加上想要的类别的分类器梯度，作为引导（或者称为条件的梯度修正），就可以导出以类别作为条件的生成。

推导看起来并不复杂，具体怎么实现呢？怎么在生成的时候加入分类器的梯度作为引导呢？这里我们参考 OpenAI 原 Classifier Guidance 给出的代码来理解：

# https://github.com/openai/guided-diffusion/blob/main/scripts/classifier_sample.py#L54
# 核心就是这里的cond_fn函数

import torch as th
import torch.nn.functional as F
classifier = ... # 加载一个（噪声）图像分类器

def cond_fn(x, t, y=None):
    assert y is not None
    with th.enable_grad():
        x_in = x.detach().requires_grad_(True)
        logits = classifier(x_in, t)
        log_probs = F.log_softmax(logits, dim=-1)
        selected = log_probs[range(len(logits)), y.view(-1)]
        return th.autograd.grad(selected.sum(), x_in)[0] * args.classifier_scale

这里的 t 是当前时间步，x 是当前步的去噪结果图，y 是类别索引。我们看到，计算分类器梯度的过程其实很简单：

1. 首先把 x 和原始的梯度断开（detach），准备计算分类器的梯度
2. 把 x_in 和 t 都输入到分类其中，得到分类器预测的类别 logits
   • 注意，这里的分类器实际上需要是一个能够分类带噪声图像的分类器，不仅需要输入图像，还要输入当前时间步 t，相当于告知分类器当前噪声的强度。所以说，在 Classifier Guidance 的方法中，我们虽然不需要重新训练 diffusion 模型，但是我们需要单独训练一个噪声图像分类器。
3. 再把预测的类别 logits 过一下 softmax，得到各类别的概率 log_probs
4. 从 log_probs 中取出我们指定的类别 y 对应的概率，即 selected
5. 最后将 selected 中各个目标类别的概率值加在一起，希望该值越大越好，取该值对于 x 的梯度，即为分类器引导的梯度。

Classifier-Free Guidance

Classifier Guidance 的方法虽然不需要重新训练 diffusion 模型，但是需要额外的训练一个噪声图像分类器，并且在采样时需要额外的梯度引导。最关键的是，其作为可控生成的方法，对结果的控制能力十分有限，仅能够支持分类器所认识的有限类别，这无疑是不能满足我们多种多样的使用需求的。我们想要的肯定是 zeroshot 的，能够直接理解自然语言的可控生成。所幸在 CLIP 之后，图像文本特征已经能够在一定程度上对齐，CV 各个方向都基于 CLIP 实现了 zeroshot / open-vocab，图像生成也不例外。

Classifier-Free Guidance 的方法训练额外的分类器，并且，可以实现各种条件的引导生成。以最火爆的文生图为例，只要结合 CLIP 文本编码器提取 prompt 的文本特征 embedding，输入到 diffusion 模型中作为条件，即可实现。目前，Classifier-Free Guidance 已经成为条件生成的主流思路。

Classifier-Free Guidance 的想法是这样的：同时训练无条件生成模型和条件生成模型（实际上这俩是一个模型，只是训练时有概率输入是有条件的，有概率是无条件的），在推理时，同时 forward 带输入条件的生成和无条件的生成吗，然后把俩结果进行线性组合外推，得到最终的条件生成结果。

直接来看一下伪代码（参考 diffusers 的 API）：

unet = ... # 加载unet去噪模型
clip_model = ...  # 加载CLIP模型

text = "a cat"  # 文本条件
text_embeddings = clip_model.text_encode(text)  # 编码条件文本，cond
empty_embeddings = clip_model.text_encode("")  # 编码空文本，uncond
text_embeddings = torch.cat(empty_embeddings, text_embeddings)  # concat到一起，只做一次forward

input = torch.randn((1, 3, sample_size, sample_size), device="cuda") # 采样初始噪声

for t in scheduler.timesteps:
    # 用 unet 推理，预测噪声
    with torch.no_grad():
        # 这里同时预测出了有文本的和空文本的图像噪声
        noise_pred = unet(input, t, encoder_hidden_states=text_embeddings).sample

    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)  # 拆成无条件和有条件的噪声
    
		# Classifier-Free Guidance 引导 
    noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

    # 用预测出的 noise_pred 和 x_t 计算得到 x_t-1
    latents = scheduler.step(noise_pred, t, latents).prev_sample

代码里的写法是 $\overline{ϵ}={ϵ}_u+s({ϵ}_c-{ϵ}_u)$， 论文里的公式是 $\overline{ϵ}=(1+w){ϵ}_c-w{ϵ}_u$，二者是等价的，只是做了下变换 $s=1+w$。个人感觉代码里这个形式更好理解一点：${ϵ}_c-{ϵ}_w$ 表示从无条件到目标条件的一个方向，$w$ 是多大程度上考虑条件的系数，在无条件 ${ϵ}_u$ 的基础上，再朝目标类别移动一定距离，即：$\overline{ϵ}={ϵ}_u+s({ϵ}_c-{ϵ}_u)$。（也可能是作者大佬的思路我没领悟到

Classifier-Free Guidance 的做法看起来并不复杂，但有几个问题值得讨论（笔者自己也很不明白，希望有大佬指点一下）：

1. 为什么不像 cvae 一样直接把 embedding 丢进去做条件生成，而是非要同时训练无条件生成的情况，再做一个线性作何外推呢？
   • 可能是因为采样时现需要有一个无条件的基准，然后像目标条件的方向再修正？
2. 关于空门大佬提到的 Classifier-Free Guidance 破坏了 Neural Diffusion Operator 的准线性性质。
   • 提到在 $w$ 很大时，采样结果会崩掉，实践中确实这种现象。但笔者目前尚在学习，还无法完全理解。贴一下大佬的文章链接：Classifer-free Guidance 是万恶之源 。

这些问题有大佬了解，可以指点一下，或者介绍下应该去补充哪些理论知识来深化理解，感激不尽。