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

Segmentation tasks differ mainly in the semantics used to group pixels, yet research has developed specialized architectures for each task. Mask2Former uses masked attention within a universal architecture and performs on par with or better than specialized architectures across three tasks and four datasets.

  • Problem

    Although segmentation tasks differ mainly in grouping semantics, research uses task-specific architectures, while existing universal architectures have lower instance-segmentation performance and are harder to train.

    Source: §1 · pp. 1–2

  • Method

    Mask2Former is a universal segmentation architecture that uses masked attention to constrain decoder attention to localized features around predicted segments.

    Source: §1 · p. 2

  • Results

    Across three segmentation tasks and four datasets, one Mask2Former architecture performs on par with or better than specialized architectures, setting 57.8 PQ on COCO panoptic segmentation.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Mask2Former saves 3× research effort compared with designing specialized models for each segmentation task.

    Source: §5 · p. 8

    Despite generalizing across tasks, Mask2Former still needs task-specific training, as panoptic-only training performs slightly worse than training with corresponding instance and semantic annotations.

    Source: §4.6 · p. 8

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

Segmentation methods are specialized by task despite sharing the goal of grouping pixels, while earlier universal models lagged specialized methods in performance and training efficiency. Mask2Former proposes one universal architecture with localized masked attention and other design improvements, and reports strong results across three segmentation tasks.

1. Introduction

  • Although segmentation tasks differ primarily in grouping semantics, specialized FCN and mask-classification architectures dominate different tasks and lack flexibility to generalize across them.This specialization duplicates research and hardware-optimization effort across tasks.Source: §1 · p. 1
  • 57.8 PQ on COCO panoptic segmentation, 50.1 AP on COCO instance segmentation, and 57.7 mIoU on ADE20K semantic segmentation set state-of-the-art with one architecture.Across panoptic, instance, and semantic tasks on COCO, Cityscapes, ADE20K, and Mapillary Vistas, the same architecture performs on par or better than specialized architectures.Source: Figure 1 · p. 1, §1 · p. 2
  • Prior universal architectures lagged instance specialists by >9 AP; MaskFormer needed 300 epochs for 40.1 AP versus Swin-HTC++'s 72 epochs.Universal models share an architecture across tasks but are trained separately for each task and dataset; MaskFormer fit only one image on a 32G GPU, hampering deployment through performance and training-efficiency costs.Source: §1 · pp. 1–2
  • Mask2Former combines a backbone, pixel decoder, and Transformer decoder with masked attention over predicted segment regions to improve convergence and performance.Multi-scale high-resolution features support small objects and regions; optimization changes improve performance without extra compute, while randomly sampled mask-loss points save 3× training memory without affecting performance.Source: §1 · p. 2

2. Related Work

Prior segmentation work developed separate approaches for semantic, instance, and panoptic tasks, while emerging set-prediction architectures offered greater generality but did not consistently match specialized models. Mask2Former is presented as outperforming specialized architectures across all considered tasks and datasets.

Related Work

  • Specialized semantic models classify each pixel, while instance models predict class-labeled masks; their task-specific advances lack flexibility and duplicate research effort.Semantic methods add context modules or self-attention, while instance methods refine boxes or generate masks dynamically. BiFPN performs better for instance segmentation and FaPN for semantic segmentation.Source: §2 · p. 2
  • Panoptic architectures combine specialized semantic and instance models or use shared objectives, but panoptic performance alone does not guarantee generalization to other tasks.The paper therefore evaluates Mask2Former on all studied tasks rather than calling architectures evaluated only on panoptic segmentation universal.Source: §2 · p. 2
  • DETR, MaskFormer, and K-Net extend set prediction toward universal segmentation, but their performance on some tasks or datasets remains below specialized models.MaskFormer [14] performs well on panoptic and semantic segmentation but cannot segment instances well; K-Net [62] extends set prediction to instance segmentation.Source: §2 · p. 2
  • Mask2Former is presented as the first architecture to outperform state-of-the-art specialized architectures on all considered tasks and datasets.The comparison addresses the limitations of prior universal architectures, which did not replace specialized models across tasks or datasets.Source: §2 · p. 2

3. Masked-attention Mask Transformer

Mask2Former builds on a mask-classification architecture, replacing its standard Transformer decoder with masked attention and adding multi-scale features and training improvements. These components aim to improve segmentation while making training more accessible.

3.1. Mask classification preliminaries

  • Mask classification represents an image as N binary masks and corresponding category labels, using a backbone, pixel decoder, and Transformer decoder to produce masks.Different segment semantics allow the architecture to address different segmentation tasks; Mask2Former adopts this meta-architecture and replaces its standard decoder.Source: §3.1 · p. 3, §3.2 · p. 3

3.2.1 Masked attention

  • Masked attention constrains each query’s cross-attention to its predicted foreground mask, using localized features to update queries while self-attention gathers context.The attention mask is a resized, binarized prediction from the previous decoder layer, thresholded at 0.5; the initial mask comes from the input queries.Source: §3.2 · p. 3, §3.2.1 · pp. 3–4

3.2.2 High-resolution features

  • A round-robin feature pyramid feeds resolutions 1/32, 1/16, and 1/8 of the original image to successive decoder layers, adding high-resolution features for small objects.The three-layer sequence is repeated L times, giving the final decoder 3L layers.Source: §3.2.2 · p. 4

3.2.3 Optimization improvements

  • The decoder switches self- and cross-attention order, learns and supervises initial query features as mask proposals, and removes dropout.The authors report that dropout is unnecessary and usually decreases performance; learnable query features function like a region proposal network [43].Source: Figure 2 · p. 3, §3.2.3 · p. 4

3.3. Improving training efficiency

  • Sampling K = 12544 points for mask losses reduces training memory 3×, from 18GB to 6GB per image.The matching loss uses the same uniformly sampled points for all prediction and ground-truth masks; the final loss uses importance sampling with different points for each matched pair.Source: §3.3 · p. 4

4. Experiments

Experiments compare Mask2Former with specialized methods, evaluate design choices through ablations on three segmentation tasks, and test generalization beyond standard benchmarks. The reported results cover panoptic, instance, and semantic segmentation, alongside ablations, implementation and training details, cross-dataset results, and limitations.

4.3. Main results

  • 57.8 PQ with Swin-L sets a new COCO panoptic state of the art, beating MaskFormer [14] by 5.1 PQ and K-Net [62] by 3.2 PQ.Mask2Former also beats MaskFormer by more than 5 PQ across backbones while converging 6× faster; the best ensemble models with extra training data are also surpassed.Source: Table 1 · p. 5, §4.3 · p. 6

4.1. Implementation details; 4.2. Training settings

  • Mask2Former uses a 3-stage, 9-layer Transformer decoder with 100 queries and a 6-layer MSDeformAttn pixel decoder by default.The implementation uses binary cross-entropy plus dice mask loss, with weights 5.0 each, and classification-loss weights 2.0 for matched predictions and 0.1 for unmatched predictions. COCO panoptic and instance training uses AdamW, 50 epochs, batch size 16, and large-scale jittering unless otherwise stated.Source: §4.1 · p. 5, §4.2 · p. 5

4.3. Main results

  • +0.6 AP and +2.1 Boundary AP over HTC++ [6] with Swin-L establish Mask2Former’s instance-segmentation gains on COCO.With ResNet, it also outperforms a strong Mask R-CNN baseline using 8× fewer training iterations; comparisons use single-scale inference and COCO train2017 only.Source: Table 2 · p. 6, §4.3 · p. 6
  • 57.7 mIoU with Swin-L and FaPN sets a new ADE20K semantic-segmentation state of the art, improving on MaskFormer [14] across backbones.Mask2Former also improves small-object APS over MaskFormer by 7.0 with ResNet-50, but its small-object performance still trails other state-of-the-art models.Source: Table 3 · p. 6, §4.3 · p. 6

4.4. Ablation studies

  • Ablations across three tasks find masked attention and high-resolution features yield the most gains, while MSDeformAttn [66] performs best across tasks.The efficient multi-scale strategy reduces FLOPs without affecting performance; point-based final loss cuts training memory by 3× without performance loss, and point-based matching improves performance across all three tasks.Source: §4.4 · pp. 6–7

4.5. Generalization to other datasets

  • Mask2Former is competitive with state-of-the-art methods on Cityscapes and other popular segmentation datasets, with results generalizing beyond COCO.The experiments include Cityscapes, ADE20K, and Mapillary Vistas beyond COCO; the supplied passages do not give specific results for those additional datasets.Source: §4.5 · p. 8

4.6. Limitations

  • A panoptic-trained Mask2Former performs slightly worse than the same model trained with task-specific annotations for instance and semantic segmentation across three datasets.Thus, it can generalize across tasks but still needs task-specific training; the authors identify single training for multiple tasks and datasets as a future goal.Source: Table 7 · p. 8, §4.6 · p. 8

5. Conclusion · Appendix

Mask2Former is presented as a universal segmentation model that outperforms specialized models across three major tasks and four datasets. The appendices extend the report with benchmark, dataset, ablation, and visualization results.

5. Conclusion

  • Mask2Former achieves top results across panoptic, instance, and semantic segmentation on four datasets, outperforming the best specialized models while remaining easy to train.The conclusion presents it as a universal model for the three major image-segmentation tasks.Source: §5 · p. 8
  • The model combines a simple meta framework [14] with a new Transformer decoder that uses masked attention.Source: §5 · p. 8
  • Mask2Former saves 3× the research effort required to design specialized models for each task and is accessible to users with limited computational resources.Source: §5 · p. 8
  • The authors caution that models trained on real-world datasets should undergo ethical review and that the approach should not be used for applications such as illegal surveillance.They report no apparent inherent biases in the technical innovations, while warning against potentially problematic predictions and uses.Source: §5 · p. 8

Appendix

  • The conclusion reports broad segmentation results, while the appendices add backbone and test-set results, additional datasets, ablations, and prediction visualizations.Source: §5 · p. 8, Appendix · p. 11
  • Appendix A covers results with different backbones and test-set performance on COCO panoptic, COCO instance, and ADE20K semantic segmentation benchmarks.Source: Appendix · p. 11
  • Appendices B–D provide results on additional datasets, further ablation studies, and Mask2Former prediction visualizations for all three segmentation tasks.Source: Appendix · p. 11

A. Additional results

The section reports additional Mask2Former evaluations across panoptic, instance, and semantic segmentation benchmarks using ResNet and Swin backbones initialized from ImageNet-pretrained checkpoints.

A. Additional results

  • Mask2Former was evaluated on COCO panoptic, COCO instance, and ADE20K semantic segmentation with ResNet-50/101 and Swin Tiny, Small, Base, and Large backbones initialized from ImageNet-pretrained checkpoints.Source: A. Additional results · p. 11

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

Across panoptic, instance, and semantic segmentation, Mask2Former reports state-of-the-art results on COCO and ADE20K validation and test sets. Its reported validation records include 57.8 PQ for panoptic segmentation and 57.7 mIoU for semantic segmentation.

A.1. Panoptic segmentation.

  • 57.8 PQ is Mask2Former’s best COCO panoptic val2017 result, surpassing all existing panoptic models across various backbones.Source: §A.1 · p. 11
  • With only standard train2017 data, Mask2Former achieves new COCO panoptic state of the art on validation and test, exceeding a competition entry using extra data and test-time augmentation.Source: §A.1 · p. 11

A.2. Instance segmentation.

  • Mask2Former outperforms HTC++ [6, 36], the best single-scale instance model; comparisons use single-scale inference because multi-scale instance inference complicates post-processing.Source: §A.2 · p. 11
  • On COCO instance segmentation, Mask2Former achieves state of the art on validation and test, surpasses the challenge winner on APL without extra techniques, but remains weaker on small objects.Source: §A.2 · p. 11

A.3. Semantic segmentation.

  • 57.7 mIoU is Mask2Former’s best ADE20K val result, surpassing all existing semantic segmentation models across various backbones.Source: §A.3 · p. 11
  • On ADE20K test, Mask2Former surpasses previous state-of-the-art methods on all metrics when trained on train and val with ImageNet-22K pretraining and multi-scale inference, following [14].Source: §A.3 · p. 11

B. Additional datasets … C. Additional ablation studies

Mask2Former is evaluated on additional datasets for multiple segmentation tasks, with results reported for Cityscapes, ADE20K, and Mapillary Vistas. The section also introduces further ablation studies using the main paper’s ResNet-50 settings.

B. Additional datasets

  • The additional-dataset section reports results on Cityscapes, ADE20K, and Mapillary Vistas, alongside more detailed training settings.The paper studies panoptic, instance, and semantic segmentation across four datasets.Source: B. Additional datasets · p. 11

B.1. Cityscapes

  • On Cityscapes, Swin-L Mask2Former outperforms Panoptic-DeepLab with SWideRnet for panoptic segmentation using single-scale inference, while Swin-B outperforms SegFormer for semantic segmentation.All three tasks use 512 × 1024 crops, batch size 16, and 90k iterations; inference uses the whole 1024 × 2048 image.Source: §B.1 · p. 11

B.2. ADE20K

  • Mask2Former with Swin-L sets a new state of the art for panoptic segmentation on ADE20K.For panoptic and instance segmentation, training uses a 640 × 640 crop; Swin-L models use 200 queries and other backbones use 100.Source: §B.2 · p. 12

B.3. Mapillary Vistas

  • Mapillary Vistas results cover panoptic and semantic segmentation on a large-scale street-view dataset with 18k training, 2k validation, and 5k test images.Training uses scale jittering from 0.5 to 2.0, horizontal flipping, 1024 × 1024 random crops, and color jittering.Source: §B.3 · p. 12
  • On Mapillary Vistas, Mask2Former is very competitive with state-of-the-art specialized models for panoptic and semantic segmentation, despite not being designed for the dataset.Training runs for 300k iterations with batch size 16; inference resizes the longer side to 2048 pixels, and the Swin-L panoptic model uses 200 queries.Source: §B.3 · p. 13

C. Additional ablation studies

C.1. Convergence analysis

Mask2Former’s Transformer decoder converges substantially faster than standard Transformer decoders under standard and large-scale jittering augmentation. It converges in 25 epochs with standard augmentation and almost converges in 50 epochs with LSJ, compared with 500 epochs for DETR and 300 for MaskFormer.

C.1. Convergence analysis

  • Mask2Former converges in 25 epochs with standard augmentation and almost converges in 50 epochs with LSJ, versus 500 epochs for DETR and 300 for MaskFormer.The faster convergence is attributed to Mask2Former’s proposed Transformer decoder compared with the standard Transformer decoder.Source: §C.1 · p. 13
  • Convergence was evaluated after 12, 25, 50, and 100 epochs using either standard scale augmentation or large-scale jittering augmentation.Source: §C.1 · p. 13

C.2. Masked attention analysis

The section analyzes masked attention in a COCO panoptic model, finding that it focuses attention on the predicted object rather than background regions. It also reports that masked attention performs strongly with fewer decoder layers than cross-attention.

C.2. Masked attention analysis

  • 20% of cross-attention weights focus on foreground on average, compared with almost 60% for masked attention across COCO val queries.Foreground is defined by the matching ground truth to each prediction.Source: §C.2 · p. 14
  • Masked attention with one Transformer decoder layer already outperforms cross-attention with nine layers on panoptic segmentation.The comparison plots performance using output from each Transformer decoder layer.Source: §C.2 · p. 14, Figure II · p. 14
  • Cross-attention spreads a cat-predicting query’s attention across the image, while masked attention focuses it on the object of interest.The authors suggest that cross-attention's softmax never reaches zero, allowing small weights over large background regions to dominate.Source: §C.2 · p. 14

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

The analyses show that effective query choices depend on the segmentation task and that direct query supervision matters. Mask2Former’s training settings and masked-attention decoder improve on MaskFormer, while Mask2Former also offers a better compute-performance trade-off.

C.3. Object query analysis

  • 100 queries give the best instance and semantic segmentation performance, while 200 queries further improve panoptic segmentation, suggesting query counts should reflect segments per image.Source: §C.3 · p. 15
  • Mask2Former uses learnable query positional embeddings and learnable query features, applying losses to the features before they enter the decoder.Query features are the decoder’s initial input and are updated through its layers; positional embeddings are added at each layer when computing attention weights.Source: §C.3 · p. 15
  • Directly supervising object queries before the decoder is important; without supervision, learnable query features perform similarly to DETR’s zero-initialized features.Source: §C.3 · p. 15

C.4. MaskFormer vs. Mask2Former

  • New training parameters significantly improve MaskFormer as well, showing that these settings are generally applicable to other models.Source: §C.4 · p. 16
  • With the same ResNet-50 backbone, FPN pixel decoder, and training parameters, Mask2Former’s masked-attention decoder improves all three segmentation tasks over MaskFormer’s standard decoder.This isolates the decoder as the only difference between the two models in the comparison.Source: §C.4 · pp. 15–16
  • 1/4 the FLOPs let the lightest Mask2Former outperform the heaviest MaskFormer, giving Mask2Former a better compute-performance trade-off.Computational efficiency was not the primary goal of the comparison.Source: §C.4 · p. 16, Figure III · p. 16

D. Visualization

The section visualizes Mask2Former with a Swin-L backbone on panoptic, instance, and semantic segmentation, showing predictions alongside ground truth and including failure cases.

D. Visualization

  • 57.8 PQ: Figure V shows COCO panoptic validation predictions from Mask2Former with Swin-L; the last row shows failure cases.The first and third columns are ground truth, and the second and fourth are predictions.Source: D. Visualization · p. 16, Figure V · p. 18
  • 50.1 AP: Figure VI shows COCO instance segmentation predictions from Mask2Former with Swin-L, using confidence scores above 0.5.The first and third columns are ground truth, the second and fourth are predictions, and the last row shows failure cases.Source: D. Visualization · p. 16, Figure VI · p. 19
  • 57.7 mIoU: Figure VII shows multi-scale ADE20K semantic segmentation predictions from Mask2Former with Swin-L; the last row shows failure cases.The first and third columns are ground truth, and the second and fourth are predictions.Source: D. Visualization · p. 16, Figure VII · p. 20

Related papers

Loading 2112.01527v3…