PolyMaX: General Dense Prediction with Mask Transformer
Authors: name: Xuan Yang; name: Liangzhe Yuan; name: Kimberly Wilber; name: Astuti Sharma; name: Xiuye Gu; name: Siyuan Qiao; name: Stephanie Debats; name: Huisheng Wang; name: Hartwig Adam; name: Mikhail Sirotenko; name: Liang-Chieh Chen
Published: 2023-11-09 22:29:21

Abstract: Dense prediction tasks, such as semantic segmentation, depth estimation, and
surface normal prediction, can be easily formulated as per-pixel classification
(discrete outputs) or regression (continuous outputs). This per-pixel
prediction paradigm has remained popular due to the prevalence of fully
convolutional networks. However, on the recent frontier of segmentation task,
the community has been witnessing a shift of paradigm from per-pixel prediction
to cluster-prediction with the emergence of transformer architectures,
particularly the mask transformers, which directly predicts a label for a mask
instead of a pixel. Despite this shift, methods based on the per-pixel
prediction paradigm still dominate the benchmarks on the other dense prediction
tasks that require continuous outputs, such as depth estimation and surface
normal prediction. Motivated by the success of DORN and AdaBins in depth
estimation, achieved by discretizing the continuous output space, we propose to
generalize the cluster-prediction based method to general dense prediction
tasks. This allows us to unify dense prediction tasks with the mask transformer
framework. Remarkably, the resulting model PolyMaX demonstrates
state-of-the-art performance on three benchmarks of NYUD-v2 dataset. We hope
our simple yet effective design can inspire more research on exploiting mask
transformers for more dense prediction tasks. Code and model will be made
available.