[Paper Review] Masked-attention Mask Transformer for Universal Image Segmentation

Bowen Cheng, Ishan Misra+3 more|arXiv (Cornell University)|Dec 2, 2021

Advanced Neural Network ApplicationsComputer Science57 references127 citations

TL;DR

Mask2Former is a universal image segmentation architecture using masked attention that outperforms specialized models across panoptic, instance, and semantic tasks on multiple datasets, with improved training efficiency.

ABSTRACT

Image segmentation is about grouping pixels with different semantics, e.g., category or instance membership, where each choice of semantics defines a task. While only the semantics of each task differ, current research focuses on designing specialized architectures for each task. We present Masked-attention Mask Transformer (Mask2Former), a new architecture capable of addressing any image segmentation task (panoptic, instance or semantic). Its key components include masked attention, which extracts localized features by constraining cross-attention within predicted mask regions. In addition to reducing the research effort by at least three times, it outperforms the best specialized architectures by a significant margin on four popular datasets. Most notably, Mask2Former sets a new state-of-the-art for panoptic segmentation (57.8 PQ on COCO), instance segmentation (50.1 AP on COCO) and semantic segmentation (57.7 mIoU on ADE20K).

Motivation & Objective

• Motivate a universal segmentation architecture that can handle panoptic, instance, and semantic tasks with a single design.
• Eliminate fragmentation by showing a single model can perform well across multiple segmentation tasks.
• Improve training efficiency and accessibility for universal segmentation without sacrificing accuracy.
• Demonstrate that masked attention and multi-scale high-resolution features enhance convergence and performance.
• Provide ablations to understand which components drive gains across tasks.

Proposed method

• Introduce a meta architecture with a backbone, pixel decoder, and Transformer decoder operating on mask queries.
• Replace standard cross-attention with masked attention that attends only within predicted mask regions.
• Use a multi-scale high-resolution feature strategy to better segment small objects.
• Apply optimization changes: reorder self- and cross-attention, make query features learnable, and remove dropout.
• Reduce training memory by computing mask loss on a limited set of randomly sampled points.

Experimental results

Research questions

• RQ1 Can a single architecture achieve state-of-the-art performance across panoptic, instance, and semantic segmentation?
• RQ2 Do masked attention and high-resolution multi-scale features speed up training and improve accuracy?
• RQ3 What optimization and training strategies enable practical training of universal segmentation models?
• RQ4 How does Mask2Former compare to specialized and other universal architectures across common datasets?
• RQ5 Is the approach robust across diverse backbones and datasets?

Key findings

• Mask2Former achieves state-of-the-art or on-par performance on panoptic, instance, and semantic segmentation across multiple datasets.
• It sets a new state-of-the-art: 57.8 PQ on COCO panoptic, 50.1 AP on COCO instance, and 57.7 mIoU on ADE20K semantic segmentation using the same architecture.
• Masked attention substantially improves performance and convergence compared to standard cross-attention.
• High-resolution multi-scale features significantly boost small-object segmentation, with an efficient multi-scale strategy reducing FLOPs.
• Training efficiency improved through learnable queries, removal of dropout, and sampling-based mask loss, reducing memory while maintaining accuracy.