PolyMaX: 마스크 트랜스포머를 활용한 일반적 밀집 예측

PolyMaX: General Dense Prediction with Mask Transformer

November 9, 2023

저자: 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

초록

의미론적 분할(semantic segmentation), 깊이 추정(depth estimation), 표면 법선 예측(surface normal prediction)과 같은 조밀 예측(dense prediction) 작업은 픽셀 단위 분류(이산 출력) 또는 회귀(연속 출력)로 쉽게 공식화될 수 있습니다. 이 픽셀 단위 예측 패러다임은 완전 합성곱 네트워크(fully convolutional networks)의 보편화로 인해 여전히 널리 사용되고 있습니다. 그러나 최근 분할 작업의 최전선에서는, 특히 마스크 트랜스포머(mask transformers)와 같은 트랜스포머 아키텍처의 등장으로 인해 픽셀 단위 예측에서 클러스터 예측(cluster-prediction)으로의 패러다임 전환이 이루어지고 있습니다. 이러한 전환에도 불구하고, 깊이 추정 및 표면 법선 예측과 같이 연속 출력이 필요한 다른 조밀 예측 작업에서는 픽셀 단위 예측 패러다임을 기반으로 한 방법들이 여전히 벤치마크를 지배하고 있습니다. DORN과 AdaBins가 연속 출력 공간을 이산화함으로써 깊이 추정에서 성공을 거둔 것에 영감을 받아, 우리는 클러스터 예측 기반 방법을 일반적인 조밀 예측 작업으로 일반화하는 것을 제안합니다. 이를 통해 마스크 트랜스포머 프레임워크와 조밀 예측 작업을 통합할 수 있습니다. 특히, 결과적으로 얻은 모델인 PolyMaX는 NYUD-v2 데이터셋의 세 가지 벤치마크에서 최첨단 성능을 보여줍니다. 우리의 간단하지만 효과적인 설계가 더 많은 조밀 예측 작업에서 마스크 트랜스포머를 활용하는 연구에 영감을 줄 수 있기를 바랍니다. 코드와 모델은 공개될 예정입니다.

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.