iBOT:使用在线标志进行图像BERT预训练.
iBOT提出了适用于视觉任务的大规模预训练方法,通过对图像使用在线 tokenizer 进行 BERT 式预训练让 CV 模型获得通用广泛的特征表达能力。
对于 CV 任务的 MIM 中,图像 patch 是连续分布的且存在大量冗余的底层细节信息,visual tokenizer 把 image patch 变成富含图像语义的 tokens,从而遮住tokens的一部分并通过模型去预测这些遮住部分的语义信息,
iBOT 认为一个能够提取图像 patch 中高层语义的 tokenizer 可帮助模型避免学习到冗余的这些细节信息。作者认为视觉的 tokenizer 应该具备两个属性:(a)具备完整表征连续图像内容的能力;(b) 具备高层语义。基于此设计了一种在线tokenizer。
作者首先将经过 mask 过的图片序列输入 Transformer 之后进行预测的过程建模为知识蒸馏的过程,从 tokenizer 中获得知识。具体地,待训练的目标网络 (比如 ViT) 输入 masked images,Online tokenizer 接收原始图像。目标是让待训练的目标网络将每个 masked patch token 恢复到其相应的 token。Online tokenizer 与目标网络一起学习,希望能够捕获到图片高维的语义信息。
在线 tokenizer 是指 tokenizer 和目标网络共享网络结构,在线即指 tokenizer 其参数从目标网络的历史参数中滑动平均得出。通过使用在线 tokenizer 监督 MIM 过程,即 tokenizer 和目标网络同步学习,能够较好地保证语义的同时并将图像内容转化为连续的特征分布。
iBOT的具体实现过程如下:
- 对一张图片$x$构造两个数据增强版本$u,v$,对它们进行 random mask 操作,得到 mask 版本$\hat{u},\hat{v}$;
- 学生网络$f_s$ (通过梯度下降更新参数) 输入mask 版本$\hat{u},\hat{v}$,得到token $\hat{u}_s,\hat{v}_s$;
- 教师网络$f_t$ (通过EMA更新参数) 输入unmask 版本$u,v$,得到token $u_t,v_t$;
- 使用$\hat{u}_s,\hat{v}_s, u_t,v_t$构造损失函数。
iBOT的损失包括:
① patch 标签上的自蒸馏。让目标网络的2个输出 tokens 和教师网络的2个输出 tokens 分别越接近越好。
\[\begin{gathered} \mathcal{L}_{\mathrm{MIM}}(\boldsymbol{u},\hat{\boldsymbol{u}})=-\sum_{i=1}^N m_i \cdot P_{\boldsymbol{\theta}^{\prime}}^{\mathrm{patch}}\left(\boldsymbol{u}_i\right)^{\mathrm{T}} \log P_{\boldsymbol{\theta}}^{\text {patch }}\left(\hat{\boldsymbol{u}}_i\right) \\ \mathcal{L}_{\mathrm{MIM}}(\boldsymbol{v},\hat{\boldsymbol{v}})=-\sum_{i=1}^N m_i \cdot P_{\boldsymbol{\theta}^{\prime}}^{\mathrm{patch}}\left(\boldsymbol{v}_i\right)^{\mathrm{T}} \log P_{\boldsymbol{\theta}}^{\text {patch }}\left(\hat{\boldsymbol{v}}_i\right) \end{gathered}\]② 不同的增强版本 [CLS] 标签上的自蒸馏,保证了在线 tokenizer 学习到高语义特征。
\[\begin{gathered} \mathcal{L}_{[\mathrm{CLS}]}(\boldsymbol{u},\hat{\boldsymbol{v}})=-P_{\boldsymbol{\theta}^{\prime}}^{[\mathrm{CLS}]}(\boldsymbol{u})^{\mathrm{T}} \log P_{\boldsymbol{\theta}}^{[\mathrm{CLS}]}(\hat{\boldsymbol{v}})\\ \mathcal{L}_{[\mathrm{CLS}]}(\boldsymbol{v},\hat{\boldsymbol{u}})=-P_{\boldsymbol{\theta}^{\prime}}^{[\mathrm{CLS}]}(\boldsymbol{v})^{\mathrm{T}} \log P_{\boldsymbol{\theta}}^{[\mathrm{CLS}]}(\hat{\boldsymbol{u}}) \end{gathered}\]为了帮助理解 MIM 想要学习的模式,将几种模式布局可视化。作者根据 ImageNet 验证集中所有图片 patch 的概率分布,可视化了部分类别中心所代表的模式。在大量的可视化结果中发现 iBOT 针对局部语义有非常好的可视化结果,如下图所示是 Patch 的 tokens 学习到的模式的可视化。左边2张图展示的模式是汽车的灯和狗的耳朵,展现了不同局部类别语义;右边2张图展示的模式是条纹和曲面展现了不同局部纹理语义。