PolyMaX : Prédiction dense générale avec un transformeur à masque

PolyMaX: General Dense Prediction with Mask Transformer

November 9, 2023

Auteurs: Xuan Yang, Liangzhe Yuan, Kimberly Wilber, Astuti Sharma, Xiuye Gu, Siyuan Qiao, Stephanie Debats, Huisheng Wang, Hartwig Adam, Mikhail Sirotenko, Liang-Chieh Chen

cs.AI

Résumé

Les tâches de prédiction dense, telles que la segmentation sémantique, l'estimation de la profondeur et la prédiction des normales de surface, peuvent être facilement formulées comme des classifications par pixel (sorties discrètes) ou des régressions (sorties continues). Ce paradigme de prédiction par pixel est resté populaire en raison de la prévalence des réseaux entièrement convolutifs. Cependant, sur le front récent des tâches de segmentation, la communauté assiste à un changement de paradigme, passant de la prédiction par pixel à la prédiction par clusters avec l'émergence des architectures de transformateurs, en particulier les transformateurs de masques, qui prédisent directement une étiquette pour un masque plutôt que pour un pixel. Malgré ce changement, les méthodes basées sur le paradigme de prédiction par pixel dominent toujours les benchmarks pour les autres tâches de prédiction dense nécessitant des sorties continues, comme l'estimation de la profondeur et la prédiction des normales de surface. Motivés par le succès de DORN et AdaBins dans l'estimation de la profondeur, obtenu en discrétisant l'espace de sortie continu, nous proposons de généraliser la méthode basée sur la prédiction par clusters à des tâches de prédiction dense générales. Cela nous permet d'unifier les tâches de prédiction dense avec le cadre des transformateurs de masques. De manière remarquable, le modèle résultant, PolyMaX, démontre des performances de pointe sur trois benchmarks du jeu de données NYUD-v2. Nous espérons que notre conception simple mais efficace inspirera davantage de recherches sur l'exploitation des transformateurs de masques pour davantage de tâches de prédiction dense. Le code et le modèle seront rendus disponibles.

English

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.