PaperlayerView exact revision on arXiv

Source-linked AI summary

MaX-DeepLab: End-to-End Panoptic Segmentation with Mask Transformers

Huiyu Wang, Yukun Zhu, Hartwig Adam, Alan Yuille, Liang-Chieh Chen

arXiv:2012.00759v3cs.CV

TL;DR

Panoptic segmentation has traditionally relied on surrogate subtasks and hand-designed components rather than directly modeling class-labeled masks. MaX-DeepLab predicts masks and classes with a mask transformer, trains them using PQ-style bipartite matching, and achieves state-of-the-art COCO performance while closing the box-free and box-based gap.

  • Problem

    Existing panoptic-segmentation systems approximate mask prediction with box-, center-, and other surrogate subtasks plus hand-designed pipeline components.

  • Method

    MaX-DeepLab directly predicts class-labeled masks with a mask transformer, using a PQ-style loss with bipartite matching and a dual-path CNN/global-memory architecture.

  • Results

    MaX-DeepLab reaches 51.3% PQ on COCO test-dev without test time augmentation and closes the gap between box-free and box-based methods.

  • Takeaways & Limitations

    End-to-end mask-set prediction can remove hand-designed box and center priors while achieving state-of-the-art panoptic segmentation on COCO.

Abstract

from arXiv · show

We present MaX-DeepLab, the first end-to-end model for panoptic segmentation. Our approach simplifies the current pipeline that depends heavily on surrogate sub-tasks and hand-designed components, such as box detection, non-maximum suppression, thing-stuff merging, etc. Although these sub-tasks are tackled by area experts, they fail to comprehensively solve the target task. By contrast, our MaX-DeepLab directly predicts class-labeled masks with a mask transformer, and is trained with a panoptic quality inspired loss via bipartite matching. Our mask transformer employs a dual-path architecture that introduces a global memory path in addition to a CNN path, allowing direct communication with any CNN layers. As a result, MaX-DeepLab shows a significant 7.1% PQ gain in the box-free regime on the challenging COCO dataset, closing the gap between box-based and box-free methods for the first time. A small variant of MaX-DeepLab improves 3.0% PQ over DETR with similar parameters and M-Adds. Furthermore, MaX-DeepLab, without test time augmentation, achieves new state-of-the-art 51.3% PQ on COCO test-dev set. Code is available at https://github.com/google-research/deeplab2.

1. Introduction

MaX-DeepLab replaces surrogate-task panoptic pipelines with direct mask and class prediction trained end-to-end. Its PQ-style matching objective and dual-path transformer support state-of-the-art COCO performance without test time augmentation.

  • Motivation: Panoptic segmentation predicts non-overlapping, class-labeled masks, but existing methods approximate this target through manually designed surrogate-task pipelines.These pipelines may include box detection, anchor operations, non-maximum suppression, and thing-stuff merging.
  • Motivation: Box-free methods based on instance-center offsets struggle with deformable objects and nearby objects having close centers, limiting COCO performance.The paper contrasts these failure modes with direct mask prediction in its case study.
  • Approach: MaX-DeepLab directly predicts non-overlapping masks and semantic labels using a mask transformer, removing hand-coded box- and center-based priors.The model is trained and inferred end-to-end rather than assembling separate surrogate subtasks.
  • Approach: The model optimizes a PQ-style objective through one-to-one bipartite matching between predicted and ground-truth class-labeled masks.The objective combines mask similarity and class similarity.
  • Architecture: The dual-path transformer lets CNN layers read and write global memory at any layer through pixel-path and memory-path communication.Its decoder aggregates multi-scale features into high-resolution masks, while another branch predicts mask classes.
  • Results: 51.3% PQ on COCO test-dev establishes a new state-of-the-art without test time augmentation, while MaX-DeepLab closes the box-free versus box-based gap.It exceeds Axial-DeepLab with TTA by 7.1% PQ and DetectoRS with TTA by 1.7% PQ.

2. Related Work

Related panoptic-segmentation methods either use boxes and merging modules or remove boxes through pixel grouping and center-based surrogates. MaX-DeepLab extends this line of work with direct mask-set prediction using a mask transformer.

  • Transformers: Transformers use attention to capture long-range interactions, with global-memory variants reducing the cost of full attention while propagating information globally.These ideas motivate the global-memory path in MaX-DeepLab.
  • Transformers: Vision research has combined attention or transformers with CNNs across classification, detection, segmentation, video recognition, generation, and panoptic segmentation.DETR is specifically cited as stacking a transformer on top of a CNN for end-to-end object detection.
  • Box-based panoptic segmentation: Box-based methods detect object bounding boxes, predict a mask for each box, and merge thing and stuff segments into panoptic output.Panoptic FPN and Mask R-CNN with FPN exemplify this pipeline.
  • Box-free panoptic segmentation: Box-free methods typically begin with semantic segments and obtain instances by grouping thing pixels using centers, watershed, Hough voting, or pixel affinity.Axial-DeepLab advances this regime with a fully axial-attention backbone.

3. Method

MaX-DeepLab directly predicts class-labeled masks for panoptic segmentation and trains them with a PQ-style similarity, matching, and loss framework. Its dual-path transformer combines CNN features with global memory for mask-set prediction and auxiliary instance-discriminative learning.

  • 3.1. MaX-DeepLab formulation: MaX-DeepLab predicts a fixed-size set of class-labeled masks directly, unifying thing and stuff classes with an ∅ no-object class.This formulation avoids anchors, boxes, instance centers, and separate thing-stuff merging.
  • 3.1. MaX-DeepLab formulation: Inference assigns each mask a class and each pixel a mask ID using two argmax operations, followed by confidence filtering.The resulting panoptic segmentation is obtained without common post-processing such as NMS and thing-stuff merging.
  • 3.2. PQ-style loss: The PQ-style objective multiplies class correctness and Dice mask similarity, then optimizes matched positive masks alongside cross-entropy for unmatched ∅ masks.The reformulation yields Dice and cross-entropy terms weighted by class and mask correctness, respectively.
  • 3.2. PQ-style loss: One-to-one bipartite matching maximizes total mask similarity, with the Hungarian algorithm assigning K positives and training the remaining N −K predictions as no-object negatives.The matching prevents multiple predictions from being optimized against the same ground-truth mask.
  • 3.3. MaX-DeepLab Architecture: The dual-path transformer adds a 1D global memory to a 2D CNN path and supports bidirectional M2P, M2M, P2M, and P2P attention.It can be inserted at any backbone layer or resolution, while P2M feedback lets CNN features refine using mask information.
  • 3.3. MaX-DeepLab Architecture: A stacked decoder fuses multi-scale features for high-resolution mask prediction, while an auxiliary per-pixel contrastive loss encourages same-instance similarity and different-instance distinction.The decoder output is multiplied with global memory features to form the mask set, and the auxiliary loss is applied to instance pixels.

4. Experiments

Experiments on COCO evaluate MaX-DeepLab’s main configurations, scaling choices, transformer components, losses, and training behavior. The model achieves strong box-free and overall panoptic segmentation results while showing which architectural and auxiliary-loss choices matter.

  • Experimental setup: MaX-DeepLab is evaluated with small and large variants, including a DETR-matched small model and a large model intended for state-of-the-art comparison.MaX-DeepLab-S matches DETR-R101 in parameters and M-Adds; MaX-DeepLab-L uses a stacked L = 2 decoder and wide axial-attention blocks.
  • Main results: MaX-DeepLab-S outperforms all other box-free methods by more than 4.5% PQ on COCO validation, with or without TTA.It surpasses single-scale Panoptic-DeepLab by 8.7% PQ without TTA and by 7.1% PQ with TTA.
  • Scaling: 49.4% PQ is reached after increasing input resolution to 1025 × 1025 with a short 100k-iteration schedule, while stacking the decoder once adds 1.4% PQ.Replacing convolutional blocks with axial-attention blocks gives the largest scaling improvement, and further decoder stacking begins to saturate.
  • Dual-path transformer: Removing P2M feedback attention drops PQ by 0.7%, whereas removing M2M self-attention reduces it by 0.6%; applying transformer blocks at stride 8 adds 1.0% PQ.The model can use stride-8 transformer blocks through its dual-path design, unlike DETR in the comparison described.
  • Loss ablation: Combining instance-discrimination and Mask-ID auxiliary losses produces a 5.6% PQ gain, while semantic segmentation adds a further 0.6% PQ.The PQ-style multiplicative similarity improves PQ by 0.8% over the additive alternative.
  • Analysis: MaX-DeepLab converges to around 46% PQ within 100k iterations, while only 10% of pixels predict wrong mask-IDs and contribute most of the PQ error.Matched masks show better class correctness than mask correctness, and visualized features assign similar colors within masks and different colors across masks.

5. Conclusion

MaX-DeepLab demonstrates end-to-end panoptic segmentation by directly predicting masks and classes with a mask transformer, removing many hand-designed priors. With PQ-style loss and a dual-path transformer, it achieves state-of-the-art results on COCO and closes the gap between box-based and box-free methods.

  • MaX-DeepLab directly predicts masks and classes with a mask transformer, removing hand-designed priors such as bounding boxes and thing-stuff merging.
  • PQ-style loss and a dual-path transformer equip MaX-DeepLab for end-to-end panoptic segmentation.
  • MaX-DeepLab achieves state-of-the-art results on the challenging COCO dataset and closes the gap between box-based and box-free methods for the first time.

A.1. Panoptic Segmentation Results

MaX-DeepLab shows robust panoptic segmentation behavior on challenging examples involving similar boxes, nearby objects, and many small objects, while retaining a specific failure mode with excessive masks.

  • MaX-DeepLab is robust to similar object bounding boxes and nearby objects with close centers, where other methods make systematic surrogate-task mistakes.
  • MaX-DeepLab shows exceptional mask quality and performs well on images containing many small objects.
  • MaX-DeepLab typically fails when there are too many object masks, similar to DETR.

A.2. Runtime

Runtime measurements on a V100 GPU use single-scale, non-flipped inputs and the built-in TensorFlow library without extra inference optimization. In the fast regime, MaX-DeepLab-S processes a typical 641×641 input in 67 ms.

  • 67 ms is the reported runtime for MaX-DeepLab-S with a typical 641×641 input in the fast regime.The measurement includes 5 ms of postprocessing and 15 ms of batch normalization that can be optimized.
  • Runtime measurements use single-scale input without flipping and built-in TensorFlow without extra inference optimization.

A.3. Mask Output Slot Analysis

Analysis of MaX-DeepLab-L's 128 mask prediction slots finds specialization by class and image region, alongside imbalanced prediction frequencies. Unlike DETR, the slots do not nearly all favor large image-wide masks.

  • N = 128 mask prediction slots in MaX-DeepLab-L have imbalanced prediction frequencies and specialize in thing and stuff classes.
  • Mask slots also specialize in particular input-image regions based on their average predicted masks.
  • Unlike DETR, MaX-DeepLab does not show a pattern in which almost all slots favor large image-wide masks.

A.4. Mask Head Visualization

With only D = 3 decoder feature channels, MaX-DeepLab’s decoder features can be visualized as RGB colors, revealing clustering among pixels belonging to the same instance.

  • D = 3 reduces performance from 45.7% PQ to 37.8% PQ but enables direct visualization of decoder features as RGB colors.The model is trained with three decoder feature channels specifically for visualization.
  • Pixels from the same instance have similar feature colors, while pixels from different instances have distinct colors.This clustering effect supports visualization of instance structure in the decoder features.

A.5. Transformer Attention Visualization

The appendix visualizes mask-to-pixel attention and compares MaX-DeepLab with representative methods on COCO, while noting its performance with many small instances and defining runtime metrics.

  • Transformer Attention Visualization: M2P attention is visualized by probing four selected output masks against all pixels in the last dual-path transformer block.The attention maps are colorized and displayed for two COCO validation images, covering eight maps in total.
  • Transformer Attention Visualization: MaX-DeepLab performs well in the presence of many small instances.
  • Transformer Attention Visualization: Figure A.1 compares MaX-DeepLab with other representative methods on the COCO validation set.
  • Transformer Attention Visualization: Table A.1 reports end-to-end runtime alongside PQ (%) on the COCO validation and test-dev sets.PQ [val] denotes COCO val PQ, while PQ [test] denotes COCO test-dev PQ.

A.6. More Technical Details

The appendix provides architectural details, visualizes mask-slot behavior and attention, and documents the D = 3 feature-color limitation that produces boundary artifacts.

  • More Technical Details: ImageNet pretraining uses only P2P attention, followed by average pooling and a linear layer, to produce the backbone for panoptic segmentation.The pretraining model omits the other attention types, FFN, and memory.
  • More Technical Details: Mask slots specialize by firing frequency and spatial region, with some slots favoring thing classes, stuff classes, or both.The appendix illustrates these behaviors through joint slot-class distributions and average slot masks.
  • More Technical Details: D = 3 provides only 37.8% PQ and can cause missing objects or boundary artifacts because there are insufficient colors for all masks.These artifacts are reported for the visualization-only extreme setting and not for the normal D = 1 setting.
  • More Technical Details: The M2P attention visualizations show some heads focusing on predicted mask regions and others gathering broader context from similar instances or boundaries.One head largely avoids the pixel path except for minor activations on selected regions.
  • More Technical Details: An Axial-Block applies two axial-attention layers sequentially along the height and width axes.
  • More Technical Details: The appendix diagrams the building blocks and detailed MaX-DeepLab architectures, including a dual-path transformer block and variants matched to DETR-R101-Panoptic.The figures also mark ImageNet pretraining and identify the borrowed Axial-Block.
Loading 2012.00759v3…