作者丨ChaucerG
来源丨集智书童
编辑丨极市平台
论文链接：https://arxiv.org/abs/2208.06366
代码链接：https://github.com/microsoft/unilm/tree/master/beit
Masked image modeling 通过恢复损坏的图像块在自监督表示学习中展示了令人印象深刻的结果。然而，大多数方法仍然对 low-level 图像像素进行操作，这阻碍了对表示模型的 high-level 语义的利用。
在这项研究中，作者建议使用语义丰富的 visual tokenizer 作为 Mask 预测的重建目标，为将 MIM 从像素级提升到语义级提供了一种系统的方法。
具体来说，引入向量量化知识蒸馏来训练 tokenizer，它将连续的语义空间离散化为 compact codes。然后，通过预测 masked image patches 的原始 visual tokenizer 来预训练 Vision Transformers 。此外，作者鼓励模型将patches信息显式聚合到全局图像表示中，这有助于线性预测。
图像分类和语义分割的实验表明，本文的方法优于所有比较的 MIM 方法。在 ImageNet-1K（224 大小）上，base-size BEIT V2 在微调时达到 85.5% 的 top-1 精度，在线性预测时达到 80.1% 的 top-1 精度。large-size BEIT V2 在 ImageNet-1K（224 大小）微调上获得 87.3% 的 top-1 准确率，在 ADE20K 上获得 56.7% 的 mIoU 用于语义分割。
1、简介
Masked image modeling 在学习视觉表示方面显示出令人印象深刻的结果，这极大地缓解了 Vision Transformer 的注释饥饿问题。给定一张图像，这些方法通常首先通过mask一些 patches 来破坏原始图像，相关任务是恢复原始图像。以开创性工作 BEiT 为例，每张图像在预训练期间都有2个视图，即图像 patches 和 visual tokenizer。原始图像首先被tokenizer为离散 token。随机采样的图像 patches 在被馈送到 Vision Transformer 之前被 Masked。预训练的目标是根据损坏的图像 patches 恢复原始 visual tokenizer。在预训练视觉编码器后，可以通过附加轻量级任务层直接在各种下游任务上微调模型。
在 mask-then-predict 框架下，与之前工作的主要区别在于重建目标，例如 visual tokenizer、原始像素和手工制作的 HOG 特征。然而，恢复low-level监督往往会浪费建模能力来预训练高频细节和短程依赖关系。例如，当掩盖戴在男人头上的“帽子”时，更喜欢模型在给定整个上下文的情况下学习被掩盖的“帽子”的高级概念，而不是在像素级细节上苦苦挣扎。相比之下，语言建模中的掩码词通常被认为具有比像素更多的语义。这促使通过在预训练期间利用语义感知监督来挖掘 MIM 的潜力。
在这项工作中，作者引入了一种自监督的视觉表示模型 BEIT V2，旨在通过学习语义感知的 visual tokenizer 来改进 BEIT 预训练。具体来说，提出了向量量化知识蒸馏（VQ-KD）算法来离散化语义空间。VQ-KD 编码器首先根据可学习的 codebook 将输入图像转换为离散token。然后解码器学习重建由教师模型编码的语义特征，以离散token为条件。在训练 VQ-KD 后，其编码器用作 BEIT 预训练的 visual tokenizer，其中离散 code 用作监督信号。
此外，作者建议通过明确鼓励 CLS token 聚合所有 patches 来预训练全局图像表示。该机制解决了 Masked image modeling 仅预训练 patch-level 表示的问题。结果，在聚合全局表示的帮助下，线性预测的性能得到了提高。
在 ImageNet-1k 上对 base 和 big Vision Transformer 进行自监督学习，并在多个下游任务上进行评估，例如图像分类、线性预测和语义分割。如图 1 所示，BEIT V2 在 ImageNet 微调上大大优于以前的自监督学习算法，例如，在 ViT-B/16 和 ViT-L/16 上都比 BEIT 提高了大约2个点。BEIT V2在 ImageNet 线性预测上优于所有比较的 MIM 方法，同时在 ADE20k 上实现语义分割的巨大性能提升。
本研究的贡献总结如下：
• 引入向量量化知识蒸馏，将Masked image modeling从像素级提升到语义级，以进行自监督表示学习。
• 提出了一种Patch聚合策略，该策略在给定 Patch-Level Masked image modeling 的情况下强制执行全局表示。
• 对下游任务进行了广泛的实验，例如 ImageNet 微调、线性预测和语义分割。实验结果表明，BEIT V2显著提高了模型大小、训练步骤和下游任务的性能。
2、BeiT v2
BEIT V2 继承了用于Masked image modeling的 BEIT 框架。具体来说，给定输入图像，使用visual tokenizer将图像标记为离散的visual token。然后屏蔽一部分图像patch并将其输入 Vision Transformer。预训练任务是根据损坏的图像恢复Mask的visual token。
2.1、Image Representations
x reshape为 N=HW/P2 个patches {xip}i=1N ，其中 xp∈RN×(P2C) 和 (P,P) 是patch-size。在实验中, 将每个 224×224 的图像分割成一个 14×14 的图像patch网格, 其中每个patch为 16×16。然后将图像patch {xip}i=1N 展平并线性投影到 Transformer 的输入嵌入中。对于 N 个 图像块, 将编码向量表示为 。{hi}i=1 。 N
2.2、Training Visual Tokenizer
visual tokenizer 将图像映射到一系列离散 tokens 。具体来说, 图像 x 被标记为 z=[z1,⋯,zN]∈V(H/P)×(W/P) ，其中 vocabulary ∨ (即 visual codebook）包含 |V| 离散 code。请注意, token 的数量与图像块的数量相同。本文提出向量量化知识蒸馏 ( VQ-KD ) 来训练 visual tokenizer。如图 2 所示, VQ-KD 在训练过程中有两个模块, 即 visual tokenizer 和解码器。
visual tokenizer 由 Vision Transformer 编码器和量化器组成。分词器首先将输入图像编码为向量。接下来, 向量量化器在 codebook 中查找每个 patch 表示 hi 的最近邻。让 e1,…,e|V| 表示 codebook 嵌入。对于第 i 个图像块, 其量化代码由下式获得:
zi=argminj‖ℓ2(hi)−ℓ2(ej)‖2
其中 l2 归一化用于 codebook 查找。上述距离相当于根据余弦相似度找 code 。
在将图像量化为 visual token 之后, 将 l2-normalized codebook embeddings {l2(ezi)}i=1N 提供给解码器。解码器也是一个多层 Transformer 网络。输出向量 {oi}i=1N 旨在重建教师模型的语义特征, 例如 DINO 和 CLIP。让 ti 表示第 i 个图像块的教师模型的特征向量。最大化解码器输出 oi 和教师指导 ti 之间的余弦相似度。
因为量化过程（等式 1）是不可微分的。如图 2 所示，为了将梯度反向传播到编码器，梯度直接从解码器输入复制到编码器输出。直观地说，量化器为每个编码器输出查找最近的code，因此codebook embeddings的梯度指示了编码器的有用优化方向。
VQ-KD的训练目标是：
max∑x∈D∑i=1Ncos(oi,ti)−‖sg[ℓ2(hi)]−ℓ2(ezi)‖22−‖ℓ2(hi)−sg[ℓ2(ezi)]‖22
其中, sg[⋅] 表示停止梯度算子, 它是向前传递时的一个身份, 而在反向传递过程中梯度为零, D表示用于 tokenizer 训练的图像数据。
提高codebook利用率
矢量量化训练的一个常见问题是 codebook 崩溃。换句话说，只使用了一小部分code。作者凭经验发现有几种技术可用于提高 codebook 利用率。于等人建议对codebook查找应用降维和 l2 归一化。等式 1 表明计算 l2 归一化距离以找到最近的 code。此外，将查找空间的维度减少到 32 维。低维codebook嵌入在馈入解码器之前被映射回高维空间。此外，采用指数移动平均线来更新codebook嵌入，这在实验中往往更加稳定。
2.3、Pretraining BEIT V2
遵循 BEIT 中的 MIM 设置, 以预训练 Vision Transformers 的图像表示。给定一个输入图像 x, 分块选择大约 40% 的图像 patch 进行 mask 。
如果 eip if
i∈M:xiM=δ(i∈M)⊙e[M]+(1−δ(i∈M))⊙xip, 其中δ(⋅)是指示函数, 则使用共享的可学习嵌入 e[M] 来替换 eip 的原始图像 patch 嵌入 eip 。随后, 在输入前准备了一个可学习的 CLS token , 即 [eCLS,{xiM}i=1N], 并将它们提供给 vision Transformer 。最终的编码向量记为 {hi}i=0N, 其中 h0 为 CLS token 。
接下来, 使用一个 mask 图像建模 head,基于损坏的图像 xM 来预测 mask 位置的 visual tokens。对于每个 mask 位置 {hi:i∈M}i=1N, 一个 softmax 分类器预测 visual tokens p(z′∣xM)=softmaxz(Wchi+bc), 其中 xM 是 mask 图像, Wc,bc 是分类器的权重。 visual tokens 是由第 2.2 节中训练的标记器获得的,该节为 mask 图像建模提供了监督。最后, MIM 的训练损失可以表述为：
LMIM=−∑x∈D∑i∈Mlogp(zi∣xM)
其中 zi 表示原始图像的 visual tokens, D表示训练前的图像。
预训练全局表示
受(Gao和Callan，2021)的启发，明确地预训练了CLS token以进行全局表示。目标是减轻patch-level预训练和 image-level 表示聚合之间的差异。
如图3所示, 构建了一个 representation bottleneck 来指导 CLS token 收集信息。对于 L-layer Transformer, 让 ${ h∧∣ _ i}∧N−{i=1} 表示第 1 层的输出向量, 其中 |=1⋯|$ 。为了预先训练最后一层的 CLS token hCLSL, 连接中间I层的 patch 向量 ${ h∧I _ i∖}∧N−{i=1} decoder，并进行mask预测。
注意，这里也计算了第l层的MIM损失，如公式3所示。所以最终的训练损失是两项的总和，即，在第l层的原始损失，和浅层Transformer decoder的MIM损失。在实现中还共享2个头的MIM Softmax权重。
直观地看，该模型倾向于将全局信息推到，因为该模型倾向于充分利用从第1层到第l+1层的参数，以减少额外的MIM损失。information-flow bottleneck鼓励CLS token比未经训练的对应token获得更可靠的全局表示。新添加的浅层解码器仅用于对CLS token进行预训练，经过预训练后将被丢弃。
