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

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

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

2.2、Training Visual Tokenizer

因为量化过程（等式 1）是不可微分的。如图 2 所示，为了将梯度反向传播到编码器，梯度直接从解码器输入复制到编码器输出。直观地说，量化器为每个编码器输出查找最近的code，因此codebook embeddings的梯度指示了编码器的有用优化方向。

VQ-KD的训练目标是：

其中，sg[·]表示停止梯度算子，它是向前传递时的一个身份，而在反向传递过程中梯度为零，D表示用于tokenizer训练的图像数据。

提高codebook利用率

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

2.3、Pretraining BEIT V2

预训练全局表示

受(Gao和Callan，2021)的启发，明确地预训练了CLS token以进行全局表示。目标是减轻patch-level预训练和 image-level 表示聚合之间的差异。

注意，这里也计算了第l层的MIM损失，如公式3所示。所以最终的训练损失是两项的总和，即，在第l层的原始损失，和浅层Transformer decoder的MIM损失。在实现中还共享2个头的MIM Softmax权重。

直观地看，该模型倾向于将全局信息推到，因为该模型倾向于充分利用从第1层到第l+1层的参数，以减少额外的MIM损失。information-flow bottleneck鼓励CLS token比未经训练的对应token获得更可靠的全局表示。新添加的浅层解码器仅用于对CLS token进行预训练，经过预训练后将被丢弃。