PaperlayerView exact revision on arXiv

Source-linked AI summary

Masked-attention Mask Transformer for Universal Image Segmentation

Bowen Cheng, Ishan Misra, Alexander G. Schwing, Alexander Kirillov, Rohit Girdhar

arXiv:2112.01527v3cs.CVcs.AIcs.LG

TL;DR

Current segmentation methods use specialized architectures for different tasks despite sharing the same pixel-grouping problem. Mask2Former introduces a universal architecture with masked attention, matching or exceeding specialized models across panoptic, instance, and semantic segmentation, including 57.8 PQ on COCO panoptic segmentation.

  • Problem

    Despite differing mainly in grouping semantics, image segmentation tasks rely on specialized architectures that lack flexibility across panoptic, instance, and semantic segmentation.

  • Method

    Mask2Former uses a universal segmentation architecture with Transformer-decoder masked attention that restricts attention to localized features around predicted segments.

  • Results

    Across three segmentation tasks and four datasets, the same Mask2Former architecture performs on par with or better than specialized models, reaching 57.8 PQ on COCO panoptic segmentation.

  • Takeaways & Limitations

    Mask2Former provides one architecture for the three major image segmentation tasks while reducing the research effort required to design specialized models.

  • Takeaways & Limitations

    Although Mask2Former generalizes across tasks, it still needs task-specific training rather than being trained once for multiple tasks.

Abstract

from arXiv · show

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).

1. Introduction

Mask2Former addresses the fragmentation and performance gap caused by task-specific segmentation architectures with a universal architecture for panoptic, instance, and semantic segmentation. Across four datasets, it matches or surpasses specialized architectures and achieves state-of-the-art results on COCO and ADE20K.

  • Problem: Image segmentation tasks differ in grouping semantics, but existing methods typically develop a specialized architecture for each task.Examples include per-pixel classification architectures for semantic segmentation and mask classification architectures for other tasks.
  • Problem: Universal architectures use one architecture, loss, and training procedure across segmentation tasks, but their performance lags behind specialized methods.They are commonly based on end-to-end set prediction objectives such as DETR.
  • Problem: > 9 AP separates universal architectures from the state-of-the-art specialized architecture for instance segmentation.Universal architectures are also harder to train, requiring more advanced hardware and longer training schedules.
  • Mask2Former: Mask2Former is a universal architecture designed to outperform specialized architectures across segmentation tasks while remaining easy to train.It builds on a meta-architecture with a backbone feature extractor, pixel decoder, and Transformer decoder, and introduces key improvements for better results and efficient training.
  • Results: 57.8 PQ on COCO panoptic segmentation, 50.1 AP on COCO instance segmentation, and 57.7 mIoU on ADE20K semantic segmentation establish new state-of-the-art results.The single architecture performs on par with or better than specialized architectures across three tasks and four popular datasets.

2. Related Work

Related work develops specialized architectures for semantic and instance segmentation. Semantic methods typically classify pixels individually, while instance methods predict class-associated binary masks.

  • Semantic segmentation architectures typically formulate the task as per-pixel classification.FCN-based architectures independently predict a category label for every pixel.
  • Follow-up semantic methods emphasize context for precise per-pixel classification.They design customized context modules or self-attention variants.
  • Instance segmentation architectures typically use mask classification, predicting binary masks associated with individual class labels.

3. Masked-attention Mask Transformer

Mask2Former builds on a general mask-classification architecture and replaces its standard Transformer decoder with masked attention, complemented by efficient multi-scale features and training improvements for accessibility.

  • Architecture: Mask2Former adopts a mask-classification meta architecture and replaces its standard Transformer decoder with a masked-attention decoder.The section presents the meta architecture first, then introduces the decoder and training improvements.
  • Masked attention: Masked attention constrains cross-attention to each query’s predicted foreground mask, extracting localized features instead of attending to the full feature map.The mask is the resized, binarized prediction from the previous decoder layer, thresholded at 0.5.
  • Multi-scale features: A multi-scale strategy feeds one feature resolution to each Transformer decoder layer, incorporating high-resolution features while controlling computational cost.The feature pyramid uses resolutions 1/32, 1/16, and 1/8 of the original image.
  • Decoder improvements: The decoder design improves computation by ordering self-attention before masked cross-attention, making query features learnable, and directly supervising them before decoding.The ordering avoids applying self-attention to image-independent query features before they receive image signals.
  • Efficient training: Point-based mask-loss computation addresses universal architectures’ memory demands by calculating mask loss on K randomly sampled points instead of the whole mask.The approach is motivated by PointRend and Implicit PointRend, while MaskFormer can fit only one image on a GPU with 32G memory.

4. Experiments

Mask2Former delivers strong results across panoptic, instance, and semantic segmentation, with masked attention and high-resolution features providing major ablation gains. Its performance generalizes across datasets, but task-specific training and small-object segmentation remain limitations.

  • Generalization: Mask2Former generalizes competitively to state-of-the-art performance across four additional datasets, indicating that its universal-model behavior extends beyond standard benchmarks.The experiments describe generalization across datasets rather than restricting evaluation to the standard benchmarks.
  • Panoptic segmentation: 57.8 PQ: Mask2Former with Swin-L sets a new COCO panoptic segmentation state of the art, exceeding prior MaskFormer by 5.1 PQ and K-Net by 3.2 PQ.It also outperforms MaskFormer by more than 5 PQ across different backbones while converging 6× faster.
  • Panoptic segmentation: Mask2Former trained only on panoptic annotations also achieves higher APTh_pan and mIoU_pan than DETR and MaskFormer, supporting its use for instance and semantic segmentation.APTh_pan evaluates the 80 thing categories, while mIoU_pan evaluates 133 categories converted from panoptic annotations.
  • Instance segmentation: 10.6 APL: Mask2Former’s largest instance-segmentation gain over MaskFormer with ResNet-50 occurs for large objects, while its 7.0 APS improvement still trails other state-of-the-art models.The authors identify small-object performance as an area for future improvement, including possible use of dilated backbones.
  • Ablations: Masked attention and high-resolution features produce the largest ablation gains, while optimization improvements increase performance without extra computation.Removing masked attention reduces the reported scores by 5.9, 4.8, and 1.7, and removing high-resolution features reduces them by 2.2, 1.7, and 1.1.
  • Limitations: Training only on panoptic annotations performs slightly worse than training with task-specific annotations across three datasets, and small-object segmentation remains difficult despite baseline improvements.The authors also report that Mask2Former cannot fully leverage multiscale features.

5. Conclusion · Appendix

Mask2Former is presented as a universal image-segmentation framework that achieves top results across panoptic, instance, and semantic tasks while reducing specialized-model research effort. The appendix extends evaluation across backbones, benchmarks, datasets, ablations, and prediction visualizations.

  • 5. Conclusion: Mask2Former addresses panoptic, instance, and semantic image segmentation within one universal framework.It is built upon a simple meta framework with a Transformer decoder using masked attention.
  • 5. Conclusion: Mask2Former obtains top results on four popular datasets across all three major image-segmentation tasks.The reported tasks are panoptic, instance, and semantic segmentation.
  • 5. Conclusion: Mask2Former outperforms the best specialized models designed for each benchmark while remaining easy to train.The comparison covers specialized models for the respective segmentation benchmarks.
  • 5. Conclusion: 3× research effort is saved compared to designing specialized models for each segmentation task.The conclusion frames this saving as a consequence of using Mask2Former rather than separate task-specific models.
  • 5. Conclusion: The conclusion emphasizes a single architecture that combines a simple meta framework with a masked-attention Transformer decoder.This architectural combination underlies Mask2Former’s universal segmentation formulation.
  • Appendix: Appendix A reports additional Mask2Former results with different backbones and test-set performance on standard benchmarks.The benchmarks include COCO panoptic for panoptic segmentation, COCO for instance segmentation, and ADE20K for semantic segmentation.
  • Appendix: Appendix B adds detailed results on additional datasets, while Appendix C supplies further ablations and Appendix D visualizes predictions for all three tasks.The appendix therefore expands dataset evaluation, component analysis, and qualitative inspection.

A. Additional results

Additional results evaluate Mask2Former across three segmentation benchmarks using multiple ResNet and Swin backbone variants initialized from ImageNet-pretrained checkpoints.

  • Evaluation setup: Mask2Former is evaluated on COCO panoptic, COCO instance, and ADE20K semantic segmentation benchmarks.The evaluations cover panoptic, instance, and semantic segmentation, respectively.
  • Backbones: The backbone study uses ResNet-50, ResNet-101, and Swin Tiny, Small, Base, and Large variants.These architectures are compared as alternative backbones for each benchmark.
  • Initialization: All backbones are initialized with ImageNet-pretrained checkpoints.ImageNet pre-training provides the initialization for the evaluated backbone variants.

A.1. Panoptic segmentation. · A.2. Instance segmentation. · A.3. Semantic segmentation.

Across panoptic, instance, and semantic segmentation, Mask2Former outperforms existing or previous state-of-the-art methods under the reported evaluation settings. It achieves 57.8 PQ on COCO panoptic validation and 57.7 mIoU on ADE20K validation, while showing particularly strong large-object instance segmentation performance.

  • A.1. Panoptic segmentation.: Mask2Former achieves new state-of-the-art performance on both COCO validation and test-dev using only standard train2017 data.It also outperforms the best COCO competition entry, which uses extra training data and test-time augmentation.
  • A.2. Instance segmentation.: On COCO val2017, Mask2Former outperforms the best single-scale instance segmentation model, HTC++ [6, 36].The comparison uses single-scale inference because multi-scale inference can require complex post-processing such as non-maximum suppression.
  • A.2. Instance segmentation.: Mask2Former achieves new state-of-the-art instance segmentation performance on both COCO validation and test-dev.The reported best model is evaluated on the test-dev set.
  • A.2. Instance segmentation.: Mask2Former outperforms the challenge winner on APL by a large margin without extra training data, model ensembles, or similar additions.Its poor performance on small objects remains an area for future improvement.
  • A.3. Semantic segmentation.: 57.7 mIoU: Mask2Former’s best model sets a new state-of-the-art on ADE20K val across evaluated backbones.It outperforms all existing semantic segmentation models with various backbones.
  • A.3. Semantic segmentation.: Mask2Former outperforms previous state-of-the-art semantic segmentation methods on all metrics in the ADE20K test-set evaluation.Following [14], training uses the union of ADE20K train and val, an ImageNet-22K pre-trained checkpoint, and multi-scale inference.

B. Additional datasets … C. Additional ablation studies

The supplementary section reports Mask2Former experiments on Cityscapes, ADE20K, and Mapillary Vistas, alongside additional training details and ablation studies. It specifies dataset-specific settings and evaluates panoptic, instance, and semantic segmentation where applicable.

  • B. Additional datasets: Additional experiments cover Cityscapes, ADE20K, and Mapillary Vistas across panoptic, instance, and semantic segmentation tasks.The section also provides more detailed training settings.
  • B.1. Cityscapes: Cityscapes contains 2,975 training, 500 validation, and 1,525 testing images across 19 classes at 1024 × 2048 resolution.All three tasks use 512 × 1024 crops, batch size 16, and 90k training iterations, with whole-image inference.
  • B.1. Cityscapes: Cityscapes training uses 512 × 1024 crops, batch size 16, and 90k iterations for all three segmentation tasks.Inference operates on the full 1024 × 2048 image.
  • B.2. ADE20K: ADE20K experiments use a 640 × 640 crop size for all backbones, with 200 queries for Swin-L panoptic and instance models and 100 queries otherwise.These settings apply to the reported panoptic, instance, and semantic segmentation experiments.
  • B.3. Mapillary Vistas: Mapillary Vistas provides 18k, 2k, and 5k images for training, validation, and testing, with resolutions ranging from 1024 × 768 to 4000 × 6000.The section reports panoptic and semantic segmentation results for this dataset.
  • B.3. Mapillary Vistas: Mapillary Vistas training runs for 300k iterations with batch size 16 and a poly learning-rate schedule, while inference resizes the longer side to 2048 pixels.The Swin-L panoptic model uses 200 queries; other backbones and semantic models use 100.
  • C. Additional ablation studies: Additional ablation studies use the main-paper settings with a single ResNet-50 backbone.The ablations are presented as supplementary studies of Mask2Former.

C.1. Convergence analysis

Mask2Former converges rapidly with both standard augmentation and large-scale jittering, reaching convergence far earlier than models using the standard Transformer decoder. With standard augmentation, it converges in 25 epochs, while large-scale jittering reaches near-convergence in 50 epochs.

  • C.1. Convergence analysis: Mask2Former is trained for 12, 25, 50, and 100 epochs using either standard scale augmentation or large-scale jittering augmentation.The experiments compare Standard Aug. with LSJ Aug.
  • C.1. Convergence analysis: 25 epochs marks convergence for Mask2Former with standard augmentation.This result is reported in Figure IV.
  • C.1. Convergence analysis: 50 epochs brings Mask2Former close to convergence with large-scale jittering augmentation.The model almost converges under LSJ augmentation.
  • C.1. Convergence analysis: 500 epochs for DETR and 300 epochs for MaskFormer contrast with Mask2Former's faster convergence using the proposed Transformer decoder.Both DETR and MaskFormer use the standard Transformer decoder.

C.2. Masked attention analysis

The COCO panoptic R50 model was analyzed quantitatively and qualitatively, including attention maps for a query predicting a cat. Cross-attention spread across the image, whereas masked attention constrained attention to predicted mask regions.

  • The analysis used the COCO panoptic model with an R50 backbone and examined its behavior both quantitatively and qualitatively.
  • For a query predicting the “cat,” cross-attention spread across the entire image, with its strongest response outside the object of interest.The authors attribute this behavior to softmax attention weights never reaching zero, allowing small weights over large background regions to dominate.
  • Masked attention instead limits attention to localized regions by constraining it within predicted mask regions.

C.3. Object query analysis · C.4. MaskFormer vs. Mask2Former

The analysis shows that object-query design should reflect segmentation-task demands, while directly supervised learnable queries improve Mask2Former. Compared with MaskFormer, Mask2Former benefits from broadly applicable training improvements, a stronger decoder, and a better compute-performance trade-off.

  • C.3. Object query analysis: For instance and semantic segmentation, 100 queries perform best, whereas 200 queries further improve panoptic segmentation.The authors relate this to panoptic images containing more segments because panoptic segmentation combines instance and semantic segmentation.
  • C.3. Object query analysis: Mask2Former makes both query positional embeddings and query features learnable, applying losses to query features before Transformer decoding.Query features are updated through decoder layers, while positional embeddings are added at every decoder layer when computing attention weights.
  • C.3. Object query analysis: Direct supervision of object queries before Transformer decoding is important for Mask2Former.Learnable queries without supervision perform similarly well as DETR’s zero-initialized queries.
  • C.4. MaskFormer vs. Mask2Former: Mask2Former uses a modified Transformer decoder with masked attention and optimization improvements, alongside more advanced training parameters.These differences are evaluated through ablations against MaskFormer’s original training parameters and standard Transformer decoder.
  • C.4. MaskFormer vs. Mask2Former: New training parameters significantly improve MaskFormer as well as Mask2Former.This indicates that the training parameters are generally applicable beyond Mask2Former.
  • C.4. MaskFormer vs. Mask2Former: The new Transformer decoder improves all three segmentation tasks over the standard decoder.MaskFormer and Mask2Former use the same ResNet-50 backbone, FPN pixel decoder, and training parameters in this comparison.
  • C.4. MaskFormer vs. Mask2Former: 1/4 the FLOPs: even the lightest Mask2Former instantiation outperforms the heaviest MaskFormer instantiation.The authors report this as evidence of a better compute-performance trade-off, although computational efficiency was not the primary goal.

D. Visualization

The section visualizes Mask2Former predictions across panoptic, instance, and semantic segmentation using a Swin-L backbone, reporting strong validation-set results. The figures juxtapose ground truth and predictions and include failure cases.

  • Panoptic segmentation: 57.8 PQ is achieved for panoptic segmentation on the COCO panoptic validation set with a Swin-L backbone.Figure V shows ground truth, predictions, and failure cases.
  • Instance segmentation: 50.1 AP is achieved for instance segmentation on the COCO validation set with a Swin-L backbone.Figure VI shows predictions with confidence scores greater than 0.5, alongside ground truth and failure cases.
  • Semantic segmentation: 57.7 mIoU is achieved for semantic segmentation on the ADE20K validation set with multi-scale inference and a Swin-L backbone.Figure VII shows ground truth, predictions, and failure cases.
Loading 2112.01527v3…