[Paper Review] Axial-DeepLab: Stand-Alone Axial-Attention for Panoptic Segmentation

Huiyu Wang, Yukun Zhu+4 more|arXiv (Cornell University)|Mar 17, 2020

Advanced Neural Network Applications105 references66 citations

TL;DR

Introduces position-sensitive axial-attention to build stand-alone axial-attention models, achieving state-of-the-art panoptic segmentation results on COCO, Mapillary Vistas, and Cityscapes, with strong efficiency gains over previous stand-alone self-attention methods.

ABSTRACT

Convolution exploits locality for efficiency at a cost of missing long range context. Self-attention has been adopted to augment CNNs with non-local interactions. Recent works prove it possible to stack self-attention layers to obtain a fully attentional network by restricting the attention to a local region. In this paper, we attempt to remove this constraint by factorizing 2D self-attention into two 1D self-attentions. This reduces computation complexity and allows performing attention within a larger or even global region. In companion, we also propose a position-sensitive self-attention design. Combining both yields our position-sensitive axial-attention layer, a novel building block that one could stack to form axial-attention models for image classification and dense prediction. We demonstrate the effectiveness of our model on four large-scale datasets. In particular, our model outperforms all existing stand-alone self-attention models on ImageNet. Our Axial-DeepLab improves 2.8% PQ over bottom-up state-of-the-art on COCO test-dev. This previous state-of-the-art is attained by our small variant that is 3.8x parameter-efficient and 27x computation-efficient. Axial-DeepLab also achieves state-of-the-art results on Mapillary Vistas and Cityscapes.

Motivation & Objective

• Motivate modeling long-range context efficiently without conventional convolution locality constraints.
• Propose position-sensitive axial-attention to enable large/global receptive fields in stand-alone models.
• Demonstrate Axial-ResNet and Axial-DeepLab as backbones for classification and panoptic segmentation.
• Show state-of-the-art performance on COCO, Mapillary Vistas, and Cityscapes with improved efficiency.

Proposed method

• Factorize 2D self-attention into sequential height- and width-axis 1D attentions (axial-attention).
• Introduce position-sensitive self-attention with query-, key-, and value-dependent positional terms (r^q, r^k, r^v).
• Replace 3x3 convs in ResNet blocks with axial-attention layers to form Axial-ResNet.
• Convert Axial-ResNet to Axial-DeepLab for segmentation by adapting strides and removing ASPP.
• Train and evaluate on ImageNet for classification and COCO, Mapillary Vistas, Cityscapes for panoptic, instance, and semantic segmentation.
• Use span m to control axial-attention extent; allow global receptive field by setting m to input size; apply two consecutive axial-attention layers for height and width.

Experimental results

Research questions

• RQ1 Can stand-alone axial-attention achieve global receptive fields with lower complexity than full 2D self-attention?
• RQ2 Does position-sensitive axial-attention improve segmentation performance over prior stand-alone attention methods?
• RQ3 How do Axial-ResNet and Axial-DeepLab perform on panoptic, instance, and semantic segmentation benchmarks compared to bottom-up state-of-the-art?
• RQ4 What is the impact of axial-attention span and model size on accuracy and efficiency across datasets?
• RQ5 Is it feasible to replace traditional convolutions with axial-attention in backbone networks without ASPP for competitive segmentation results?

Key findings

• Axial-DeepLab-L achieves 43.9 PQ on COCO test-dev, surpassing Panoptic-DeepLab by 2.7 PQ.
• Single-scale Axial-DeepLab-S outperforms DeeperLab by 8.0 PQ on COCO val and outperforms multi-scale SSAP and single-scale Panoptic-DeepLab by 5.3 and 2.1 PQ respectively.
• Axial-DeepLab-L with MS reaches 44.2 PQ on COCO test-dev, achieving state-of-the-art among bottom-up methods and closing the gap to top-down approaches.
• On Mapillary Vistas validation, Axial-DeepLab-L outperforms state-of-the-art in single-scale and multi-scale settings; with MV pretraining, Axial-DeepLab-XL attains 68.5 PQ and 44.2 AP.
• Cityscapes validation shows Axial-DeepLab variants outperform ResNet-50 baselines, with larger models and MS further improving PQ and mIoU.
• Across experiments, axial-attention with position-sensitivity yields consistent gains in PQ, AP, and mIoU compared to prior stand-alone attention methods.