3 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]
3.1 本文的定位
我们从相关工作中尝试描述本文的定位。
该文的相关工作分为三部分:(1)传统分割模型;(2)基于Vision Transformer的分割;(3)改善Vision Transformer的训练
其中:(1)主要介绍了传统的各个分割任务的方法,例如语义分割FCN方法,实例分割Mask R-CNN方法,相比对于每个任务独立优化和提出独立架构的传统做法,MP-Former是一个统一分割模型(与MaskFormer,Mask2Former一样);(2)回顾受到DETR[1]启发后基于Transformer架构的图像分割模型,以及近期一些统一分割模型,如MaskFormer,Mask2Former。(3)指出Transformer架构相比CNN,其优势在于全局的注意力,但是该优势也为训练带来难度——因为全局搜索目标是困难的,导致Transformer训练往往是缓慢收敛。作者总结如[6][3]等工作是如何加速训练的:一是给予查询一个准确的意义避免模棱两可,二是在注意力机制中增加局部限制,减少运算量。MP-Former实现加速训练的途径更倾向思路一。
因此,从相关工作可以看出本文的定位是一个基于Transformer的统一的图像分割模型,同时设法改善了Transformer的训练。
3.2 本文的方法
如图5所示,本文的方法可分为以下三点:(1)多层掩膜引导的注意力;(2)掩膜加噪处理;(3)标签引导的训练。
- 多层掩膜引导的注意力
MP-Former在Transformer解码器部分增加了额外的查询嵌入向量和注意力掩膜,其中查询嵌入向量是类别真值对应的嵌入向量,注意力掩膜是查询对应的真值掩膜。在训练阶段,MP-Former在每一个Transformer解码器层都进行上述操作。这是因为,即便在第一层解码器层加入真值掩码作为注意力范围,其输出仍然是不完整,容易误导随后的解码器层。如图6所示,(a)为送入第一层解码器层的真值掩膜(完整地覆盖目标),(b)为第一层解码器层的预测掩膜(仅覆盖部分目标),如果将(b)作为第二层解码器层的注意力掩膜将会干扰其预测结果。因此,MP-Former提出在训练过程中增加多层掩膜引导部分,对Transformer解码器每一层的注意力掩膜都使用真值掩膜,而选择上一层解码器的预测掩膜(Mask2Former做法)。
图6:掩膜可视化
- 掩膜加噪处理
为缓解多个解码器层预测不一致问,MP-Former提供真值掩膜作为掩膜注意力的范围,使得掩膜预测能够逐层细化。但是这可能使得解码器的任务过于容易,阻碍解码器掩膜逐层优化的鲁棒性。因此,作者们提出送入带有噪声的真值掩膜,迫使解码器重建去噪后的真值掩膜,提升模型的鲁棒性。在噪声类型选取上,MP-Former实践了多种噪声,最终发现点噪声能够提高模型表现(即在掩膜上随机抹去一部分前景点和增加一些背景点)。
- 标签引导的训练
MP-Former使用实例对应的真值类别的嵌入向量作为掩膜引导(MP)部分中的查询,同时为了避免分类损失函数过于容易,也引入了一定的噪声。本文没有很清楚解释将类别嵌入向量作为查询的动机,但是考虑到解码器层预测掩膜的过程其实就是查询嵌入向量同图像特征点积计算,在解码器初始几层查询往往不够准确,导致初始几层掩膜不准确,引入真值类别嵌入向量作为查询等同于一个很好的初始点,能够帮助网络更快找到目标,同时和其他类别的目标能够区分。此外,给予查询一个准确的意义避免模棱两可,提高了Transformer的训练效率。
3.3 本文的实验
本文是对Mask2Former的继承和改进,所以作者通过COCO数据集上的语义分割与实例分割,ADE20K和Cityscapes数据集上的语义、实例、全景分割的表现证明了优化带来的提升显著。
图7:在ADE20K数据集上各图像分割任务性能
为验证设计的合理性,作者开展了相应的消融实验,验证(多层)掩膜引导部分和(类别嵌入、掩膜)加噪所带来的提升。通过测算,MP-Former相比Mask2Former显著加快训练过程。
3.4 本文的评价
本文的网络架构与Mask2Former[3]一致,推动了图像分割通用模型的发展,有助于集合研究资源更高质量改进图像分割技术。同时,针对Transformer架构应用到分割领域中的不足,如训练时间过长,以及针对具体工作Mask2Former中连续的Transformer解码器层预测掩膜不一致的问题,本文提出MP-Former,从训练方法上发力。具体地,MP-Former在Mask2Former原有的解码器层上增加掩膜引导部分,通过每一层注入带有噪声的真值掩膜引导注意力,既提高了掩膜预测的一致性,也让模型在“去噪”的过程中强化鲁棒性。针对Transformer解码器初始几层效果不佳,引入了类别真值嵌入向量作为查询,此类给予查询确定含义的做法被证明能够加速Transformer训练。作者通过三个常见的图像分割数据集验证了模型在训练效率和精度上的优势,以及应有的消融实验,总体来说实验部分是充分和具有逻辑的。此外,MP-Former的改进仅在训练时的设计上,在推理时没有引入额外的计算开销。
本文的核心贡献在于:提出了影响Mask2Former精度和训练效率的关键之一是其Transformer解码器层预测的不一致性,并通过定义合适的指标量化地验证这一问题。随后通过借鉴一些优秀的工作思路,一定程度缓解这一问题,严谨的实验显示提高了精度和训练效率。
开山之作MaskFormer提出了一个简明的框架,统一了图像分割任务,但也留下了许多优化空间。改进的总体方向是更高的精度和Transformer训练效率。回顾从MaskFormer到Mask2Former:为了让Transformer关注局部,掩膜注意力被提出。通过限制注意力的计算范围,省去没必要的计算,提高了模型的精度和训练效率。为了处理小目标但同时控制计算量,不同分辨率的多尺度特征被轮番送入Transformer解码器与查询基于交叉注意力相互作用。而MP-Former则通过引入真值掩膜进一步改善注意力的计算范围,同时为了一定的随机性,增加了点噪声扰动。
参考文献
- 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.
- 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.
- 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.
- 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.
- Ren, Shaoqing, et al. "Faster r-cnn: Towards real-time object detection with region proposal networks." Advances in neural information processing systems 28 (2015).
- 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