PaperlayerView exact revision on arXiv

Source-linked AI summary

Axial-DeepLab: Stand-Alone Axial-Attention for Panoptic Segmentation

Huiyu Wang, Yukun Zhu, Bradley Green, Hartwig Adam, Alan Yuille, Liang-Chieh Chen

arXiv:2003.07853v2cs.CVcs.LG

TL;DR

Convolution is efficient but misses long-range context, while unrestricted self-attention is expensive and local attention limits receptive field. The paper combines factorized axial attention with position-sensitive design, achieving strong results across image classification and segmentation benchmarks, with slower runtime than convolutional counterparts.

  • Problem

    Self-attention must balance the computational cost of global interactions, the receptive-field limits of local attention, and the loss of positional information needed for spatial structures.

  • Method

    The paper factorizes 2D self-attention into two 1D axial attentions and adds position-sensitive attention to build models with large or global receptive fields.

  • Results

    Axial-attention models achieve state-of-the-art results among stand-alone self-attention models on ImageNet and on several segmentation benchmarks, including COCO, Mapillary Vistas, and Cityscapes.

  • Takeaways & Limitations

    Axial attention is an effective building block for modern computer vision models across image classification and segmentation tasks.

  • Takeaways & Limitations

    Despite saving M-Adds, axial-attention models run slower than convolutional counterparts because specialized accelerator kernels are currently lacking.

Abstract

from arXiv · show

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.

1 Introduction

Convolution is efficient because of locality but struggles with long-range relations, motivating axial attention as a scalable alternative for fully attentional vision models. The paper introduces position-sensitive axial attention and applies it across classification and segmentation, reporting strong benchmark performance.

  • Convolution reduces parameters and M-Adds through locality but makes modeling long-range relations challenging.
  • Local constraints make stand-alone self-attention computationally feasible but limit its receptive field, especially for high-resolution segmentation.
  • The paper proposes position-sensitive axial attention to support stand-alone attention models with large or global receptive fields.
  • Axial attention is presented as a backbone for image classification, panoptic segmentation, instance segmentation, and semantic segmentation.
  • Axial-DeepLab improves significantly over bottom-up state-of-the-art on COCO, matches the performance of two-stage methods, and surpasses previous state-of-the-art methods on Mapillary Vistas and Cityscapes.

2 Related Work

Related work addresses long-range interactions through expanded convolutions, image pyramids, attention, and proposal-based or proposal-free segmentation pipelines. Self-attention offers direct long-range interactions and parallelism, while efficient variants target its computational cost.

  • Top-down panoptic segmentation: Top-down panoptic segmentation methods generate object proposals and process each proposal sequentially.
  • Bottom-up panoptic segmentation: Bottom-up panoptic segmentation predicts semantic labels and groups thing pixels into clusters to obtain instance segmentation.
  • Self-attention: Research on long-range interactions includes atrous convolutions, larger kernels, image pyramids, and attention mechanisms.
  • Self-attention: Self-attention directly encodes long-range interactions and supports parallel computation across language, vision, and other tasks.

3 Method

The method combines position-sensitive self-attention with axial factorization to capture long-range or global context efficiently while preserving positional information. Axial-DeepLab adapts these layers for dense prediction without atrous attention or ASPP.

  • Position-Sensitive Self-Attention: Self-attention’s global computation costs O(h^2w^2) and its global pooling lacks positional information needed for spatial structures.Local constraints reduce computation but limit the receptive field and retain quadratic complexity in region length.
  • Position-Sensitive Self-Attention: Position-sensitive self-attention adds key-dependent positional bias and lets outputs retrieve relative positions alongside values.The design addresses the loss of precise spatial structure when using large receptive fields or memory banks.
  • Position-Sensitive Self-Attention: The proposed position-sensitive self-attention captures long-range interactions with precise positional information at reasonable computation overhead.Relative positional vectors are shared across attention heads, limiting additional parameters.
  • Axial-Attention: Axial-attention factorizes spatial processing by applying one-dimensional attention along an axis, then sequentially along height and width.Two consecutive multi-head layers capture global information across both axes.
  • Axial-Attention: Axial-attention reduces complexity to O(hwm) while enabling a global receptive field by setting the span m to the whole input.A fixed m can instead reduce memory usage on very large feature maps.
  • Axial-DeepLab: Axial-DeepLab removes the final-stage stride, uses output stride 16, and omits atrous attention and ASPP because axial-attention encodes global or multi-scale information.The method is built by converting Axial-ResNet for segmentation and is evaluated without ASPP against Panoptic-DeepLab variants.

4 Experimental Results

Experiments evaluate Axial-ResNet and Axial-DeepLab across ImageNet, COCO, Mapillary Vistas, and Cityscapes, showing strong accuracy, efficiency, and scale-robustness results. Ablations indicate that position sensitivity and axial attention are especially useful for segmentation with large inputs.

  • ImageNet: Axial-ResNet experiments compare accuracy, parameters, and M-Adds across network widths under matched ImageNet training protocols.The models include Conv-Stem and fully attentional variants, with channel-scaled S, M, L, and XL configurations.
  • ImageNet: 0.5% higher accuracy, 1.44× fewer parameters, and 1.09× fewer M-Adds distinguish Full Axial-Attention from Full Attention.Conv-Stem + PS-Attention improves performance by 0.4%, while Conv-Stem + Axial-Attention matches Conv-Stem + Attention with greater efficiency.
  • COCO: 43.9% PQ is achieved by multi-scale Axial-DeepLab-L on COCO validation, outperforming Panoptic-DeepLab by 2.7% PQ.Single-scale Axial-DeepLab-S exceeds multi-scale Panoptic-DeepLab by 0.6% PQ while using 3.8× fewer parameters and 27× fewer M-Adds.
  • COCO: 44.2% PQ on COCO test-dev surpasses Panoptic-DeepLab by 2.8% PQ and sets a new state-of-the-art among bottom-up approaches.Larger backbones produce consistent improvements across Axial-DeepLab variants.
  • COCO: 30% relative gain on huge images and 40% on small images in the COCO stress test show improved performance across extreme resolutions.The test includes out-of-distribution input resolutions and compares Axial-DeepLab-L with Panoptic-DeepLab X-71.
  • Mapillary Vistas: Axial-DeepLab-L outperforms prior state-of-the-art methods on Mapillary Vistas, including gains of 0.8% PQ, 0.3% AP, and 0.8% mIoU in a multi-scale comparison.The model also reports a 2.4% PQ advantage over the previous best single-scale Panoptic-DeepLab.
  • Cityscapes and ablations: More than 1% gains on all three Cityscapes segmentation tasks result from axial attention relative to position-sensitive attention, without ASPP and with fewer parameters and M-Adds.The gains occur in PQ, AP, and mIoU, supporting the value of long-range context for large input images.
  • Cityscapes and ablations: Larger axial spans consistently improve performance at marginal cost in Cityscapes ablations.The span m controls the spatial extent of local regions in an axial block.

5 Conclusion and Discussion

The paper presents position-sensitive axial attention as an effective building block for classification and segmentation. It reports state-of-the-art performance across several benchmarks while noting a runtime disadvantage versus convolutional counterparts.

  • Conclusion: Position-sensitive axial attention is effective for image classification and segmentation, including panoptic, instance, and semantic segmentation.Axial-ResNet and Axial-DeepLab extend the approach across ImageNet, COCO, Mapillary Vistas, and Cityscapes.
  • Discussion: Axial attention factorizes attention across axes to provide large or global receptive fields while reducing M-Adds relative to full attention.Its design is related to decoupled convolution but retains content-based relations and positional bias.
  • Limitations: Despite saving M-Adds, axial-attention models run slower than convolutional counterparts because specialized accelerator kernels are currently lacking.The paper identifies kernel support as a possible avenue for improving runtime.

Appendix A Runtime

The runtime analysis shows that Conv-Stem Axial-ResNet-L is faster than a comparable stand-alone attention model, while TPU graph transposes remain an optimization opportunity.

  • Runtime: 16.54 ms: Conv-Stem Axial-ResNet-L inference on a V100 for 224x224 inputs with batch size 1.The measurement includes input standardization and the final projection to 1000 logits, averaged over five runs.
  • Caveat: TPU benchmarking graphs transpose between NCHW and NHWC around almost every conv2d operation, and further optimization could improve inference speed.
  • Runtime: Axial-ResNet-L runs faster than Conv-Stem Stand-Alone-L despite splitting one layer into two.The paper attributes this to axial-attention making better use of existing accelerator kernels.
  • Kernel efficiency: Width-axis attention is parallelizable over the height axis, forming a large batch of 1D row operations.The batch size for these row operations is the input height.
  • Kernel efficiency: Axial attention avoids extracting padded 2D memory blocks and their splits and concatenations, which are inefficient on accelerators.

Appendix B Axial-Decoder

The appendix replaces dual convolutional decoders with axial-attention blocks at upsampling stages and finds comparable segmentation results across decoder and output-stride settings.

  • Axial-Decoder: Axial-DeepLab’s axial-decoder applies one axial-attention block at each upsampling stage.The design uses three such blocks, analogous to Panoptic-DeepLab’s three 5×5 convolutions.
  • Ablation: Axial-decoder yields results similar to dual convolutional decoders on the Cityscapes validation ablation.The table varies output stride and decoder type; AD denotes use of the axial-decoder.
  • Ablation: Axial-DeepLab models are robust to output stride in the reported Cityscapes validation experiments.

Appendix C COCO Visualization

COCO visualizations show robustness to occlusion, attention heads specializing across local and long-range context, and improved fitting of position-sensitive offset predictions.

  • COCO Visualization: Axial-DeepLab-L demonstrates robustness to occlusion compared with Panoptic-DeepLab on COCO validation examples.
  • Attention Maps: Some attention heads focus on local details while others capture long-range context in low- and high-level blocks.The visualizations use stage 3 block 2 and stage 4 block 3, both at 65×65 resolution under output stride 16.
  • COCO Visualization: The visualizations include cases where Axial-DeepLab captures occluded body parts or cables absent from ground-truth labels.
  • COCO Visualization: Axial-DeepLab correctly distinguishes one person occluding another where the ground truth treats them as one instance.
  • Training Loss: Axial-DeepLab fits the three COCO training losses better than Panoptic-DeepLab, especially for offset prediction.The paper connects this result to position-sensitive attention and axial-attention’s long-range modeling ability.

Appendix D Raw Data

The appendix provides raw ImageNet and COCO performance data corresponding to the paper’s trade-off and scale-stress-test comparisons.

  • ImageNet: Table 10 reports ImageNet validation results by model width and identifies Full models as stand-alone self-attention models without spatial convolutions.
  • COCO: Figure 10 reports raw COCO validation performance for Axial-DeepLab-L and Panoptic-DeepLab in the scale stress test.
Loading 2003.07853v2…