Masked-attention Mask Transformer for Universal Image Segmentation
图像分割是关于将不同语义的像素分组，例如，类别或实例成员关系，其中每个语义的选择定义了一个任务。虽然只是每个任务的语义不同，但目前的研究重点是为每个任务设计专门的体系结构。我们提出了一种新的架构Mask -attention Mask Transformer (Mask2Former)，能够解决任何图像分割任务(全景、实例或语义)。它的关键组成部分包括Mask -attention，通过约束掩模区域内的交叉注意来提取局部特征。除了将研究工作量减少至少三倍之外，它在四个流行的数据集上的性能显著优于最好的专门架构。最值得注意的是，Mask2Former设置了一个新的最先进的全景分割(COCO上57.8 PQ)，实例分割(COCO上50.1 AP)和语义分割(ADE20K上57.7 mIoU)。
Mask2Former
本文提出通用的图像分割架构,它在不同的分割任务中优于专门的架构，同时仍然很容易在每个任务上训练。
Mask2Former vs MaskFormer
在maskformer的基础上： ? ? ? ? ? ?? ? ? ? ? ? ? ? ? ? ?
-
增加了masked attention机制
-
?使用多尺度高分辨率特征来帮助模型分割小目标/区域
-
调整了decoder部分的self-attention和cross-attention的顺序 ?
-
通过在少量随机采样点上计算mask loss，来加快训练速度
Mask2Former overview
Mask2Former采用与MaskFormer相同的元架构，具有backbone、像素解码器和transformer解码器。
1、backbone从图像中提取低分辨率特征的主干。
2、像素解码器，从主干的输出中逐渐上采样低分辨率的特征，以生成高分辨率的逐像素嵌入。
3、Transformer解码器，它通过图像特征处理对象查询。
Masked attention
背景特征已经被证明对图像分割非常重要。然而，最近的研究表明，基于transformer的模型收敛缓慢是由于交叉注意层的全局上下文，交叉注意学习关注局部对象区域需要许多训练时间。我们假设局部特征足以更新查询特征，通过自我注意可以收集上下文信息。为此，我们提出了掩码注意，这是交叉注意的一种变体，对每个查询只关注预测掩码的前景区域。
标准交叉注意(带有残差路径)计算
?masked attention模型的注意力矩阵计算：
特征位置(x, y)的注意掩码为 :
?
是前面第(l?1)个transformer解码器层调整掩码预测的二值化输出(阈值为0.5)。它被调整到相同的分辨率kl，M0为从X0得到的二进制掩码预测，即在向Transformer解码器输入查询特征之前的预测。
高分辨率特征
高分辨率特性提高了模型的性能，特别是对于小目标。然而，这需要计算。因此，我们提出了一种有效的多尺度策略，在控制计算量增加的同时引入高分辨率特征。我们不总是使用高分辨率特征图，而是使用一个特征金字塔，它包括低分辨率特征和高分辨率特征，并一次向一个Transformer decoder层提供一个多尺度特征的分辨率。
具体来说，我们使用pixel decoder生成的特征金字塔，分辨率分别为原始图像的1/32、1/16和1/8。对于每个分辨率，我们添加一个正弦位置嵌入，紧跟DETR，以及一个可学习尺度级嵌入，紧跟Deformable DETR。我们使用这些，从最低分辨率到最高分辨率的对应Transformer解码器层。我们重复这个3层Transformer解码器L次。因此，我们最终的Transformer解码器有3L层。更具体地说，前三层得到的特征图分辨率为H1 = H/32, H2 = H/16, H3 = H/8, W1 = W/32, W2 = W/16, W3 = W/8，其中H和W为原始图像分辨率。此模式以轮询方式对以下所有层重复。
优化改进
标准的transformer解码器（左图） vs.?本文的transformer解码器（右图）
? ? ? ??
标准的Transformer解码器层由三个模块组成，依次处理查询特性:自注意模块、交叉注意模块和前馈网络(FFN)。此外，查询特性(X0)在被输入到Transformer解码器之前是零初始化的，并且与可学习的位置嵌入相关联。此外，dropout应用于残差连接和注意图。
为了优化Transformer解码器的设计，本文做了以下三点改进。
-
首先，我们交换一下self-attention和cross-attention(我们的新“masked attention”)顺序,使计算更有效:第一个自我注意层的查询特征还不依赖于图像特征，因此应用自我注意不会产生任何有意义的特征。
-
其次，我们使查询特征(X0)也可学习(我们仍然保持可学习的查询位置嵌入)，可学习的查询特征在被用于Transformer解码器预测掩码(M0)之前是直接监督的。我们发现这些可学习的查询功能类似于区域提议网络RPN，并且能够生成mask proposal。
-
最后，我们发现dropout是不必要的，通常会降低性能。因此，我们在我们的解码器完全消除了dropout。
采样点损失函数
受 PointRend 和 Implicit PointRend 的启发，通过在 K个随机采样点而不是整个掩码上计算掩码损失来训练分割模型。在本文中设置K=12544，即112×112点。 本文在匹配损失和最终损失计算中使用采样点计算掩码损失：
1、在构建二分匹配成本矩阵的匹配损失中，对所有预测和真实掩码的相同的 K 点集进行均匀采样。
2、在预测之间的最终损失及其匹配的ground truth中，使用importance sampling为不同的预测和ground truth对不同的K点进行采样。
这种训练策略有效地将训练内存减少了 3 倍，从每张图像 18GB 到 6GB。