BeiT v2 来袭 | BeiT升级，全面超越 MAE，实现 Vision Transformer 微调自由！

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 用于语义分割。

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显著提高了模型大小、训练步骤和下游任务的性能。

BEIT V2继承了用于Masked image modeling的BEIT框架。具体来说，给定输入图像，使用visual tokenizer将图像标记为离散的visual token。然后屏蔽一部分图像patch并将其输入Vision Transformer。预训练任务是根据损坏的图像恢复Mask的visual token。

使用Vision Transformer作为主干网络来获得图像表示。给定输入图像，将图像reshape为个patches，其中和是patch-size。在实验中，将每个 224×224 的图像分割成一个 14×14 的图像patch网格，其中每个patch为 16×16。然后将图像patch展平并线性投影到 Transformer 的输入嵌入中。对于 N 个图像块，将编码向量表示为。

visual tokenizer将图像映射到一系列离散tokens。具体来说，图像 x 被标记为，其中vocabularyV（即visual codebook）包含 |V| 离散code。请注意，token的数量与图像块的数量相同。本文提出向量量化知识蒸馏（VQ-KD）来训练visual tokenizer。如图 2 所示，VQ-KD在训练过程中有两个模块，即visual tokenizer和****。visual tokenizer由Vision Transformer编码器和量化器组成。分词器首先将输入图像编码为向量。接下来，向量量化器在codebook中查找每个patch表示的最近邻。让表示codebook嵌入。对于第 i 个图像块，其量化代码由下式获得：其中归一化用于codebook查找。上述距离相当于根据余弦相似度找code。在将图像量化为visual token之后，将-normalized codebook embeddings提供给****。****也是一个多层Transformer网络。输出向量旨在重建教师模型的语义特征，例如DINO和CLIP。让表示第 i 个图像块的教师模型的特征向量。最大化****输出和教师指导之间的余弦相似度。因为量化过程（等式 1）是不可微分的。如图 2 所示，为了将梯度反向传播到编码器，梯度直接从****输入复制到编码器输出。直观地说，量化器为每个编码器输出查找最近的code，因此codebook embeddings的梯度指示了编码器的有用优化方向。VQ-KD的训练目标是：其中，sg[·]表示停止梯度算子，它是向前传递时的一个身份，而在反向传递过程中梯度为零，D表示用于tokenizer训练的图像数据。

矢量量化训练的一个常见问题是codebook崩溃。换句话说，只使用了一小部分code。作者凭经验发现有几种技术可用于提高codebook利用率。于等人建议对codebook查找应用降维和 l2 归一化。等式 1 表明计算 l2 归一化距离以找到最近的code。此外，将查找空间的维度减少到 32 维。低维codebook嵌入在馈入****之前被映射回高维空间。此外，采用指数移动平均线来更新codebook嵌入，这在实验中往往更加稳定。

遵循BEIT中的MIM设置，以预训练Vision Transformers的图像表示。给定一个输入图像 x，分块选择大约40%的图像patch进行mask。如果if，其中δ（·）是指示函数，则使用共享的可学习嵌入来替换的原始图像patch嵌入。随后，在输入前准备了一个可学习的CLS token，即，并将它们提供给vision Transformer。最终的编码向量记为，其中为CLS token。接下来，使用一个mask图像建模head，基于损坏的图像来预测mask位置的visual tokens。对于每个mask位置，一个softmax分类器预测visual tokens，其中是mask图像，，是分类器的权重。visual tokens是由第2.2节中训练的标记器获得的，该节为mask图像建模提供了监督。最后，MIM的训练损失可以表述为：其中表示原始图像的visual tokens，D表示训练前的图像。

受(Gao和Callan，2021)的启发，明确地预训练了CLS token以进行全局表示。目标是减轻patch-level预训练和image-level表示聚合之间的差异。如图3所示，构建了一个representation bottleneck来指导CLS token收集信息。对于L-layer Transformer，让${h^l_i\}^N_{i=1}表示第层的输出向量，其中l=1····l$。为了预先训练最后一层的CLS token，连接中间l层的patch向量${h^l_i\}^N_{i=1}，即S=[h^L_{CLS},h^l_1,···,h^l_N]$。然后，将S输入一个浅层（例如，两层）Transformer decoder，并进行mask预测。注意，这里也计算了第l层的MIM损失，如公式3所示。所以最终的训练损失是两项的总和，即，在第l层的原始损失，和浅层Transformer decoder的MIM损失。在实现中还共享2个头的MIM Softmax权重。直观地看，该模型倾向于将全局信息推到，因为该模型倾向于充分利用从第1层到第l+1层的参数，以减少额外的MIM损失。information-flow bottleneck鼓励CLS token比未经训练的对应token获得更可靠的全局表示。新添加的浅层****仅用于对CLS token进行预训练，经过预训练后将被丢弃。

[1].BEIT V2: Masked Image Modeling with Vector-Quantized Visual Tokenizers

本文仅做学术分享，如有侵权，请联系删文。