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

要旨

セマンティックセグメンテーション、深度推定、表面法線予測などの密な予測タスクは、ピクセルごとの分類（離散出力）または回帰（連続出力）として容易に定式化できます。このピクセルごとの予測パラダイムは、完全畳み込みネットワークの普及により、長らく人気を保ってきました。しかし、最近のセグメンテーションタスクの最前線では、特にマスクトランスフォーマーの登場により、ピクセルごとの予測からクラスター予測へのパラダイムシフトが起こっています。マスクトランスフォーマーは、ピクセルではなくマスクに対して直接ラベルを予測します。このシフトにもかかわらず、深度推定や表面法線予測など、連続出力を必要とする他の密な予測タスクでは、ピクセルごとの予測パラダイムに基づく手法が依然としてベンチマークを支配しています。DORNやAdaBinsが連続出力空間を離散化することで深度推定で成功を収めたことに着想を得て、我々はクラスター予測ベースの手法を一般的な密な予測タスクに一般化することを提案します。これにより、マスクトランスフォーマーフレームワークを用いて密な予測タスクを統一的に扱うことが可能になります。驚くべきことに、結果として得られたモデルPolyMaXは、NYUD-v2データセットの3つのベンチマークで最先端の性能を発揮しました。我々のシンプルでありながら効果的な設計が、より多くの密な予測タスクにおいてマスクトランスフォーマーを活用する研究を促進することを期待しています。コードとモデルは公開予定です。

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.