NeurIPS 2023 | RevCol V2: 当解耦学习遇上自监督训练，视觉预训练的思路

©作者 | Qi Han

单位 | 旷视研究院

RevColV2 基于可逆列网络，首次在 MIM 视觉预训练中实现了特征解耦学习，统一了上下游任务的网络结构，在 COCO 检测和 ADE20K 分割任务上达到了 62.1 的 AP 和 60.4 的 mIoU。

论文标题:

RevColV2: Exploring Disentangled Representations in Masked Image Modeling

论文链接:

https://arxiv.org/pdf/2309.01005.pdf

代码链接:

https://github.com/megvii-research/RevCol

前言

vision pre-training 作为视觉基础模型必不可少一个环节，始终承担着建立基础表征的职责。从监督训练，到对比学习训练，再到 MIM 预训练，视觉预训练一直在建立更”通用“的表征的道路上砥砺前行，时至 2023 年，vision pre-trianing 已经成为一个标准流程。那么，现有的"标准流程"中有什么问题吗？

▲ 图1. 主流预训练方式众所周知，视觉预训练旨在学到可以被很好泛化到各种下游任务上的通用特征，这就要求了 pre-training 阶段的表征要承载丰富多样的信息，同时又要语义丰富。图 1 展示了几种主流的预训练方法，其中黄色表示 detail 信息，深蓝色表示更 semantic 的信息。
可以看到，在监督学习中，由于 label 是 highly semantic 的分类标签，模型在输出层保留了高度的语义化的特征；而 MAE 这种利用 MIM 方式训练的 auto-encoders, 需要在 pre-training 阶段利用 decoder 来重建图像表征，然后在下游任务重抛弃掉 decoder，利用 encoder-only 的形式来建立表征。
这种 encoder-decoder 的人为界定，约束了语义信息出现的位置。而 BEiT 类模型在预训练中采用 discrete token 预测的形式，使用 encode-only 的架构，无需人为界定语义信息的位置，但在输出层杂糅了语义与细节信息，导致模型 feature 的直接利用性较差（体现为 linear porb 性能较低）。
虽然这几种架构都学到了有效的语义信息表征，但通用模型的角色要求了模型的表征不仅仅有丰富的语义表征，也需要有足够的细节 detail 来支撑各种下游任务。MAE 这种从 details - semantic - details 的结构，通过手动去掉后面 decoder 来让模型输出更多的 semantic 信息，势必会对 detail 特性有所损害，且这种预训练（auto-encoders）与下游（encoder-only）的模型差异，会丢失掉一些学到的 distribution 特性。

架构统一与解耦学习

为什么 MAE 这种 autoencoder 模型不能保留整体的 encoder-decoder 在下游使用呢？上一段的图 1 已经给出了答案，建模好的 semantic 特征在 decoder 中有所丢失，让下游任务的性能变差。在 NLP 任务上，encoder-decoder 架构是一种非常主流的结构，但在 CV 任务上，由于 pre-training 和 fine-tune 的输出空间不一致，需要在 pre-trian 的阶段中，既保留 pre-training 的输出空间，也要兼顾 fine-tune 的输出空间。
但视觉任务的下游并不固定，所以在 pre-training 阶段能保留下尽可能多的信息，同时又要建模丰富的语义特征，才是必要的。有没有什么网络架构能保留所有图像信息呢，自然可以想到可逆神经网络，在通用模型的表征学习中，reversible column network 是一种可逆的在输出空间保留全部信息同时又能够建模丰富语义的网络。详情可以参考：https://zhuanlan.zhihu.com/p/606594311
RevCol 作为一个现有的可以同时建模丰富语义信息且保留全部图像信息的模型，在通用模型领域开辟了一个新的思路。但 RevCol 的训练依赖于大量 labeled 数据来进行监督训练，例如原文章中使用 168M 数据进行训练。这里我们依托 RevCol 模型提出 V2 版本，在 MIM 领域探索解耦学习，让信息保留特性在自监督预训练中绽放光彩，且由于兼顾了预训练和下游的输出空间（语义信息+完整信息），模型可以无损的应用到下游任务重，不需要人为界定一个 decoder 来抛弃，统一了上下游模型结构。
RevColV2 最大的贡献就是 MIM 中的统一与解耦，在进行 MIM 预训练时，上下游采用相同的模型结构，保留了完整的 decoder 且兼顾了语义特性和 detail 信息。设计原则如图所示：

▲ 图2. RevColV2模型设计概览RevColV2 继承了之前的基础单元，仍然以 reversible column 为基础模块单元，通过堆叠不同的 column 来构建整个网络。在 V2 中，一部分 columns 承担了 encoder 的角色，另一部分扮演了 decoder 角色。
但这里的 encoder 和 decoder 并没有本质区别，decoder 也不是单纯的从语义信息回复图像本身。由于 reversible 的特性，每个 column 都保留了图像的完整信息，所以 encoder 和 decoder 都在同时建模 semantic 且保留 details，完整的信息流动如图 2 所示，从左到右的 columns 语义信息一直在加强，且 detail 信息依旧保留，最终形成了 output 特征空间中解耦的表达。

RevCol V2训练范式

▲ 图3. RevCol V2结构与训练方式
在这种由 columns 构成的 encoder decoder 结构中进行 MIM pre-trianing，unmasked patches 首先通过 encoder columns，然后在每个 level 的 output 上都和 MAE 一样加上 mask token，在 decoder columns 中继续进行建模。在最后一个 column 的输出中，不同的 level 联合构建了图像的全部信息，因此我们在 pre-training 阶段的 level-0 上进行图像补全的自监督，这种偏 detail 的监督信号，会在 reversible 的结构下迫使更高的 levels 建立更 semantic 的信息。
因此在下游任务重，相应的使用不同 level 的 features。对于分类任务，只需要最 semantic 的特征即可，因此只使用 level-3 的 output，对于检测、分割等需要 both high-level 特征，则使用所有 level 的 output，但这些特征是以解耦的形式表达了图像的全部信息，因此具有更好的泛化性。
更多的方法细节，欢迎阅读原 paper。

实验结果

RevColV2 的基础组件为 reversible columns，其中包含多个 levels，每个 level 由一个 transformer block 构成。在预训练阶段，采用 MIM 的训练方式在 ImageNet 上进行训练，然后迁移到下游进行评测。
在分类任务上，ImageNet1K 测试结果：

▲ 图4. 下游ImageNet分类结果

得益于结构的特征和统一的上下游结构，RevColV2 的 base 和 large 模型分别取得了 87.5 和 88.4 的 acc（with 22K ft），同时在分割和检测任务上，也有较为明显的提升。

▲ 图5. 下游COCO和ADE20K的检测和分割结果

更直接的joint pre-training
图 3 所示的 RevColV2 整体训练 pipeline 中，在 pre-training 中重建 raw image 来保持 detail 特征，通过 RevCol 解耦学习的特性来使 semantic 信息保留在其他 level。更进一步，可以显示的同时监督 semantic 部分与 detail 部分，在 detail 部分（bottom level）重建原始图像，而在最 semantic 部分（top level）重建 semantic feature。这里 semantic feature 我们采用 CLIP 来进行实验，使用两种不同的自监督信号 joint 建模了 detail+semantic 特征。
采用 OpenCLIP-L 作为 semantic feature 的重建，同时我们尝试把自监督阶段的数据 scale 到更大的量级（LAION400M），训练同样的 iterations，效果如下图所示：

▲ 图6. Joint Pre-training的下游结果联合训练的 RevColV2-L 模型在不经过 object365 pre-train 的情况下，在 COCO detection 上达到了 62.1 的 AP（cascade Mask RCNN）在 ADE20K 上也达到了 60.4 的 mIoU。

总结
RevColV2 作为一个视觉通用模型，首次将解耦学习成功应用到了 MIM 自监督训练中，通过解耦特征，在上下游任务上保持了相同的网络结构，增强了特征的泛化性。相比 RevCol 的 V1 版本，V2 不需要大量的 labeled 数据集来进行 supervised training，同时得益于 isotropic 的 ViT 结构，速度上也有了一些提升。
相比 MAE 等经典的自监督训练框架，RevColV2 首次在上下游中统一使用了 encoder-decoder 结构，不再受限于输出空间的限制。我们相信，这种将解耦特征学习与自监督预训练相结合的形式，能够为视觉预训练领域打开全新的思路。

更多阅读