自监督学习系列（四）：基于蒸馏的图片掩码学习
前文
好久不见！自监督系列文章继续更新啦！在前几期的文章我们介绍了基于辅助任务，对比学习，和图片掩码学习的自监督学习范式
(对比学习，图片掩码学习其实也可以归属于基于辅助任务的自监督学习，由于这两类是目前最流行的研究方向，所以我们在此进行这样的区分)。
有感兴趣的读者可以查看前几期的文章:
https://zhuanlan.zhihu.com/p/470914640
https://zhuanlan.zhihu.com/p/474847821
https://zhuanlan.zhihu.com/p/475952825
由于 BEiT 和 MAE 的推出，目前自监督预训练领域的发展仍然是：改进现有的图片掩码学习模型。从前一期的文章可以知道，很多掩码学习的工作都以纯 RGB 像素值作为预测的目标，这些方法的具体操作步骤为:
- 将一张图片划分为多个 patch；
- 随机覆盖掉其中的一部分 patch, 并将剩余的 patch 送入网络；
- 用网络预测的被覆盖掉的 patch 去拟合这些 patch 的真实的像素值。
这些方法以 MAE 为代表，这样做的好处是整个模型的结构简单，计算量小。但是纯 RGB 像素作为预测目标也有两个明显的劣势:
- 纯像素值是 low-level 的特征，模型往往需要预训练更长的迭代次数才能够达到一个不错的效果；
- 模型容易通过训练一些捷径，例如像素插值，来完成这样一个重建的任务，所以采用纯 RGB 像素作为预测目标的模型往往需要一个非常大的 mask ratio，例如，MAE 采用了 0.75，SimMIM 采用了 0.6。
目前大多数方法主要是针对下面两个方面进行优化:
- 重建目标，将 RGB 像素替换为具有丰富语义信息的模型输出的特征，例如，CLIP 输出特征；
- 引入额外的任务，例如，CMAE 除了重建任务以外，引入了一个对比学习的任务。
今天这篇文章就主要围绕重建目标介绍几篇典型的文章：MILAN、BEiT-v2、MASKDISTILL。下表展示了这几个算法的概要。
| Method | Backbone | Decoder | Target | lin | ft | seg |
|---|---|---|---|---|---|---|
| MILAN | ViT-B | Prompt-decoder | CLIP | 78.9 | 85.4 | 52.7 |
| MASKDISTILL | ViT-B | FC | CLIP | - | 85.5 | 54.3 |
| BEiT-v2 | ViT-B | FC | VQ-KD | 80.1 | 85.0 | 52.7 |
MASKDISTILL
TL;DR: 图片掩码学习的统一范式
MASKDISTILL 站在一个更高角度审视现有的图片掩码模型算法，将其抽象为下图的结构，简而言之就是，一张图片经过 mask 送入一个 encoder，这些 encoder 的选择多种多样，例如 ViT 和 Swin；然后将 encoder 输出的特征送入一个 head，这个 head 可能是一个线性层，也有可能是一个具有多层的 decoder。
图 1
而对于 target 特征，也存在多种选择，例如纯像素、CLIP 之类的，然后 target 经过一个 normalization 层，最终输出，并完成和预测特征的匹配。
基于这个统一的框架，MASKDISTILL 使用 CLIP 作为 target，以及使用 smooth-l1 loss，在没有其他额外技术的加持下，也可以取得不错的效果。下表展现了 MASKDISTILL 与目前大多数主流自监督学习方法的一个比较：
MILAN
划重点：CLIP target, Prompting Decoder, Semantic Aware Sampling
另外 MMSelfSup 现已支持 MILAN，欢迎体验 ：
https://github.com/open-mmlab/mmselfsup/tree/dev-1.x/configs/selfsup/milan
从下图可以看出 MILAN 具有三个重要部分:
- CLIP target
- Prompting Decoder
-
Semantic Aware Sampling
图 2
CLIP target
CLIP 为 OpenAI 推出的多模态大规模预训练模型。其通过在大量的图文对上进行大规模的对比学习，使得 image encoder 分支具有捕捉图片高维语义的能力，其训练范式如下图所示：
图 3
使用 CLIP 特征作为预测 target 可以极大地提升模型的性能，从下表第 1 和第 2 组消融实验可以看出。同时，相较于其他的 target，例如 SLIP，CLIP 具有明显的优势，这一点从下表的第 8 组和第 9 组实验可以看出。
其实，相较于这些模型，之前的模型例如 BEiT 也采用了预训练模型（DALLE）输出的特征作为预测的，但是由于这些预训练模型本身性能的问题，使得最终模型的性能也是不尽人意。
Prompt Decoder
相较于 MAE 采用的使用自注意力机制作为 decoder 的基础模块，MILAN 采用的是一种类似于 CAE 所使用的 cross attention 机制。
这种机制最主要的特征是在解码阶段只会去更新 mask token 的特征，而保持 encoder 输出的可见 token 的特征不变。
这样做的一个明显的好处就是可以降低解码阶段的计算量。同时在与 CLIP target 配合使用的时候，可以进一步提升模型的精度。
Semantic Aware Sampling
目前大多数图片掩码学习的算法都是使用的随机掩码策略，这种策略虽然简单，但是对需要 mask 的区域进行无差别的对待，可能使得很多具有更多语义信息的地方被 mask 掉（如，前景物体）。
所以 MILAN 采用了一个 semantic aware 的 mask 策略。具体而言就是，通过 CLIP 输出的特征，计算各个地方被 mask 的概率，依赖这种概率去 mask 图片，使得具有丰富语义信息的地方得以保留。
由于 CLIP 的特征具有提取图片语义的能力，所以这种策略是合理的。通过上表的第 4 组实验也可以看出，这种策略也是有效的。
BEiT-v2
TL;DR: VQ-KD 训练, Class Token Pretraining
MMSelfSup 现已支持 BEiT( https://github.com/open-mmlab/mmselfsup/tree/dev-1.x/configs/selfsup/beit ), 同时 BEiT-v2 也在加速支持中
BEiT-v2 是 BEiT 的拓展，两者主要的差距在于以下两点:
- BEiT-v2 使用 CLIP 作为 teacher，自己设计了一套范式蒸馏得到了一个具有更好性能的目标提取器；
- 为了使得 class token 具有整张图的全局表达能力，将 encoder 最后一层的 cls token 连同其中某一层的特征一起送入了一个比较轻量的 decoder，从而引入另外一个重建损失。
VQ-KD
VQ-KD 的训练方式如下图所示，对于每一张图片，encoder 吐出的特征去一个预先定义好的 Codebook 寻找一个离当前特征最近的一个特征，然后将这个找到的特征喂入后面的 decoder 重建 CLIP 输出的特征。当训练结束之后，只有图中的 tokenizer encoder 被用于后面预训练中的目标特征提取。
图 4
Class Token Pretraining
如图 5 的右侧所示，Class Token Pretraining 的具体做法是将 encoder 最后一层的 class token 连接到encoder 的其中某一层输出的 token，这新组成的一组输入喂入额外的一个 encoder 中进行解码，最终完成整张图特征的重建。这样做的目的是为了改善 class token 的全局特征表达能力。
图 5
这样做，本人认为基于的一个假设应该是：只有最后一层的 class token 具有良好的全局表达能力，其才可能在解码阶段将图片的全局语义传递给其他 token，从而完成整张图的重建。通过下表的消融实验也可以发现，class token pretraining 可以极大地改善模型的 linear probing 精度。
总结
今天我们讲了目前几种基于 CLIP 的图片掩码学习算法，总结一下，采用 CLIP 的特征作为目标特征，由于这种特征包含丰富的语义信息，相较于使用纯像素的方法，其收敛更快（例如，MILAN 预训练 400 epoch 也能够达到不错的效果）, 下游表现更好。
但是，采用这种方式也会让整个预训练的 pipeline 变得更加复杂，同时引入额外的计算量。
目前，MMSelfSup 已经支持了目前主流的图片掩码学习算法，如 MAE, SimMIM, CAE，以及上面提到的 MILAN, 对于上文提到的其他几种蒸馏类型的图片掩码学习算法也在积极支持中, 欢迎大家使用 MMSelfSup，如果发现不错，也希望大家多多 Star、Issue 和 PR！