【论文笔记】Mask2Former: Masked-attention Mask Transformer for Universal Image Segmentation

论文地址：http://arxiv.org/abs/2112.01527
项目地址：https://bowenc0221.github.io/mask2former

本文提出了用于通用图像分割（全景、实例或语义）的 Masked-attention Mask Transformer (Mask2Former)。

Mask2Former建立在一个简单的元框架 （MaskFormer）和一个新的 Transformer 解码器上，其关键组成部分为掩码注意力（Masked-attention），通过将交叉注意力限制在预测的掩码区域内来提取局部特征。

与为每个任务（全景、实例或语义）设计专门模型相比，Mask2Former 节省了 3 倍的研究工作，并且有效节省计算资源。

Mask2Former 在全景分割（COCO 上的 57.8 PQ）、实例分割（COCO 上的 50.1 AP）和语义分割（ADE20K 上的 57.7 mIoU）上都实现了SOTA。

本文提出的关键改进在于：
1、在 Transformer 解码器中使用屏蔽注意力，将注意力限制在以预测段为中心的局部特征上（根据分组的特定语义，可以是对象或区域）。与关注所有位置的标准 Transformer 解码器中使用的交叉注意力相比，掩码注意力可以得到更快的收敛和更好的性能。
2、使用多尺度高分辨率特征，帮助模型分割小对象/区域。
3、提出了优化改进，例如切换自注意力和交叉注意力的顺序，使查询特征可学习，去除dropout；所有这些都无需额外计算即可提高性能。
4、通过在K个随机采样点上计算掩码损失，在不影响性能的情况下节省了 3 倍的训练内存。

模型整体架构

本文的模型整体架构来自于MaskFormer的简单的元架构，由三个组件组成：

1. 主干特征提取器：从图像中提取低分辨率特征的；
2. 像素解码器：从主干的输出中逐渐上采样低分辨率特征生成高分辨率的每像素嵌入；
3. Transformer 解码器：对图像特征进行操作以处理对象查询。 最终的二元掩码预测是从带有对象查询的每像素嵌入中解码的。

本文主要改进在Transformer 解码器中。

屏蔽注意力（Masked Attention）

Masked Attention 操作符，通过将交叉注意力限制在每个查询的预测掩码的前景区域内来提取局部特征，而不是关注完整的特征图。

标准交叉注意（带有残差路径）计算公式如下

掩码注意力通过添加掩码值调节注意力矩阵

高分辨率特征

高分辨率特征可以提高模型性能，特别是对于小目标。但增加了对计算的要求。
本文提出了一种有效的多尺度策略来引入高分辨率特征，同时控制计算量的增加。
利用一个由低分辨率和高分辨率特征组成的特征金字塔，并将多尺度特征的不同尺度特征分别馈送到不同的Transformer解码器层。
对于每个分辨率，添加了一个正弦位置嵌入以及一个可学习的尺度级嵌入。
重复3层Transformer 解码器L次，最终的 Transformer 解码器因此具有 3L层。

Transformer 解码器优化改进

为了优化 Transformer 解码器设计，进行了以下三点改进：

1、切换自注意力和交叉注意力（新的“掩蔽注意力”）的顺序以使计算更有效：对第一个自注意力层的查询特征还不依赖于图像特征，因此应用自注意力不会产生任何意义。
2、使查询特征（X0）也可学习（仍然保留可学习的查询位置嵌入），并且可学习的查询特征在用于 Transformer 解码器的预测掩码（M0）之前直接受到监督。 作者发现这些可学习的查询特征功能类似于region proposal network [43]，并且能够生成掩码proposals。
3、dropout 不是必需的，通常会降低性能。作者因此在解码器中完全删除了 dropout。

提高训练效率

训练通用架构的一个限制是由于高分辨率掩码预测而导致的大量内存消耗，这使得它们比对内存更友好的专用架构更难以接收。例如，MaskFormer 只能在具有 32G 内存的 GPU 中容纳单个图像。

受 PointRend [30] 和 Implicit PointRend [13] 的启发，可以通过在 K个随机采样点而不是整个掩码上计算掩码损失来训练分割模型。在本文中设置K=12544，即112×112点。

本文在匹配损失和最终损失计算中使用采样点计算掩码损失。

• 在构建二分匹配成本矩阵的匹配损失中，对所有预测和真实掩码的相同的 K 点集进行均匀采样。
• 在预测之间的最终损失及其匹配的ground truth中，使用重要性采样[30]为不同的预测和ground truth对不同的K点集进行采样。

这种训练策略有效地将训练内存减少了 3 倍，从每张图像 18GB 到 6GB。

训练细节

使用 AdamW [38] 优化器和步进式学习率计划。
对所有主干使用 0.0001 的初始学习率和 0.05 的权重衰减。
模型训练 50 个epochs，批次大小为 16。
对于数据增强，使用大规模抖动 (LSJ) 增强 ，随机尺度从 0.1 到 2.0 范围内采样，然后是固定大小的裁剪到 1024×1024。