MaskCLIP:通过掩码自蒸馏提升对比语言-图像预训练.
0. TL; DR
本文介绍了一种名为 MaskCLIP 的新型视觉-语言对比学习框架,通过引入掩码自蒸馏(masked self-distillation)来提升对比语言-图像预训练的效果。MaskCLIP 的核心思想是从完整图像中提取的表征来指导从掩码图像中预测的表征,从而学习局部语义信息。实验表明,MaskCLIP 在多种下游任务中表现出色,包括线性探测、微调和零样本学习。该方法在 ImageNet-1K 分类、ADE20K 语义分割和 MS-COCO 目标检测等任务中均取得了显著的性能提升。
1. 背景介绍
近年来,视觉-语言对比学习(VL contrastive learning)在预训练领域取得了显著进展。通过利用大规模的图像-文本对,模型能够通过图像和文本之间的对齐学习强大的语义先验。然而,现有的方法主要关注全局表征的学习,忽略了图像中的局部细节和语义信息。这些细节对于下游任务可能非常关键,但往往没有得到充分利用。
为了解决这一问题,MaskCLIP 提出了一种掩码自蒸馏方法,通过随机掩码输入图像的一部分,强制视觉编码器关注剩余的可见区域,从而学习局部补丁的表征。此外,MaskCLIP 还引入了局部语义监督,使图像和文本编码器能够从语言中获得间接监督,进一步提升模型的表征能力。
2. MaskCLIP 模型
MaskCLIP 的核心是将掩码自蒸馏与视觉-语言对比学习相结合,以提升视觉编码器的表征能力。
MaskCLIP 使用一个基于 Transformer 的文本编码器和一个视觉编码器(如 Vision Transformer, ViT),通过对比学习目标对齐图像和文本的全局表征。给定一个图像-文本对 $I, T$,模型提取视觉特征和文本特征,并通过投影头将它们映射到一个度量空间,计算对比损失:
\(L_I=− \frac{1}{B} \sum_{i=1}^B \log \frac{\exp(e^I_i e^T_i/\sigma)}{\sum_{j=1}^B \exp(e^I_i e^T_j/\sigma)} \\ L_T =− \frac{1}{B} \sum_{i=1}^B \log \frac{\exp(e^T_i e^I_i/\sigma)}{\sum_{j=1}^B \exp(e^T_i e^I_j/\sigma)}\) 其中,$B$ 是训练批次中的图像-文本对数量,$σ$ 是对比损失的温度参数。
MaskCLIP 引入了掩码自蒸馏,通过一个“教师模型”(EMA 模型)和“学生模型”来学习局部语义。教师模型的参数是学生模型参数的指数移动平均(EMA),用于生成目标特征。学生模型则从掩码图像中预测特征,并通过最小化与教师模型输出的交叉熵损失来学习:
\(L_{Dist}=− \frac{1}{|M|} \sum_{k∈M} -\overline{h} (\overline{f}_k)^T\log h(f_k^{′′})\) 其中,$M$ 是掩码的补丁索引,$h$ 和 $\overline{h}$ 分别是学生和教师的在线量化器,用于将特征映射到软编码词分布。
MaskCLIP 还引入了掩码语言建模(MLM)来增强文本编码器。通过随机掩码文本中的部分单词,并使用一个小的解码器预测掩码单词的特征,模型能够学习文本的局部语义:
\[L_{MLM}=− \frac{1}{|M_T|}\sum_{k∈M_T} -t_k^T \log t_k^{′′}\]其中,$M_T$ 是掩码的文本单词索引。
3. 实验分析
图像分类:在 ImageNet-1K 零样本分类任务中,MaskCLIP 达到了 44.5% 的准确率,比 CLIP 提升了 6.9%。在 ImageNet-1K 的线性探测任务中,MaskCLIP 的准确率达到了 73.7%,比 CLIP 提升了 7.2%。在微调任务中,MaskCLIP 的 Top-1 准确率达到了 83.6%,比 CLIP 提升了 1.3%。
语义分割:在 ADE20K 数据集上,MaskCLIP 的 mIoU 达到了 50.5%,比 CLIP 提升了 2.7%。这表明 MaskCLIP 学习到的局部语义表征在密集预测任务中具有优势。
图像-文本检索:在 Flickr30K 和 MS-COCO 数据集上,MaskCLIP 的图像到文本检索和文本到图像检索的 Rank@1 准确率分别达到了 70.1%/41.4% 和 45.6%/25.5%,显著优于 CLIP 和其他对比方法。
通过可视化图像特征和文本特征之间的相似性,MaskCLIP 能够更好地对齐图像和文本的语义信息。例如,在 MS-COCO 数据集上,MaskCLIP 能够更准确地区分图像中的不同对象。
通过消融实验,MaskCLIP 证明了掩码自蒸馏和掩码语言建模的重要性。例如,移除掩码语言建模会导致图像-文本检索任务性能下降,而移除掩码自蒸馏则会导致所有任务性能下降。此外,实验还表明,使用交叉熵损失比均方误差损失更有效,且浅层解码器在图像和文本分支中均表现更好。