字节&约翰斯·霍普金斯&上交提出iBOT框架，基于MIM进行自监督训练，在ImageNet-1K上达到86.3%的微调精度！

CV君

发布于 2021-12-13 14:26:59

6870

发布于 2021-12-13 14:26:59

▊ 写在前面

语言Transformer的成功主要归功于masked language modeling（MLM） 的预训练任务，其中文本首先被标记为语义上有意义的片段。在这项工作中，作者研究了masked image modeling（MIM） ，并指出了使用语义上有意义的视觉标记器（visual tokenizer） 的优势和挑战。作者提出了一个自监督的框架iBOT ，它可以通过在线标记器（online tokenizer） 执行mask预测。

具体而言，作者对masked patch tokens进行自蒸馏，并将教师网络作为在线标记器，同时对class token进行自蒸馏以获得视觉语义。在线标记器可与MIM目标共同学习，无需多阶段训练pipeline，其中tokenizer需要事先进行预训练。

作者通过在ImageNet-1K上实现81.6%的linear probing精度和86.3%的微调精度，展示了iBOT的有效性。除了SOTA的图像分类结果外，作者还强调了局部语义模式，这有助于模型获得针对常见损坏的强大鲁棒性，并在密集的下游任务（例如，目标检测、实例分割和语义分割）上取得SOTA的结果。

▊ 1. 论文和代码地址

iBOT: Image BERT Pre-Training with Online Tokenizer

论文地址：https://arxiv.org/abs/2111.07832

代码地址：未开源

▊ 2. Motivation

Masked Language Modeling（MLM） 是一种流行的语言模型预训练范式，它首先随机mask然后重建一组输入token。MLM预训练的Transformer已证明了其可扩展到大容量模型和数据集，目前成为语言任务的有效方法。

然而，视觉Transformer（ViT）的潜力最近开始彻底改变计算机视觉研究，但在很大程度上还没有得到充分的开发。目前视觉任务中的大多数无监督训练都是基于全局视图，忽略了图像的内部结构，这与MLM的局部建模是不同的 。

在本文中，作者探索了与MLM相似的Masked Image Modeling（MIM）来进行实现更好的视觉Transformer。

MLM中最关键的组成部分之一是语言标记器（lingual tokenizer） ，它将语言拆分为语义上有意义的token，如BERT中的WordPiece。类似地，MIM的关键在于视觉标记器的正确设计，该标记器将mask patch转换为监督目标模型的信号（如上图所示）。

然而，与词频统计分析自然产生的语言语义不同，由于图像的连续性，视觉语义无法如此容易地提取。视觉语义是通过引导在线表示逐渐出现的，它强制执行损坏图像视图的相似性。

该属性直观地表示一个多阶段的训练pipeline，在训练目标模型之前，需要首先训练一个现成的语义丰富的标记器。然而，由于获取视觉语义是标记器和目标模型的共同目的，因此，可以联合优化标记器和目标模型的单阶段训练pipeline有待进一步探索。

为此，作者提出了一个新的框架iBOT（image BERT pre-training with Online Tokenizer） ，通过标记器很好地处理上述挑战来执行MIM。作者通过将MIM表述为知识提取（KD）来激励iBOT，学习从标记器中提取知识，并进一步在作为在线标记器的twin教师的帮助下为MIM执行self-distillation。

目标网络接收一个mask图像，而在线标记器接收到原始图像。目标是让目标网络将每个mask patch token恢复到其相应的tokenizer输出 。

online tokenizer主要解决了两大挑战。一方面，tokenizer通过在class token上强制执行交叉视图图像的相似性来捕获高级视觉语义 。另一方面，本文的tokenizer不需要额外的训练阶段作为预处理设置，因为它通过动量更新与MIM联合优化 。

online tokenizer使iBOT能够在特征表示方面获得优异的性能。具体而言，iBOT将k-NN、linear probing和微ﬁne-tuning设置下的ImageNet-1K分类基准分别提高到77.1%、79.5%和83.8%（使用ViT Base/16），比之前的最佳结果高出1.0%、1.3%和0.2%。

当使用ImageNet-22K进行预训练时，使用ViT-L/16的iBOT可实现81.6%的linear probing精度和86.3%的微调精度，两者均比以前的最佳结果高0.3%。除此之外，当迁移到其他数据集或在半监督和非监督分类设置下时，这种提升也是有效的。

本文提出的方法可以帮助模型在全局和局部尺度上进行图像识别。作者发现，patch token中学习到的语义模式有助于该模型在线性分类和抗常见图像损坏的鲁棒性方面得到提升。

当它被迁移到下游任务时，在与图像分类、目标检测、实例分割和语义分割相关的下游任务中，iBOT超过了以前的方法。实验结果表明，iBOT在很大程度上弥补了masked modeling预训练在语言和视觉Transformer之间的差距。

▊ 3. 方法

在本节中，作者提出了iBOT，将自蒸馏作为token生成的自监督目标，并通过自蒸馏执行MIM。在上图中说明了iBOT的框架，下图中展示了伪代码。

3.1 Framework

首先，作者对两个增强视图u和v进行mask，得到他们的mask视图和。以为例，学生网络输出mask视图的结果，教师网络输出没有mask视图的结果。MIM在iBOT的训练目标为：

\mathcal{L}_{\mathrm{MIM}}=-\sum_{i=1}^{N} m_{i} \cdot P_{\boldsymbol{\theta}^{\prime}}^{\mathrm{patch}}\left(\boldsymbol{u}_{i}\right)^{\mathrm{T}} \log P_{\boldsymbol{\theta}}^{\mathrm{patch}}\left(\hat{\boldsymbol{u}}_{i}\right)

对于，作者也采用来对称的方式来计算和的CE loss。

由主干网络加投影头组成的教师网络是一个visual tokenizer，它为每个mask patch token生成在线token分布。iBOT中使用的tokenizer可以联合学习MIM目标，而无需在额外阶段进行预训练。

为了确保online tokenizer在语义上有意义，作者对交叉视图图像的[CLS] token执行自蒸馏，这样就可以通过引导获得视觉语义，就像大多数自监督方法所实现的那样。

为了进一步利用从[CLS] token上的self-distillatin获得的语义抽象能力，作者共享[CLS] token和patch token的投影头参数，即，。根据实验结果，作者发现它比使用单独的head产生更好的结果。

与语义几乎确定的tokenized words不同，图像patch的语义是模糊的。因此，tokenization作为一种one-hot discretization对于图像来说可能是次优的。在iBOT中，作者使用softmax之后的token分布，而不是一个one-hot id作为监督信号，这在iBOT预训练中起着重要作用。

3.2 Implementation

作者使用具有不同参数量的Vision Transformer和Swin Transformer，ViT-S/16、ViT-B/16、ViT-L/16和SwinT/ {7,14} 作为主干网络f。默认情况下，作者使用AdamW优化器和1024的batch size在ImageNet-1K训练集上预训练iBOT。作者预训练了ViT-S/16 800个epoch，ViT-B/16 400个epoch，Swin-T/ { 7,14 } 300个epoch。

▊ 4.实验