构建图像分割任务的统一架构——从 MaskFormer 到 Mask2Former 再到 MP-Former

受到大语言模型的启发，视觉领域也掀起构建“大一统”框架的热潮。基于选定的MP-Former论文（CVPR 2023），本文回溯该工作的基础Mask2Former (CVPR 2022)和这一系列工作的起源MaskFormer(NeuraIPS 2021) ，介绍与分析了从MaskFormer到MP-Former的“变化”与“不变”。“不变”的是统一语义分割、实例分割、全景分割等图像分割任务的初心， “变化”的是更优的分割精度与更高的训练效率。

背景

统一的通用模型避免了对特定任务的设计需求，它们通过单一架构学习和执行各种任务，可以利用大规模多样化的数据语料库，有效地在不同任务之间转移概念知识，甚至可以执行在设计和训练时未知和未观察到的任务。在自然语言处理（NLP）中，尽管数据集和基准的领域非常广泛和多样化，但它们的输入和输出通常可以统一表示为一系列的Token序列。序列到序列（Seq2Seq）的架构被提出用于处理这种序列，广泛应用于各种NLP任务如文本分类、语言翻译、命名实体识别等。基于这种架构的统一模型一直是近期NLP进展的核心，包括近期大热的大语言模型（LLM）。

然而，构建用于计算机视觉（CV）的统一模型被证明是非常具有挑战性的，因为视觉任务具有极其多样化的输入和输出表示形式。例如，目标检测在图像中产生包围物体的边界框，图像分割生成图像中目标的二值掩膜，视觉问答将答案生成为文本，深度估计预测一张显示每个像素距离相机的距离的图像。仅图像分割领域，又包含语义分割、实例分割、全景分割等不同子任务。针对不同子任务，研究者们通常提出不同的框架，例如语义级分割常被建模成逐像素分类，而实例级分割（实例分割、全景分割）通常被建模为掩膜预测和分类，即预测一组二值掩膜，每个掩膜与单个类别相关联。不同的任务有着自己的最优网络架构，例如，将语义分割中常见的全卷积网络（FCN）架构应用于实例分割，往往效果不佳。

不同任务的割裂，导致了重复的资源投入。而统一的架构省去了为每个特定的任务设计一个特定的架构，能够大大节约研究投入，推动CV研究作为一个整体发展。

MaskFormer

在此背景下，研究者们提出了MaskFormer(NeuraIPS 2021)[2]，贡献了他们的洞察：掩膜分类足够通用，能够以统一的方式（相同的模型、损失、训练过程）解决语义分割、实例分割、全局分割。具体地，如图1所示，受到DETR[1]的启发，MaskFormer使用Transformer解码器预测一系列配对，每个配对由1个类别预测和1个掩膜嵌入向量组成。掩膜嵌入向量是用于同逐像素特征（像素解码器的输出）逐一做点积以获得预测的二值掩膜。见图2，无论对于语义分割还是实例分割，MaskFormer的监督信号都来自每一个掩膜。即，对于每个掩膜，给定真值，可计算二值掩膜预测的逐像素损失和类别预测的损失。针对不同具体任务，MaskFormer仅在推理策略上存在差异。

 图1：MaskFormer[2]

 图2：逐像素分类机制和逐掩膜分类机制的对比

MaskFormer是“掩膜分类足够解决任一图像分割任务”的朴素实践，其结构十分简洁，主要可以分为三部分：1）骨干网络：从图像中提取低分辨率特征；2）像素解码器：从主干输出中逐步上采样低分辨率特征，生成高分辨率的逐像素特征；3）Transformer解码器：在图像特征上进行操作，以处理对象查询，即输入对象查询的嵌入向量和图像特征计算交叉注意力。

Mask2Former

图3：Mask2Former[2]

基于MaskFormer，Mask2Former（CVPR 2022）进一步被提出优化统一架构的表现。如图3右侧所示，作者提出一个新的Transformer解码器，通过将网络注意力限制在每个查询对应的掩膜内（由查询嵌入向量同逐像素特征计算点积后阈值提取得到），提取对应的局部特征,而不是同全局像素特征计算注意力，最后经过大量训练周期学会关注目标的局部区域。掩膜注意力的设计一定程度缓解Transformer模型收敛缓慢的问题。

为了提高方法在小目标场景的表现，Mask2Former使用了一个高效的多尺度特征策略：利用像素解码器中的图像特征，构建一个由低分辨率和高分辨率特征组成的特征金字塔，每次将多尺度特征金字塔的一个层级送入Transformer解码器的一层。低分辨率和高分辨率特征交替送入Transformer解码器使得网络能更好捕捉不同尺度的目标，同时没有引入过多计算量。

Mask2Former在保证查询特征（查询嵌入向量）的位置编码是可学习的同时，还将查询特征设置为可学习的（MaskFormer中仅仅初始化为零向量），指出这能起到类似于Faster R-CNN[5]中RPN（区域建议网络）的效果。在训练效率提升上，为了减少内存开销，Mask2Former取K个随机采样点计算损失，而不是整个掩模上计算损失，来训练分割模型，相比MaskFormer实现其三分之一的内存消耗。

Mask2Former在提出时击败了为不同分割任务精心设计的专有架构，验证了统一架构的未来。

MP-Former: Mask-Piloted Transformer for Image Segmentation

MP-Former[4]是由港科大、IDEA研究院于CVPR 2023提出的基于的Mask2Former改进工作。作者们观察到Mask2Former在连续解码器层之间存在不一致的掩码预测（见图4），这导致了不一致的优化目标和解码器查询的低利用率。而且在解码器层早期，掩膜预测结果往往不够准确，但是在Mask2Former中会作为下一解码层的掩膜注意力范围，影响预测的准确性。这是因为：在每个解码器层中，预测的掩码是通过查询特征和逐像素特征计算点积获得，即每次从头开始构建，而不是逐层细化，导致不同解码器层中同一查询预测的掩码可能会发生巨大变化。本文通过计算连续的Transformer解码器层间同一查询预测的掩膜的交并比，衡量不同解码器层预测掩膜的一致性，定量地发现连续解码器层之间来自同一查询的预测之间存在严重的不一致问题。同时本文计算解码器每一层和最后一层匹配同一查询的真值掩码的比例，用以衡量该层查询的利用率，发现Mask2Former初始几层的查询利用率非常低，发生了很大的浪费。作者指出，掩膜预测的不稳定和查询的低利用率会导致优化目标不稳定和训练效率低。

分割和检测往往有很多相似之处，借鉴DN-DETR[6]中的去噪训练，作者们提出了一种掩码引导（Mask-Piloted）的训练方法，用于改善Mask2Former中掩码注意力的缺陷，稳定优化目标和提高训练效率。

图4：连续解码器层不一致的预测

图5：MP-Former[4]

参考文献

1. Carion, Nicolas, et al. “End-to-End Object Detection with Transformers.” Computer Vision – ECCV 2020,Lecture Notes in Computer Science, 2020, pp. 213–29, End-to-End Object Detection with Transformers | SpringerLink.
2. Cheng, Bowen, Alex Schwing, and Alexander Kirillov. "Per-pixel classification is not all you need for semantic segmentation." Advances in Neural Information Processing Systems 34 (2021): 17864-17875.
3. Cheng, Bowen, et al. "Masked-attention mask transformer for universal image segmentation." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.
4. Zhang, Hao, et al. "MP-Former: Mask-piloted transformer for image segmentation." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.
5. Ren, Shaoqing, et al. "Faster r-cnn: Towards real-time object detection with region proposal networks." Advances in neural information processing systems 28 (2015).
6. Feng Li, Hao Zhang, Shilong Liu, Jian Guo, Lionel M Ni, and Lei Zhang. Dn-detr: Accelerate detr training by introducing query denoising. arXiv preprint arXiv:2203.01305, 2022. 2, 3, 5