利用扩散模型的反馈优化图片理解模型
自从 DDPM 以来,扩散模型已经成为了绝对主流的图片生成模型。最近,有研究者开始探索如何将扩散模型的生成能力用于提升判别模型的能力。主要有三种用法:生成模型反演用作判别模型、生成模型用于数据增强、生成模型直接用作特征提取器。
本文介绍几篇利用扩散模型的梯度反馈来优化图片理解模型(判别模型)的方法。这些方法应该划归到上述的第一类用法:基于反演(Inversion-based)的方法。给定图片 x x x 和条件生成模型 p ϕ ( x ∣ c ) p_\phi(x|\mathbf{c}) pϕ(x∣c),通过最大化生成模型的似然,来反演得到条件嵌入 c \mathbf{c} c。
将判别模型的输出作为文生图模型的条件输入,利用文生图模型的梯度反馈来更新判别模型,从而改善判别模型的表征质量。有点 GAN 反着用的感觉。GAN 是判别模型倒逼生成模型生成更加真实的样本,而 DIffTTA 是生成模型倒逼判别模型生成质量更高的表征。
Diffusion TTA
Diffusion TTA 是将判别模型(图片理解模型)的输出,经过某种变换,作为文生图模型的文本条件,然后对原图加噪声,和文本条件一起送入到文生图模型中,计算损失,更新判别模型的参数。
作者讨论了三种判别模型以及对应的转换为文本条件的方式。
- 图片分类模型:分类模型接受一个图片输入,并输出 L L L 个类别的概率分布 y ∈ [ 0 , 1 ] L , y T 1 L = 1 y\in[0,1]^L,y^T\mathbf{1}_L=1 y∈[0,1]L,yT1L=1。对于分类模型,构造出的文本条件为为各个类别的文本嵌入 l j ∈ R d , j ∈ [ 1 , L ] l_j\in\mathbb{R}^d,j\in[1,L] lj∈Rd,j∈[1,L] 按照输出概率分布的加权和 c = ∑ j = 1 L y j ⋅ l j \mathbf{c}=\sum_{j=1}^Ly_j\cdot l_j c=∑j=1Lyj⋅lj。
- 语义分割模型:语义分割模型相当于是对输入图片的每个像素都进行分类,输出为 y = { y u ∈ [ 0 , 1 ] L , y u T 1 L = 1 } y=\{y^u\in[0,1]^L,{y^u}^T\mathbf{1}_L=1\} y={yu∈[0,1]L,yuT1L=1} 其中 u ∈ x u\in x u∈x 是图像中的各个像素位置。对于各分割模型,构造出的文本条件为 c = { ∑ j = 1 L y k u ⋅ l j , u ∈ x } \mathbf{c}=\{\sum_{j=1}^Ly_k^u\cdot l_j,u\in x\} c={∑j=1Lyku⋅lj,u∈x}。
- 深度预测模型:深度预测模型的输出是与图像相等大小的深度图 y ∈ R + w × h y\in{\mathbb{R}^{+}}^{w\times h} y∈R+w×h,构造的文本条件就是其本身 c = y \mathbf{c}=y c=y。
由于各种判别模型构造出的文本条件
c
\mathbf{c}
c 都关于其权重
θ
\theta
θ 是可微的,我们就可以根据扩散模型的损失来更新参数
θ
\theta
θ:
L
(
θ
,
ϕ
)
=
E
t
,
ϵ
∣
∣
ϵ
ϕ
(
α
ˉ
t
x
+
1
−
α
ˉ
ϵ
,
t
,
c
)
−
ϵ
∣
∣
2
2
L(\theta,\phi)=\mathbb{E}_{t,\epsilon}||\epsilon_\phi(\sqrt{\bar\alpha}_tx+\sqrt{1-\bar\alpha}\epsilon,t,\mathbf{c})-\epsilon||_2^2
L(θ,ϕ)=Et,ϵ∣∣ϵϕ(αˉtx+1−αˉϵ,t,c)−ϵ∣∣22
Diffusion TTA 具体的执行过程如算法 1 所示。
- 给定输入图片 x x x,判别模型参数 θ \theta θ,扩散生成模型参数 ϕ \phi ϕ,adaption 步数 N N N,批尺寸 B B B 和学习率 η \eta η;
- 在每一步中:
- 首先将图片输入到判别模型,得到输出 y = f θ ( x ) y=f_\theta(x) y=fθ(x);
- 使用上述设计的转换方法,将判别模型的输出转换为扩散模型的条件输入 c = transform ( y ) \mathbf{c}=\text{transform}(y) c=transform(y)
- 采样时间步 { t i } i = 1 B \{t_i\}_{i=1}^B {ti}i=1B 并计算对应的噪声 { ϵ i } i = 1 B \{\epsilon_i\}_{i=1}^B {ϵi}i=1B
- 按照上式计算损失 L ( θ , ϕ ) L(\theta,\phi) L(θ,ϕ)
- 更新判别模型的参数 θ ← θ − η ∇ θ L ( θ , ϕ ) \theta\leftarrow \theta-\eta\nabla_\theta L(\theta,\phi) θ←θ−η∇θL(θ,ϕ);
- (可选)更新生成模型的参数 ϕ ← ϕ − η ∇ ϕ L ( θ , ϕ ) \phi\leftarrow \phi-\eta\nabla_\phi L(\theta,\phi) ϕ←ϕ−η∇ϕL(θ,ϕ);
- 在迭代完成后:
- 得到最终的判别模型参数 θ \theta θ,计算输出 y = f θ ( x ) y=f_\theta(x) y=fθ(x)
在 Diffusion TTA 的项目主页上,有一个动图来演示这个迭代的过程。可以看到,随着迭代的进行,扩散模型的损失不断下降,生成图片与原图也越来也接近,这一切都是因为判别模型权重在更新,其输出,也就是输入到生成模型的条件,越来越准确。这样,就倒逼判别模型的准确率不断上升。
DIVA
DIVA 是近些天智源研究院提出的一中使用扩散模型反馈来增强 CLIP 能力的方法。众所周知,由于训练目标和训练数据的原因,CLIP 存在一些固有的缺陷,其仅聚焦于图像的高层语义,而对图像细节特征感知较差。有研究(参考)指出,CLIP 很难识别出两张图片的细微差别,而这些差别在人类来看是显而易见的。在实际中,我们也经常发现,CLIP 在朝向、颜色、数量和结构等方面,经常出现错误。
DIVA 其大致思路与 Diffusion TTA 差不多,也是利用扩散模型的梯度反馈来更新判别模型(在 DIVA 这里,就是指 CLIP)。具体来说,DIVA 方法的整体结构如下图所示。主要有两部分组成:待优化的 CLIP 模型和预训练的扩散模型(用于提供梯度反馈)。首先分别用 CLIP 的图像和文本编码器对输入图片和空文本进行编码,然后经过一种转换方法,得到条件输入到扩散模型中,进行去噪生图,并根据扩散模型的损失来更新 CLIP 模型。
这里的关键就是从视觉、文本表示,到扩散模型条件的转换方法的设计。DIVA 的作者认为,CLIP 本身仅关注图片高层语义,很大一部分原因是因为只使用了一个 ViT CLS token 来作为整张图片的表征。因此,作者这里提出了 Visual Dense Recap Scheme,将各个图像块 patch token 也添加到合成扩散模型条件的过程中,提升模型对局部区域特征的表达,从而使得模型可以根据扩散模型的反馈来改善输出的表征。
下面的图表展示了使用 DIVA 改进的 CLIP 在细粒度理解 benchmark 上的定性/定量性能提升。可以看到,对于不同的 CLIP 版本,DIVA 平均有着 3-6 个点的提升,还是非常显著的。在定性展示中,CLIP 原来一些简单视觉错误也都得到了修正。
另外一个比较重要的消融实验是关于将 CLIP embedding 转换为扩散模型条件的方式,实验结果如下图所示。首先可以看到,无论使用什么样的条件构造方式,DIVA 对 CLIP 的细粒度理解能力都有着普遍的提升。其中 3、4 两行是仅使用文本 embedding 作为条件,这就类似于 Diffusion TTA 的做法了。下方三行则是 DIVA 结合视觉、文本嵌入,并使用 Visual Dense Recap Scheme 构造的条件,其中 G 表示 ViT 中全局的 CLS token,L 表示局部 patch token。可以看到,当使用全局 CLS token 和部分(而非全部)patch token 来构造条件嵌入时,提升最为显著,达到 6.6 个点。这说明 Visual Dense Recap Scheme 中视觉局部 patch token 的密度既不能太低也不能太高。密度太低会导致任务过于困难,无法充分利用扩散模型反馈来改善输出的表征,密度太高则导致任务过于简单,机会不需要学习,扩散模型就能重构出图像。
总结
将判别模型的输出作为文生图模型的条件输入,利用文生图模型的梯度反馈来更新判别模型,从而改善判别模型的表征质量。有点 GAN 反着用的感觉。GAN 是判别模型倒逼生成模型生成更加真实的样本,而 DIffTTA 是生成模型倒逼判别模型生成质量更高的表征。但是文生图模型的语义理解能力也是来自 CLIP 的呀,为什么反过来还能提升 CLIP 的表征质量呢?感觉 DIVA 还是靠密集条件的策略?不知道哪里有更多理论分析。