ICLR2022 | IBOT：使用online tokenier进行mask预测的自监督框架

189****6672

论文标题：IBOT : IMAGE BERT PRE-TRAINING WITH ONLINE TOKENIZER

摘要

引出主题：语言Transformer的成功主要归功于mask语言建模（MLM）的前置任务（Devlin et al.，2019），在这项任务中，文本首先被tokenized成语义上有意义的片段。
现存问题：该论文研究了mask图像建模（MIM），并指出了视觉tokenizer的优势和挑战。
解决方法：论文提出了一个自监督的框架iBOT，它可以通过online tokenizer执行掩码预测。具体来说，在mask patch tokens上形成自蒸馏，并将教师网络作为在online tokenizer。
实验结果：通过在ImageNet-1K上实现82.3%的linear probing精度和87.8%的fine-tuning精度，并在密集的下游任务（例如，对象检测、实例分割和语义分割）上取得领先的结果。

整体模型架构如上图所示：

首先对输入图像进行两种数据增强得到视图 $u$ 和视图 $v$
分别对 $u$ 和 $v$ 进行mask操作得到他们的masked view $u^\hat{\boldsymbol{u}}$ 和 $v^\hat{\boldsymbol{v}}$
然后对 $u^\hat{\boldsymbol{u}}$ 和 $v^\hat{\boldsymbol{v}}$ 进行patch embedding操作
online tokenizer不进行mask操作，相应的输入patch为
mask patch输入到学生分支，unmask patch输入到教师/online分支，另外，教师分支不进行梯度反传，通过EMA根据学生分支中的网络权重进行动态更新
对输出结果进行 $L_{MIM}$ 和 $LCLS\mathcal{L}_{\mathrm{CLS}}$ 损失计算

其中， $LMIM\mathcal{L}_{\mathrm{MIM}}$ 是比较教师分支输出和学生分支输出中的unmask token部分，公式如下：

为了确定online tokenizer是具有语义意义的，对CLS token进行自蒸馏计算。论文借鉴DINO论文中的算法，使用如下公式进行 $LCLS\mathcal{L}_{\mathrm{CLS}}$ 损失计算

网络结构：

使用具有不同数量的参数的ViT（Dosovitskiy et al.，2021）和Swin Transformer（Liu et al.，2021b）：ViT-S/16、ViT-B/16、ViT-L/16和Swin-T/{7,14}作为主干网络 $f$ 。对于ViT，/16表示面片大小为16。对于Swins，/{7,14}表示窗口大小为7或14。
使用大小为224的图像对Transformer进行预训练和微调，patch总数是196。
投影头 $h$ 是一个三层MLP，并使用 $l_2$ 正则化的bottleneck，与在DINO（Caron等人，2021年）一致

预训练设置：

使用AdamW优化器和1024个batchsize在ImageNet-1K训练集上预先培训iBOT。
使用ViT-S/16预训练iBOT800个epochs，用ViT-B/16预训练400个epochs，用ViT-L/16预训练250个epochs，用Swin-T/{7,14}预训练iBOT300个epochs。还使用了ViT-B/16对ImageNet-22K训练集进行80个epochs的预训练，并使用ViT-L/16进行了50个epochs预训练。
在前10个epochs，学习率线性上升至其基准值，基准值与batchsize成比例： $lr=5e−4×batchsize/256\mathrm{lr}=5 e^{-4} \times batchsize / 256$ 。