[Paper Review] MaX-DeepLab: End-to-End Panoptic Segmentation with Mask Transformers

Huiyu Wang, Yukun Zhu+3 more|arXiv (Cornell University)|Dec 1, 2020

Advanced Neural Network Applications98 references32 citations

TL;DR

MaX-DeepLab is the first end-to-end panoptic segmentation model that directly predicts class-labeled masks with a mask transformer and a dual-path CNN+memory architecture, achieving state-of-the-art PQ on COCO without test-time augmentation.

ABSTRACT

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.

Motivation & Objective

• Motivate an end-to-end panoptic segmentation approach that avoids hand-crafted priors like anchors, boxes, and NMS.
• Directly predict a set of class-labeled masks and optimize with a PQ-inspired bipartite matching loss.
• Introduce a dual-path transformer that enables global memory communication with CNN layers.
• Show that end-to-end mask prediction closes the gap between box-based and box-free methods on COCO.

Proposed method

• Directly predict a fixed-size set of class-labeled masks and their probabilities for each image.
• Train with a PQ-style loss based on a one-to-one bipartite matching between predicted and ground-truth masks.
• Introduce a dual-path transformer that couples a CNN pixel-path with a 1D global memory path enabling four attention modes (M2P, M2M, P2M, P2P).
• Use a stacked-hourglass-style decoder to fuse multi-scale features and form mask predictions.
• Auxiliary losses include per-pixel instance discrimination, mask-ID cross-entropy, and semantic segmentation losses to aid training.

Experimental results

Research questions

• RQ1 Can an end-to-end model predict panoptic masks directly without surrogate subtasks like boxes or centers?
• RQ2 Does optimizing a PQ-style objective via bipartite matching improve panoptic segmentation performance compared to box-based or box-free approaches?
• RQ3 Does the dual-path transformer effectively integrate CNN features with a global memory to improve mask prediction quality?
• RQ4 What is the impact of decoder depth, input resolution, and attention types on PQ and related metrics on COCO?
• RQ5 Do auxiliary losses (instance discrimination, mask-ID cross-entropy, semantic segmentation) provide measurable gains?

Key findings

• MaX-DeepLab achieves state-of-the-art PQ on COCO test-dev without test-time augmentation (51.3% PQ for the large model).
• In the box-free regime, MaX-DeepLab surpasses Axial-DeepLab by 7.1% PQ and outperforms box-based DetectoRS under strong conditions, bridging the gap between box-based and box-free methods.
• A lightweight variant (MaX-DeepLab-S) outperforms DETR with a similar parameter count by 3.3% PQ on val and 3.0% PQ on test-dev.
• A dual-path transformer enables formation and communication between CNN features and a global memory at multiple resolutions, improving mask predictions.
• Auxiliary losses, especially per-pixel instance discrimination and mask-ID prediction, yield significant PQ gains when combined with the PQ-style loss.