Paper: Your Diffusion Model is Secretly a Zero-Shot Classifier

Website: diffusion-classifier.github.io/

相关阅读

• diffusion model（一）DDPM技术小结 (denoising diffusion probabilistic)
• diffusion model（二）—— DDIM技术小结
• diffusion model（三）—— classifier guided diffusion model
• diffusion model（四）文生图diffusion model（classifier-free guided）
• diffusion model（五）stable diffusion底层原理（latent diffusion model， LDM

背景

最近，出现了一系列大规模的文生图模型，它们极大地增强了我们通过文字生成图片的能力。这些模型可以根据各种提示生成逼真的图片，展现出惊人的综合创作能力。到目前为止，几乎所有的应用都只关注了模型的生成功能，但实际上，这些模型还能提供条件密度估计，这对于处理图像生成之外的任务也很有用。

本篇文章指出类似stable diffusion这样的大规模文本转图像模型所计算出的密度估计，可以被用来进行“零样本分类” (zero-shot classification)，而不需要额外的训练。

方法大意

diffusion model的背景知识

从前面diffusions系列文章中我们知道，diffuison model的去噪过程是一个马尔可夫过程
$$\begin{gathered} p_{\theta }(x_0)=\int p_{\theta }(x_{0:T})d{x}_{1:T} \\ \text{其中:}{p}_{\theta }(x_{0:T}):=p_{\theta }(x_T)\prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t) \\ {p}_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\sigma }_t\text{I})\text{\hspace{1em}(1)} \end{gathered}$$
即：
$$\begin{gathered} p_{\theta }(x_0)=\int p_{\theta }(x_T)\prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t)d{x}_{1:T} \\ {p}_{\theta }(x_{t-1}|x_t)=\mathcal{N}(x_{t-1};{\mu }_{\theta }(x_t,t),{\sigma }_t\text{I})\text{\hspace{1em}(2)} \end{gathered}$$
$p_{\theta }(x_T)$服从正态分布$\mathcal{N}(x_T;0,\mathrm{I})$, 其与$\theta$无关，可记作$p(x_T)$

当给定条件$c$时，采样过程可以表述为
$$\begin{gathered} p_{\theta }(x_0|c)=\int p(x_T)\prod _{t=1}^T{p}_{\theta }(x_{t-1}|x_t,c)d{x}_{1:T} \\ \text{\hspace{1em}(3)} \end{gathered}$$

由于涉及到积分，直接最大化$p_{\theta }(x_0|c)$很难求解，因此diffusion model的训练采用了最小化对数似然的证据下界(Evidence Lower Bound, ELBO)。通过推导得出：最大化$p_{\theta }(x_0|c)$相当于优化下界【预测噪声和实际添加噪声差异的期望越小越好】。详细过程可参考文献[1]中的式32-45， 86-92
$$\begin{array}{rl}\log p_{\theta }(x_0|c)& \ge {\mathbb{E}}_q[\log \frac{p_{\theta }(x_{0:T}|c)}{q(x_{1:T}|x_0)}]\\ & ={\mathbb{E}}_q[\log \frac{p(x_T){\prod }_{t=1}^T{p}_{\theta }(x_{t-1}|x_t,c)}{{\prod }_{t=1}^{T}q(x_t|x_{t-1})}]\\ & ={\mathbb{E}}_q[\log \frac{p(x_T)p_{\theta }(x_0|x_1,c){\prod }_{t=2}^T{p}_{\theta }(x_{t-1}|x_t,c)}{q(x_T|x_{T-1}){\prod }_{t=1}^{T-1}q(x_t|x_{t-1})}]\\ & \cdots \\ & =-{\mathbb{E}}_{ϵ}[\sum _{t=2}^T\underset{\text{当训练时以均匀分布采样时间步时}{w}_t=1}{\underbrace{w_t}}\Vert ϵ-{ϵ}_{\theta }(x_t,c){\Vert }^2-\underset{\text{当T足够大时，该项}\to 0}{\underbrace{\log p_{\theta }(x_0|x_1,c)}}]+\underset{\text{常数与c无关}}{\underbrace{C}}\\ & \stackrel{去除无关项}{\approx }-{\mathbb{E}}_{ϵ,t}\left[\Vert ϵ-{ϵ}_{\theta }(x_t,c){\Vert }^2\right]\end{array}\text{\hspace{1em}(4)}$$

如何将diffusion model应用到zero-shot classification

对于一个分类模型，给定输入$x$,模型输出类别的概率向量$c$, 即$p_{\theta }(c|x)$，为了用diffusion model求解$p_{\theta }(c|x)$，需要用到贝叶斯公式
$$p_{\theta }(c_i|x)=\frac{p(c_i)p_{\theta }(x|c_i)}{{\sum }_{j}p(c_j)p_{\theta }(x|c_j)}\text{\hspace{1em}(5)}$$
不妨假设各个类别的先验概率相同，有$p(c_1)=p(c_2)=\cdots =p(c_N)=\frac{1}{N}$

式5可写作
p θ ( c i ∣ x ) = p θ ( x ∣ c i ) ∑ j   p θ ( x ∣ c j ) = exp ⁡ { log ⁡ ( p θ ( x ∣ c i ) ) } ∑ j   exp ⁡ { log ⁡ ( p θ ( x ∣ c j ) ) } (6) \begin{aligned} p _ { \theta } ( c _ { i } \vert x ) &= \frac { p _ { \theta } ( x \vert c _ { i } ) } { \sum _ { j } \, p _ { \theta } ( x \vert c _ { j } ) } \\ & = \frac { \exp \{ \log {( p _ { \theta } ( x \vert c _ { i } ) } ) \} } { \sum _ { j } \, \exp \{ \log {( p _ { \theta } ( x \vert c _ { j } ))} \} } \end{aligned} \tag{6} pθ​(ci​∣x)​=∑j​pθ​(x∣cj​)pθ​(x∣ci​)​=∑j​exp{log(pθ​(x∣cj​))}exp{log(pθ​(x∣ci​))}​​(6)
根据式4，我们知道 $\log p_{\theta }(x_0|c)\approx -{\mathbb{E}}_{ϵ,t}\left[\Vert ϵ-{ϵ}_{\theta }(x_t,c){\Vert }^2\right]$， 带入上式得
p θ ( c i ∣ x ) ≈ exp ⁡ { − E ϵ , t [ ∥ ϵ − ϵ θ ( x t , c i ) ∥ 2 ] } ∑ j   exp ⁡ { − E ϵ , t [ ∥ ϵ − ϵ θ ( x t , c j ) ∥ 2 ] } (7) \begin{aligned} p _ { \theta } ( c _ { i } \vert x ) & \approx \frac { \exp \{ - \mathbb{E} _ { \epsilon, t} \left[ \Vert \epsilon - \epsilon _ { \theta } ( x _ { t } , c_i ) \Vert ^ { 2 } \right]\} } { \sum _ { j } \, \exp \{ - \mathbb{E} _ { \epsilon, t} \left[ \Vert \epsilon - \epsilon _ { \theta } ( x _ { t } , c_j ) \Vert ^ { 2 } \right] \} } \end{aligned} \tag{7} pθ​(ci​∣x)​≈∑j​exp{−Eϵ,t​[∥ϵ−ϵθ​(xt​,cj​)∥2]}exp{−Eϵ,t​[∥ϵ−ϵθ​(xt​,ci​)∥2]}​​(7)
由此我们推导出了基于diffusion model的classifier。

如何求解

我们看到求式7的关键是不同类别下，预测的噪声和实际噪声差异的期望。这里面有两个随机变量，分别是$ϵ,t$，其中$ϵ\sim \mathcal{N}(0,\mathrm{I})$, $t\sim \mathrm{uniform}(0,T)$。可以用蒙特卡诺采样对上述期望进行估计，假定依概率对上述两个随机变量采样$K$次，得到$\{({ϵ}_i,t_i)|i=0,1,\cdots K\}$，可将式7转化为
$${\mathbb{E}}_{ϵ,t}\left[\Vert ϵ-{ϵ}_{\theta }(x_t,c_j){\Vert }^2\right]=\frac{1}{K}\sum _{i=1}^K\Vert {ϵ}_i-{ϵ}_{\theta }(x_{t_i},c_j){\Vert }^2\text{\hspace{1em}(8)}$$
当我们求出每一个类别$j$下的$p_{\theta }(c_j|x)$,值最大的就是预测出来的类别。

细心的同学发现了，为了准确的估计期望，需要用蒙特卡诺方法采样较多的样本，一个样本意味着需要用diffusion model推理一次得到预测的噪声，当样本量较大时，推理时间会非常大。总的推理次数为$K\ast \#c$, $K$为蒙特卡诺的采样的样本数目，$\#c$为类别数目, $\#c=N$。

在实践中为了减少推理速度，作者修改了对$t_i,{ϵ}_i$这两个随机变量的采样逻辑，也将上面的one-stage分类的范式转化为two-stage。感兴趣的读者可以阅读原文。本文简单介绍核心思路：

对于第一个提速方案：修改采样逻辑。主要基于实验观测，如下图。

（注：横轴表示${ϵ}_i$的采样数目。Uniform: $t=[0,1,2,..,1000]$, 0, 500, 1000:$t=[0,500,1000]$, even 10: $t=[0,10,20,...,1000]$）

对于第二个提速方案：作者将$N=\{({ϵ}_i,t_i)|i=0,1,\cdots K\}$划分了成两个集合$K_1,K_2$, 首先在$K_1$中根据式7估计$x$的类别。保留概率最高的$M$个类别。随后在集合$K_2$上对前$M$个可能的类别继续用式7计算概率。此时的计算量从$K\ast \#c$变为$K_1\ast \#c+(K_2)\ast M$。

实验

作者对比同为zero-shot classifier的CLIP，结果如下。zero-shot的能力以及接近了基于renset50的CLIP。但与openCLIP ViT-H/14还有较大差距。其它更多的实验对比请见原始论文。

参考文献

[1]: Luo, Calvin. “Understanding diffusion models: A unified perspective.” arXiv preprint arXiv:2208.11970 (2022).
[2]: Li, Alexander C., et al. “Your diffusion model is secretly a zero-shot classifier.” arXiv preprint arXiv:2303.16203 (2023).