图像分割高分作：Mask2Former框架一统分割任务

来源：投稿 作者：xin
编辑：学姐

论文标题：Masked-attention Mask Transformer for Universal Image Segmentation

论文链接：https://arxiv.org/pdf/2112.01527.pdf

作者信息：Facebook AI Research

Motivation：

当前对于分割任务的研究（如语义分割、实例分割、全景分割等分割任务）主要集中在为每个任务设计特定的框架，而本文致力于提出能媲美单个任务模型的统一模型。

Method：

故本文提出了Mask2Former的框架，用来实现分割任务的一统。

本文的创新点如下：

1、提出了Mask Attention机制，将注意力计算限制在局部特征中，而不是像cross-attention一样计算每个像素与整张图像的关系。

2、使用多尺度高分辨率特征帮助分割小目标。

3、提出了关于Transformer的优化方法；使用可学习的query以及随机采样等方法提高计算效率，节约内存。

在介绍本文提出的Mask2Former前，先回顾一下MaskFormer的基础结构。

整体框架如图所示，分为pixel-level模型，Transformer模型以及分割模型。首先通过backbone提取出图像特征，再将其送入解码器以生成像素嵌入特征。在Transformer模型中，使用低分辨率的图像特征作为K和V，结合0初始化的查询向量query通过Transformer decoder生成N个预分割嵌入向量Q。通过MLP将Q一支送去分类，另一只映射到像素嵌入的空间，进行mask的预测。最后将mask预测和类别预测进行矩阵相乘送去分割。

回顾完MaskFormer后，介绍本文提出的Mask2Former，该模型在MaskFormer上进行改进。首先也是最重要的改进，提出了Mask Attention机制，这使得每个通过Transformer的特征只在前景区域进行特证间的交互，而不是每个像素对整张图像都做注意力。本文基于此想法提出了假设：假设局部前景特征足以保证query向量的更新，以及局部前景特征通过自注意力机制可以获取有保证的上下文信息。相比于传统的Transformer的计算方法：

为了帮助分割小目标，本文使用了一种特征金字塔结构，该结构采用特征由低像素到高像素的策略，每次将不同分辨率的特征送入一个Transformer Decoder。具体而言是使用pixel decoder产生的1/32，1/16，1/8的特征，结合正弦位置嵌入

以及可学习的尺度规模嵌入

依次由低到高的送入Transformer Decoder，3次为一个循环，共循环L次（也即共送入3L个Transformer Decoder）。

在进行完上述改进后，本文还针对Transformer进行了进一步的改进。首先本文交换了self-attention和cross-attention（mask attention）的位置，因为第一层自注意力的查询特征与图像无关，没有来自图像的信号，故用自注意力不太可能丰富信息，所以进行交换。其次，本文将MaskFormer中0初始化的Query特征换为可学习的参数，并以次直接监督的生成。最后，作者发现Transformer中的Dropout是非必要的，故将Transformer Decoder中的Dropout其全部去掉。

最后为了提升训练效率，且受到PointRend和Implicit PointRend的启发，在匹配损失计算时，对所有预测值和mask掩码进行统一采样相同的K个点组成集合进行计算。在计算整体loss时，在不同对的预测值和真实值通过重要性采样选取不同的K个点进行损失的计算。K = 12544这样节省了3倍的内存，由18G->6G。

关注下方【学姐带你玩AI】🚀🚀🚀

回复“CVPR”领取图像分割方向必读论文

码字不易，欢迎大家点赞评论收藏！