摘要
掩码语言建模(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
r
r采样一组随机mask
m
∈
{
0
,
1
}
N
m \in \{ {0,1}\}^N
m∈{0,1}N,其中
N
N
N是token的数量。对于patch token
x
i
x_i
xi、若对应
m
i
=
1
m_i=1
mi=1,则masked patch token表示为
x
˜
≜
{
x
i
∣
m
i
=
1
}
\~{x} \triangleq\{ { x_i |m_i=1} \}
x˜≜{xi∣mi=1} ,对应位置用mask token
e
[
M
A
S
K
]
e_{[MASK]}
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+mie[MASK]}i=1N。此时,MIM的目标是从损坏的图像
x
^
\hat x
x^中恢复masked patch token
x
˜
\~{x}
x˜, 故目标为最大化:
l
o
g
q
θ
(
x
˜
∣
x
^
)
≈
∑
i
=
1
N
m
i
⋅
l
o
g
q
θ
(
x
i
∣
x
^
)
log \ q_{\theta}(\~{x}|\hat x)≈\sum^N_{i=1}m_i \cdot log \ q_{\theta}(x_i|\hat x)
log qθ(x˜∣x^)≈i=1∑Nmi⋅log qθ(xi∣x^)
其中 ≈ ≈ ≈具有独立性假设,即每个mask token 可以单独重建。
在BEiT中,
q
θ
q_{\theta}
qθ被建模为分类分布,任务是最小化:
−
∑
i
=
1
N
m
i
⋅
P
ϕ
(
x
i
)
T
log
P
θ
(
x
^
i
)
-\sum^N_{i=1}m_i \cdot P_{\phi}(x_i)^T \log P_{\theta}(\hat x_i)
−i=1∑Nmi⋅Pϕ(xi)TlogPθ(x^i)
其中 P ( ⋅ ) P(\cdot) P(⋅)将输入转换为 K K K维概率分布, ϕ \phi ϕ是离散VAE的参数,该参数将图像patch聚类为 K K K个类别,并为每个patch token分配一个识别其类别的one-hot编码。这种损失的表述类似于知识蒸馏,其中知识从由 ϕ \phi ϕ参数化的前缀 tokenizer蒸馏到由 θ \theta θ参数化的模型。
自蒸馏
类似于DINO,给定训练集
I
I
I,均匀采样图像
x
∼
I
x \sim I
x∼I。对
x
x
x应用两个随机增强,产生两个扭曲视图
u
u
u和
v
v
v放入一个教师-学生框架中,以从[CLS] token:
v
t
[
C
L
S
]
=
P
θ
′
[
C
L
S
]
(
v
)
v^{[CLS]}_t=P^{[CLS]}_{θ'} (v)
vt[CLS]=Pθ′[CLS](v)、
u
t
[
C
L
S
]
=
P
θ
[
C
L
S
]
(
u
)
u^{[CLS]}_t=P^{[CLS]}_\theta(u)
ut[CLS]=Pθ[CLS](u)中获得预测类别分布。知识通过最小化交叉熵从教师向学生传递,表示为:
L
[
C
L
S
]
=
−
P
θ
′
[
C
L
S
]
(
v
)
T
log
P
θ
[
C
L
S
]
(
u
)
L_{[CLS]}=-P_{\theta'}^{[CLS]}(v)^T \log P_{\theta}^{[CLS]}(u)
L[CLS]=−Pθ′[CLS](v)TlogPθ[CLS](u)
其中教师和学生共享一个主干网络 f f f和一个投影头 h [ C L S ] h^{[CLS]} h[CLS] 架构。教师网络 θ ′ θ' θ′通过学生网络 θ \theta θ的参数指数移动平均(EMA)学习。
iBOT
整体框架如图3。结合MIM和自蒸馏方法,框架包含patch tokens和CLS两部分。
patch tokens部分在两个增强视图
u
u
u和
v
v
v执行blockwise masking(同BEiT)得到masked views
u
^
\hat u
u^和
v
^
\hat v
v^。输入教师-学生网络后,学生网络输出其patch tokens的masked view
u
^
\hat u
u^的投影
u
^
s
p
a
t
c
h
=
P
θ
p
a
t
c
h
(
u
^
)
\hat u^{patch}_s=P^{patch}_\theta(\hat u)
u^spatch=Pθpatch(u^),教师网络输出其patch tokens的non-masked view
u
u
u的投影
u
t
p
a
t
c
h
=
P
θ
′
p
a
t
c
h
(
u
)
u_t^{patch}=P^{patch}_{θ'}(u)
utpatch=Pθ′patch(u)。在此将MIM在iBOT中的训练目标定义为:
L
M
I
M
=
−
∑
i
=
1
N
m
i
⋅
P
θ
′
p
a
t
c
h
(
u
i
)
T
log
P
θ
p
a
t
c
h
(
u
^
i
)
L_{MIM}=-\sum^N_{i=1}m_i \cdot P_{\theta'}^{patch}(u_i)^T \log P_{\theta}^{patch}(\hat u_i)
LMIM=−i=1∑Nmi⋅Pθ′patch(ui)TlogPθpatch(u^i)
同样的方法可得到 v ^ s p a t c h \hat v^{patch}_s v^spatch和 v t p a t c h v_t^{patch} vtpatch的损失,这两项求和平均后可得到最终的 L M I M L_{MIM} LMIM。故教师骨干网络和投影头 h t p a t c h ⋅ f t h^{patch}_t \cdot f_t htpatch⋅ft是一个visual tokenizer,其为每个学生网络的masked patch token生成online token分布。iBOT中使用的tokenizer可以联合到MIM目标,故其的特征是可以从当前数据集中提取领域知识,而不是到指定的数据集。
CLS部分,为了确保online tokenizer具有语义意义,故需对跨视图图像的[CLS] token进行自蒸馏学习,以便其获得视觉语义。在实践中,采用DINO的 L [ C L S ] L_{[CLS]} L[CLS] 训练,且使用 u ^ s [ C L S ] \hat u_s^{[CLS]} u^s[CLS]、 v ^ s [ C L S ] \hat v_s^{[CLS]} v^s[CLS]为学生网络的输入。为了进一步采用从[CLS]上自蒸馏获得的语义抽象能力,可共享CLS和patch tokens的投影头参数,令 h s [ C L S ] = h s p a t c h h^{[CLS]}_s=h_s^{patch} hs[CLS]=hspatch、 h t [ C L S ] = h t p a t c h h^{[CLS]}_t=h_t^{patch} ht[CLS]=htpatch,其比使用单独的头产生更好的结果。最后的监督信号使用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 h h是3层mlp,在DINO之后有 l 2 − n o r m a l i z e d l_2-normalized l2−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期间线性上升到其基值,并随总批处理大小缩放: l r = 5 e − 4 × b a t c h s i z e / 256 lr=5e^{-4} \times batchsize/256 lr=5e−4×batchsize/256。使用随机MIM,0.5的概率将预测率 r r r设置为0,另外0.5的概率并从范围[0.1, 0.5]中均匀采样。总loss将 L [ C L S ] L_{[CLS]} L[CLS]和 L M I M L_{MIM} LMIM相加,不进行缩放。
实验
定量对比
上图显示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 AP^b APb从46.2提高到49.4, A P m AP^m APm从40.1提高到42.6,超过了有监督的swit、ViT-S/16和自监督MoBY。图右对比ViT-B/16,iBOT有51.2 A P b AP^b APb,44.2 A P m AP^m APm,大大超过了以前的最好成绩。
语义分割是像素级别的分类问题,主要在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
×
16
16 \times 16
16×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
r=0.3
r=0.3,且不进行 Multi-Crop增强。
iBOT通过跨视图图像 L [ C L S ] L_{[CLS]} L[CLS]对[CLS] token进行自蒸馏,以获得视觉语义。为了验证该设置的有效性,在没有 L [ C L S ] L_{[CLS]} L[CLS]的情况下执行MIM,或使用其他模型作为视觉tokenizer。 ◦ ◦ ◦表示单独使用DINO, △ △ △表示预训练的DALL-E编码器。
观察到,在没有 L [ C L S ] L_{[CLS]} 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.