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. Мы надеемся, что наш простой, но эффективный дизайн вдохновит на дальнейшие исследования по использованию масковых трансформеров для большего числа задач плотного предсказания. Код и модель будут доступны.

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.