PolyMaX: Predicción Densa General con Transformador de Máscaras

PolyMaX: General Dense Prediction with Mask Transformer

November 9, 2023

Autores: 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

Resumen

Las tareas de predicción densa, como la segmentación semántica, la estimación de profundidad y la predicción de normales de superficie, pueden formularse fácilmente como clasificación por píxel (salidas discretas) o regresión (salidas continuas). Este paradigma de predicción por píxel ha permanecido popular debido a la prevalencia de las redes completamente convolucionales. Sin embargo, en la frontera reciente de la tarea de segmentación, la comunidad ha sido testigo de un cambio de paradigma desde la predicción por píxel hacia la predicción por clúster con el surgimiento de arquitecturas transformadoras, particularmente los transformadores de máscaras, que predicen directamente una etiqueta para una máscara en lugar de un píxel. A pesar de este cambio, los métodos basados en el paradigma de predicción por píxel aún dominan los puntos de referencia en otras tareas de predicción densa que requieren salidas continuas, como la estimación de profundidad y la predicción de normales de superficie. Motivados por el éxito de DORN y AdaBins en la estimación de profundidad, logrado mediante la discretización del espacio de salida continuo, proponemos generalizar el método basado en predicción por clúster a tareas generales de predicción densa. Esto nos permite unificar las tareas de predicción densa con el marco de los transformadores de máscaras. Notablemente, el modelo resultante PolyMaX demuestra un rendimiento de vanguardia en tres puntos de referencia del conjunto de datos NYUD-v2. Esperamos que nuestro diseño simple pero efectivo pueda inspirar más investigaciones sobre la explotación de transformadores de máscaras para más tareas de predicción densa. El código y el modelo estarán 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.