PolyMaX: Algemene Dichte Voorspelling met Mask Transformer

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

Samenvatting

Dichte voorspellingstaken, zoals semantische segmentatie, diepteschatting en oppervlaktenormaalvoorspelling, kunnen eenvoudig worden geformuleerd als per-pixel classificatie (discrete uitvoer) of regressie (continue uitvoer). Dit per-pixel voorspellingsparadigma is populair gebleven vanwege de prevalentie van volledig convolutionele netwerken. Echter, op het recente front van segmentatietaken heeft de gemeenschap een paradigmaverschuiving gezien van per-pixel voorspelling naar clustervoorspelling met de opkomst van transformer-architecturen, met name de mask transformers, die direct een label voor een masker voorspellen in plaats van voor een pixel. Ondanks deze verschuiving domineren methoden gebaseerd op het per-pixel voorspellingsparadigma nog steeds de benchmarks voor andere dichte voorspellingstaken die continue uitvoer vereisen, zoals diepteschatting en oppervlaktenormaalvoorspelling. Gemotiveerd door het succes van DORN en AdaBins in diepteschatting, bereikt door het discretiseren van de continue uitvoerruimte, stellen we voor om de clustervoorspellingsmethode te generaliseren naar algemene dichte voorspellingstaken. Dit stelt ons in staat om dichte voorspellingstaken te verenigen met het mask transformer-framework. Opmerkelijk is dat het resulterende model PolyMaX state-of-the-art prestaties laat zien op drie benchmarks van de NYUD-v2 dataset. We hopen dat ons eenvoudige maar effectieve ontwerp meer onderzoek kan inspireren naar het benutten van mask transformers voor meer dichte voorspellingstaken. Code en model zullen beschikbaar worden gesteld.

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.