ICLR2022 | IBOT:使用online tokenier进行mask预测的自监督框架
-
ICLR2022 | IBOT:使用online tokenier进行mask预测的自监督框架
论文标题:IBOT : IMAGE BERT PRE-TRAINING WITH ONLINE TOKENIZER
论文地址:https://arxiv.org/pdf/2111.07832.pdf
论文代码:https://github.com/bytedance/ibot
摘要
- 引出主题: 语言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精度,并在密集的下游任务(例如,对象检测、实例分割和语义分割)上取得领先的结果。
算法
整体模型架构如上图所示:
-
首先对输入图像进行两种数据增强得到视图uuu和视图vvv
-
分别对uuu和vvv进行mask操作得到他们的masked view u^\hat{\boldsymbol{u}}u^和v^\hat{\boldsymbol{v}}v^
-
然后对u^\hat{\boldsymbol{u}}u^和v^\hat{\boldsymbol{v}}v^进行patch embedding操作
-
online tokenizer不进行mask操作,相应的输入patch为
-
mask patch输入到学生分支,unmask patch输入到教师/online分支,另外,教师分支不进行梯度反传,通过EMA根据学生分支中的网络权重进行动态更新
-
对输出结果进行LMIML_{MIM}LMIM和LCLS\mathcal{L}_{\mathrm{CLS}}LCLS损失计算
其中,LMIM\mathcal{L}_{\mathrm{MIM}}LMIM 是比较教师分支输出和学生分支输出中的unmask token部分,公式如下:
为了确定online tokenizer是具有语义意义的,对CLS token进行自蒸馏计算。论文借鉴DINO论文中的算法,使用如下公式进行LCLS\mathcal{L}_{\mathrm{CLS}}LCLS损失计算
网络结构:
- 使用具有不同数量的参数的ViT(Dosovitskiy et al.,2021)和Swin Transformer(Liu et al.,2021b):ViT-S/16、ViT-B/16、ViT-L/16和Swin-T/{7,14}作为主干网络fff。对于ViT,/16表示面片大小为16。对于Swins,/{7,14}表示窗口大小为7或14。
- 使用大小为224的图像对Transformer进行预训练和微调,patch总数是196。
- 投影头hhh是一个三层MLP,并使用l2l_2l2正则化的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 / 256lr=5e−4×batchsize/256 。
实验
ImageNet-1K上的分类任务
下游任务
COCO数据集中的目标检测和实例分割任务以及ADE20K数据集中的语义分割任务
迁移任务
可视化
消融实验