PolyMaX: Predição Densa Geral com 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

Resumo

Tarefas de predição densa, como segmentação semântica, estimativa de profundidade e predição de normais de superfície, podem ser facilmente formuladas como classificação por pixel (saídas discretas) ou regressão (saídas contínuas). Esse paradigma de predição por pixel permanece popular devido à prevalência de redes totalmente convolucionais. No entanto, na fronteira recente da tarefa de segmentação, a comunidade tem testemunhado uma mudança de paradigma, da predição por pixel para a predição por cluster, com o surgimento de arquiteturas de transformadores, particularmente os transformadores de máscara, que preveem diretamente um rótulo para uma máscara em vez de um pixel. Apesar dessa mudança, métodos baseados no paradigma de predição por pixel ainda dominam os benchmarks em outras tarefas de predição densa que exigem saídas contínuas, como estimativa de profundidade e predição de normais de superfície. Motivados pelo sucesso de DORN e AdaBins na estimativa de profundidade, alcançado pela discretização do espaço de saída contínuo, propomos generalizar o método baseado em predição por cluster para tarefas gerais de predição densa. Isso nos permite unificar tarefas de predição densa com o framework de transformadores de máscara. Notavelmente, o modelo resultante, PolyMaX, demonstra desempenho de última geração em três benchmarks do conjunto de dados NYUD-v2. Esperamos que nosso design simples, porém eficaz, inspire mais pesquisas sobre a exploração de transformadores de máscara para mais tarefas de predição densa. O código e o modelo serão disponibilizados.

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.