【论文精读】IBOT

摘要

       掩码语言建模(MLM)是一种流行的语言模型预训练范式，在nlp领域取得了巨大的成功。然而，它对视觉Transformer (ViT)的潜力尚未得到充分开发。为在视觉领域延续MLM的成功，故而探索掩码图像建模(MIM)，以训练更好的视觉transformer，使其可以像NLP一样作为标准组件。

       MLM最关键的问题是语言标记器（lingual tokenizer），其功能是将语言分成语义上有意义的标记。类似，MIM的关键也在于视觉标记器（vision tokenizer）的设计，以及实际要训练的目标网络（target network）。但是，由于图像的连续特性，视觉语义的提取非常困难，故而以往Beit等算法都是分步先计算tokenizer，在根据tokenizer训练target network。然而，由于获取视觉语义是tokenizer和target network的共同目的，理论上两者可联合优化。

       本文提出iBOT，即用online tokenizer进行图像BERT预训练。通过将MIM制定为知识蒸馏(KD)来激励iBOT，使target network学会从tokenizer中提取知识（知识蒸馏）。具体，目标网络用masked图像作为输入，而online tokenizer用原始图像作为输入。 目标是让target network将每个masked patch token恢复到其相应的tokenizer输出（上图）。归纳为：

• tokenizer通过对[CLS] token执行跨视图图像的相似性来逐步学习高级视觉语义
• 所提出的tokenizer不需要额外预处理，其通过EMA与MIM联合优化

框架

掩码图像建模(MIM)

       给定图像token patches x = { x i } i = 1 N x=\{{x_i}\}^N_{i=1} x={xi​}i=1N​，MIM首先根据预测比率$r$采样一组随机mask m ∈ { 0 , 1 } N m \in \{ {0,1}\}^N m∈{0,1}N，其中$N$是token的数量。对于patch token $x_i$、若对应$m_i=1$，则masked patch token表示为 x ˜ ≜ { x i ∣ m i = 1 } \~{x} \triangleq\{ { x_i |m_i=1} \} x˜≜{xi​∣mi​=1} ，对应位置用mask token $e_{[MASK]}$替换，则生成损坏的图像 x ^ ≜ { x ^ i ∣ ( 1 − m i ) x i + m i e [ M A S K ] } i = 1 N \hat x \triangleq\{ {\hat x_i |(1-m_i)x_i+m_ie_{[MASK]}} \}^N_{i=1} x^≜{x^i​∣(1−mi​)xi​+mi​e[MASK]​}i=1N​。此时，MIM的目标是从损坏的图像$\hat{x}$中恢复masked patch token $\tilde{x}$, 故目标为最大化：
$$log{q}_{\theta }(\tilde{x}|\hat{x})\approx \sum _{i=1}^N{m}_i\cdot log{q}_{\theta }(x_i|\hat{x})$$

       其中$\approx$具有独立性假设，即每个mask token 可以单独重建。

       在BEiT中，$q_{\theta }$被建模为分类分布，任务是最小化：
$$-\sum _{i=1}^N{m}_i\cdot P_{\varphi }(x_i{)}^T\log P_{\theta }({\hat{x}}_i)$$

       其中$P(\cdot )$将输入转换为$K$维概率分布，$\varphi$是离散VAE的参数，该参数将图像patch聚类为$K$个类别，并为每个patch token分配一个识别其类别的one-hot编码。这种损失的表述类似于知识蒸馏，其中知识从由$\varphi$参数化的前缀 tokenizer蒸馏到由$\theta$参数化的模型。

自蒸馏

       类似于DINO，给定训练集$I$，均匀采样图像$x\sim I$。对$x$应用两个随机增强，产生两个扭曲视图$u$和$v$放入一个教师-学生框架中，以从[CLS] token：$v_t^{[CLS]}=P_{{\theta }^{\prime }}^{[CLS]}(v)$、$u_t^{[CLS]}=P_{\theta }^{[CLS]}(u)$中获得预测类别分布。知识通过最小化交叉熵从教师向学生传递，表示为：
$$L_{[CLS]}=-P_{{\theta }^{\prime }}^{[CLS]}(v{)}^T\log P_{\theta }^{[CLS]}(u)$$

       其中教师和学生共享一个主干网络$f$和一个投影头$h^{[CLS]}$ 架构。教师网络${\theta }^{\prime }$通过学生网络$\theta$的参数指数移动平均(EMA)学习。

iBOT

       整体框架如图3。结合MIM和自蒸馏方法，框架包含patch tokens和CLS两部分。

       patch tokens部分在两个增强视图$u$和$v$执行blockwise masking（同BEiT）得到masked views $\hat{u}$和$\hat{v}$。输入教师-学生网络后，学生网络输出其patch tokens的masked view $\hat{u}$的投影${\hat{u}}_s^{patch}=P_{\theta }^{patch}(\hat{u})$，教师网络输出其patch tokens的non-masked view $u$的投影$u_t^{patch}=P_{{\theta }^{\prime }}^{patch}(u)$。在此将MIM在iBOT中的训练目标定义为：
$$L_{MIM}=-\sum _{i=1}^N{m}_i\cdot P_{{\theta }^{\prime }}^{patch}(u_i{)}^T\log P_{\theta }^{patch}({\hat{u}}_i)$$

       同样的方法可得到${\hat{v}}_s^{patch}$和$v_t^{patch}$的损失，这两项求和平均后可得到最终的$L_{MIM}$。故教师骨干网络和投影头$h_t^{patch}\cdot f_t$是一个visual tokenizer，其为每个学生网络的masked patch token生成online token分布。iBOT中使用的tokenizer可以联合到MIM目标，故其的特征是可以从当前数据集中提取领域知识，而不是到指定的数据集。

       CLS部分，为了确保online tokenizer具有语义意义，故需对跨视图图像的[CLS] token进行自蒸馏学习，以便其获得视觉语义。在实践中，采用DINO的$L_{[CLS]}$ 训练，且使用${\hat{u}}_s^{[CLS]}$、${\hat{v}}_s^{[CLS]}$为学生网络的输入。为了进一步采用从[CLS]上自蒸馏获得的语义抽象能力，可共享CLS和patch tokens的投影头参数，令$h_s^{[CLS]}=h_s^{patch}$、$h_t^{[CLS]}=h_t^{patch}$，其比使用单独的头产生更好的结果。最后的监督信号使用softmax之后的token分布。

算法伪代码

训练配置

       骨干网络使用具有不同数量参数的Vision transformer和Swin transformer, ViT-S/16、ViT-B/16、ViT-L/16和Swin-t/{7,14}。 对于ViT/16表示patch大小为16。对于Swin-s/{7,14}表示窗口大小为7或14。使用224大小的图像对transformer进行预训练和微调，因此patch token的总数量为196。投影头$h$是3层mlp，在DINO之后有$l_2-normalized$。模型共享CLS和patch tokens的投影头参数，共享头的输出维度设置为8192。

       使用AdamW优化器和1024的batchsize在ImageNet-1K上预训练iBOT，ViT-S/16预训练iBOT为800 epoch, ViT-B/16为400 epoch, ViT-L/16为250 epoch, Swin-T/{7,14}为300 epoch。在ImageNet-22K上预训练则，ViT-B/16用80 epoch, ViT-L/16用50 epoch。学习率在第一个10 epoch期间线性上升到其基值，并随总批处理大小缩放：$lr=5e^{-4}\times batchsize/256$。使用随机MIM，0.5的概率将预测率$r$设置为0，另外0.5的概率并从范围[0.1, 0.5]中均匀采样。总loss将$L_{[CLS]}$和$L_{MIM}$相加，不进行缩放。

实验

定量对比

       上图显示iBOT与其他sota无监督算法在ImageNet上的线性检测精度，相比与ResNet类算法，iBOT可在同等规模参数下得到更高的精度。

       上图为KNN和线性检测结合iBOT在ImageNet预训练上的测试结果。观察到ViT-S/16达到了77.9％的线性检测精度；ViT-B/16的线性检测精度为79.5％；ViT-L/16的线性检测精度为81.0％，KNN精度为78.0％，实现了最先进的性能。使用Swin-T/{7,14}，iBOT分别实现了78.6％和79.3％的线性检测精度。

       使用ViT-L/16和ImageNet-22K作为预训练数据，iBOT进一步实现了线性检测精度82.3％，超过EsViT的Swin-B/14达到的81.3％；ViT-B/16的线性检测精度为79.5％，与SimCLRv2的79.8％相当，但参数量少10%。

       随着参数的增加，与DINO相比的性能提升更大(0.9％w/ ViT-S vs 1.3％w/ ViT-B)，这表明iBOT对更大的模型更具可扩展性。

       上图为研究在ImageNet-1K上的微调并重点比较transformer的自监督方法及其监督基准线。

       如表2：iBOT分别在ViT-S/16, ViT-B/16和ViT-L/16上取得了82.3％，84.0％和84.8％的top-1准确率。

       如表3：iBOT用ImageNet-22K预训练的ViT-B/16和ViT-L/16分别实现了84.4％和86.6％的top-1精度，比ImageNet-22K预训练的BEiT分别高出0.7％和0.6％。当在512的图像大小上进行微调时，达到了87.8％的准确率。使用ViT-L/16，iBOT在使用1K数据时比BEiT差0.4％，但在使用22K数据时好0.6％，意味着iBOT在更多的数据来训练更大的模型可达到更好的性能。

       上图比较了遵循无监督预训练、有监督微调范式的各种方法。iBOT在使用ImageNet-1K1％和10％的数据时，比DINO分别提高了1.6％ 和0.8％，表明具有更高的标签效率。

       上图对无监督学习测试，使用标准的评估指标，包括精度(ACC)、调整随机指数(ARI)、归一化互信息(NMI)和Fowlkes-Mallows指数(FMI)。将 iBOT与SimCLRv2、Self-label、InfoMin和SCAN进行了比较。实现了32.8％ 的NMI，比之前的最先进水平高出1.8％，表明MIM有助于模型在全局范围内学习更强的视觉语义。

       上图在COCO上测试目标检测与实例分割及在ADE20K的测试语义分割。目标检测和实例分割需要同时进行目标定位和分类，级联Mask R-CNN可以在COCO数据集上同时产生边界框和实例掩码。

       左图对比有监督Swin-T 及其自监督对应对象MoBY的结果，该结果与有监督ViT-S/16类似的结果。iBOT则将ViT-S的$A{P}^b$从46.2提高到49.4，$A{P}^m$从40.1提高到42.6，超过了有监督的swit、ViT-S/16和自监督MoBY。图右对比ViT-B/16，iBOT有51.2 $A{P}^b$，44.2 $A{P}^m$，大大超过了以前的最好成绩。

       语义分割是像素级别的分类问题，主要在ADE20K数据集上测试。 可以看到iBOT在ViT-S/16上以0.9mIoU的差距提高了其监督基线，且超过了Swin-T。通过ViT-B/16，iBOT用UperNet将之前的最佳方法DINO提高了3.2mIoU。 注意到使用线性头的BEiT的性能下降，这表明BEiT的特征缺乏局部语义。与具有线性头的监督基线在mIoU上相比，iBOT强局部语义的特性产生了2.9mIoU的增益。

       上图测试了迁移学习。iBOT在ImageNet-1K上进行预训练，在几个较小的数据集上进行微调。观察到虽然在几个数据集(CIFAR10、CIFAR100、Flowers和Cars)上的结果几乎停滞不前，但iBOT与其他SSL框架相比始终表现良好，实现了最先进的迁移结果。

       在更大的数据集(如iNaturalist18和iNaturalist19)中，观察到比DINO更大的性能提升。对于更大的模型，比DINO获得更大的性能增益(在iNaturalist18上ViT/S-16为1.7％，而在iNaturalist18上ViT/B-16为2.0％，在iNaturalist19上ViT/S-16为0.3％，而在iNaturalist19上ViT/B-16为1.0％)。

模式布局

       上图4可视化在ImageNet-1K验证集上自蒸馏的投影头置信度最高的patch token，模型采用800 epoch预训练的ViT-S/16，并为每个$16\times 16$ patch可视化5倍上下文(橙色)。可以观察到高级语义和低级细节的出现，左1、2观察到高级语义大灯和狗耳，3、4突出了低级纹理。

       相对比，图16从BEiT和DINO的patch提取的布局中，虽然可以看到更复杂的纹理，但大多数补丁共享类似的局部细节，而不是高级语义。

自注意力判别

       用ViT-S/16可视化自注意力图，选择[CLS] token作为查询，并使用不同颜色可视化最后一层不同头部的注意力图。观察到iBOT分离不同对象或一个对象的不同部分的能力。在最左边的图中， iBOT明显地区别了鸟和树枝。iBOT注意力主要集中在物体的区分部分(汽车的轮子，鸟的喙)。这些属性对于iBOT擅长图像识别至关重要，特别是在有物体遮挡或分散注意力的情况下。

       上图显示iBOT最后一层的多个头部的自注意力图可视化。与DINO相比，iBOT提供更细致的可视化结果，在视觉上显示出更强的分离不同物体或一个物体的不同部分的能力。例如，在第五列中，iBOT中有一个注意力头，只负责狐狸的耳朵；在第八列中，iBOT将蘑菇分成语义上更有意义的部分。

       Correspondence between two views of one image是从一幅图像的两个视图中采样的图像对，在比例和颜色上增强。观察到从iBOT中提取的图像间语义匹配大部分正确。

       Correspondence between two images of one class是从一个类别的两幅图像中采样的图像对。第一行具有显著物体但大小、位置和纹理不同的图像。第二行来自动物的图像，可以更清楚地观察到iBOT正确匹配了动物的语义部分(狐狸的尾巴，鸟的喙)。第三行以人体或服装为中心的图像。第四行是自然或家庭场景，其中显著的物体是不可见的。虽然没有明确的语义部分可以匹配到人类的理解，但仍然可以观察到iBOT可以根据它们的纹理或颜色，招牌和盒子的木质纹理)提取语义对应。这些可视化结果表明iBOT在局部尺度上具有很强的局部检索和匹配能力。

鲁棒性

       上图从3方面对鲁棒性进行了定量测试：背景变化、遮挡和out-of-distribution的例子，用ViT-S/16在800 epoch进行了预训练，然后对100 epoch进行了线性评估。

       对于背景变化，在ImageNet-9 (IN-9)数据集上研究7种变化类型下的图像。IN-9通过混合不同图像的前景和背景，包括9种粗粒度类和7种变体。Only-FG (O.F.)、Mixed-Same (M.S.)、Mixed-Rand (M.R.)和Mixed-Next (M.N.)原始前景存在，但背景被修改；而No-FG (N.F.)、Only-BG-B (O.BB.)和Only-BG-T (O.BT.)前景被掩盖。观察到除了O.BT以外的性能相比DINO全部得到提升。表明iBOT对背景变化具有鲁棒性。

       对于遮挡，研究了信息损失率为0.5下显著性和非显著性斑块下降的线性精度，iBOT在这两种设置下的性能下降较小。

       对于out-of-distribution的例子，研究了ImageNet-A 中的自然对抗样本和ImageNet-C 中的图像损坏，iBOTImageNet-A上有更高的准确性，在ImageNet-C上有更小的平均损坏误差(mCE)。

       上图实验提供了在三种丢弃设置下(随机、显著和非显著)不同信息损失率的图像遮挡鲁棒性研究结果。展示了iBOT端到端微调或在预训练骨干上用线性头的结果。包括了ViT-S/16和ResNet-50的监督结果以进行比较。

       与CNN对应的ResNet-50相比，ViT显示出更高的鲁棒性，因为transformer的动态感受野使其较少依赖于图像的空间结构。经验发现iBOT与其监督基线相比，对遮挡具有更强的鲁棒性，MIM有助于用自注意力对图像块序列之间的相互作用进行建模，使得丢弃部分元素不会显著降低性能。

       上图实验对输入图像patch块打乱来研究模型对空间结构的敏感性。展示了iBOT端到端微调或在预训练骨干上用线性头的结果，包括了ViT-S/16和ResNet-50的监督结果以进行比较。（打乱网格大小1表示没有打乱，打乱网格大小196表示所有patch token都打乱了）

       观察到iBOT比其监督基线和ResNet-50保持更好的准确性，表明iBOT更少地依赖于位置嵌入来保持正确的分类决策的全局图像上下文。

消融实验

       上图消融研究了使用300 epoch预训练的ViT-S/16的语义有意义的tokenizer的重要性，该tokenizer的预测比率为$r=0.3$，且不进行 Multi-Crop增强。

       iBOT通过跨视图图像$L_{[CLS]}$对[CLS] token进行自蒸馏，以获得视觉语义。为了验证该设置的有效性，在没有$L_{[CLS]}$的情况下执行MIM，或使用其他模型作为视觉tokenizer。$◦$表示单独使用DINO，$△$表示预训练的DALL-E编码器。

       观察到，在没有$L_{[CLS]}$的情况下执行MIM会导致9.5％KNN精度和29.8％的线性精度的不良结果，这表明仅使用MIM很难获得视觉语义。虽然DINO作为tokenizer会出现视觉语义，但其没有达到一个不错的结果(KNN的准确率为44.3％vs69.1％)。 将iBOT与DINO和BEiT (DINO+BEiT)的多任务处理进行比较，看到将自蒸馏获得的视觉语义与tokenizer合并的优势，iBOT在线性探测方面有11.5％的提高，在微调方面有0.3％的提高。根据经验观察到使用[CLS] token和patch token的共享投影头(SH)可实现性能改进，因为其将[CLS] token中获得的语义共享给MIM。

reference

Zhou, J. , Wei, C. , Wang, H. , Shen, W. , Xie, C. , & Yuille, A. , et al. (2021). Ibot: image bert pre-training with online tokenizer. arXiv e-prints.