# PolyMaX: General Dense Prediction with Mask Transformer

arXiv: https://arxiv.org/abs/2311.05770 · HTML: https://mlsanity.com/paper/2311.05770/

## Key facts

- Cited by 14 papers in the map (data through 14 Oct 2026).
- Builds on 98 of its 120 references that are in the map.
- First posted on arXiv 9 Nov 2023.

## Builds on (most cited first, 30 of 98)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Fully Convolutional Networks for Semantic Segmentation](https://mlsanity.com/paper/1411.4038/) (2014)
- [The Cityscapes Dataset for Semantic Urban Scene Understanding](https://mlsanity.com/paper/1604.01685/) (2016)
- [LAION-5B: An open large-scale dataset for training next generation image-text models](https://mlsanity.com/paper/2210.08402/) (2022)
- [A ConvNet for the 2020s](https://mlsanity.com/paper/2201.03545/) (2022)
- [Deformable DETR: Deformable Transformers for End-to-End Object Detection](https://mlsanity.com/paper/2010.04159/) (2020)
- [DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs](https://mlsanity.com/paper/1606.00915/) (2016)
- [Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation](https://mlsanity.com/paper/1802.02611/) (2018)
- [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers](https://mlsanity.com/paper/2105.15203/) (2021)
- [Aggregated Residual Transformations for Deep Neural Networks](https://mlsanity.com/paper/1611.05431/) (2016)
- [Pyramid Scene Parsing Network](https://mlsanity.com/paper/1612.01105/) (2016)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [Rethinking Atrous Convolution for Semantic Image Segmentation](https://mlsanity.com/paper/1706.05587/) (2017)
- [The Kinetics Human Action Video Dataset](https://mlsanity.com/paper/1705.06950/) (2017)
- [Vision Transformers for Dense Prediction](https://mlsanity.com/paper/2103.13413/) (2021)
- [Towards Robust Monocular Depth Estimation: Mixing Datasets for Zero-shot Cross-dataset Transfer](https://mlsanity.com/paper/1907.01341/) (2019)
- [SegNet: A Deep Convolutional Encoder-Decoder Architecture for Image Segmentation](https://mlsanity.com/paper/1511.00561/) (2015)
- [Deep High-Resolution Representation Learning for Visual Recognition](https://mlsanity.com/paper/1908.07919/) (2019)
- [Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers](https://mlsanity.com/paper/2012.15840/) (2020)

## Built on by (most cited first, 14 of 14)

- [Metric3Dv2: A Versatile Monocular Geometric Foundation Model for Zero-shot Metric Depth and Surface Normal Estimation](https://mlsanity.com/paper/2404.15506/) (2024)
- [GeminiFusion: Efficient Pixel-wise Multimodal Fusion for Vision Transformer](https://mlsanity.com/paper/2406.01210/) (2024)
- [ViTamin: Designing Scalable Vision Models in the Vision-Language Era](https://mlsanity.com/paper/2404.02132/) (2024)
- [StitchFusion: Weaving Any Visual Modalities to Enhance Multimodal Semantic Segmentation](https://mlsanity.com/paper/2408.01343/) (2024)
- [COCONut: Modernizing COCO Segmentation](https://mlsanity.com/paper/2404.08639/) (2024)
- [EVP: Enhanced Visual Perception using Inverse Multi-Attentive Feature Refinement and Regularized Image-Text Alignment](https://mlsanity.com/paper/2312.08548/) (2023)
- [HAPNet: Toward Superior RGB-Thermal Scene Parsing via Hybrid, Asymmetric, and Progressive Heterogeneous Feature Fusion](https://mlsanity.com/paper/2404.03527/) (2024)
- [Towards In-the-wild 3D Plane Reconstruction from a Single Image](https://mlsanity.com/paper/2506.02493/) (2025)
- [Enhancing Monocular Depth Estimation with Multi-Source Auxiliary Tasks](https://mlsanity.com/paper/2501.12824/) (2025)
- [PDDM: Pseudo Depth Diffusion Model for RGB-PD Semantic Segmentation Based in Complex Indoor Scenes](https://mlsanity.com/paper/2503.18393/) (2025)
- [Neural USD: An object-centric framework for iterative editing and control](https://mlsanity.com/paper/2510.23956/) (2025)
- [DiffPixelFormer: Differential Pixel-Aware Transformer for RGB-D Indoor Scene Segmentation](https://mlsanity.com/paper/2511.13047/) (2025)
- [CrossWeaver: Cross-modal Weaving for Arbitrary-Modality Semantic Segmentation](https://mlsanity.com/paper/2604.02948/) (2026)
- [Hierarchical Awareness Adapters with Hybrid Pyramid Feature Fusion for Dense Depth Prediction](https://mlsanity.com/paper/2604.03339/) (2026)

## Abstract (arXiv)

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.