Polarized Self-Attention: Towards High-quality Pixel-wise Regression

Abstract

Pixel-wise regression is probably the most common problem in fine-grained computer vision tasks, such as estimating keypoint heatmaps and segmentation masks. These regression problems are very challenging particularly because they require, at low computation overheads, modeling long-range dependencies on high-resolution inputs/outputs to estimate the highly nonlinear pixel-wise semantics. While attention mechanisms in Deep Convolutional Neural Networks(DCNNs) has become popular for boosting long-range dependencies, element-specific attention, such as Nonlocal blocks, is highly complex and noise-sensitive to learn, and most of simplified attention hybrids try to reach the best compromise among multiple types of tasks. In this paper, we present the Polarized Self-Attention(PSA) block that incorporates two critical designs towards high-quality pixel-wise regression: (1) Polarized filtering: keeping high internal resolution in both channel and spatial attention computation while completely collapsing input tensors along their counterpart dimensions. (2) Enhancement: composing non-linearity that directly fits the output distribution of typical fine-grained regression, such as the 2D Gaussian distribution (keypoint heatmaps), or the 2D Binormial distribution (binary segmentation masks). PSA appears to have exhausted the representation capacity within its channel-only and spatial-only branches, such that there is only marginal metric differences between its sequential and parallel layouts. Experimental results show that PSA boosts standard baselines by $2-4$ points, and boosts state-of-the-arts by $1-2$ points on 2D pose estimation and semantic segmentation benchmarks.

References in corpus (10)

• A Survey on Deep Learning in Medical Image Analysis
• The Cityscapes Dataset for Semantic Urban Scene Understanding
• Hierarchical Multi-Scale Attention for Semantic Segmentation
• GCNet: Non-local Networks Meet Squeeze-Excitation Networks and Beyond
• Residual Non-local Attention Networks for Image Restoration
• Auto-DeepLab: Hierarchical Neural Architecture Search for Semantic Image Segmentation
• Learning a Discriminative Feature Network for Semantic Segmentation
• PAD-Net: Multi-Tasks Guided Prediction-and-Distillation Network for Simultaneous Depth Estimation and Scene Parsing
• Convolution, attention and structure embedding
• Semantic Correlation Promoted Shape-Variant Context for Segmentation

Cited by in corpus (6)

• Efficient Multi-Scale Attention Module with Cross-Spatial Learning
• DINOv2: Learning Robust Visual Features without Supervision
• Wideband Beamforming for RIS Assisted Near-Field Communications
• Diabetic Foot Ulcer Grand Challenge 2022 Summary
• Improving 2D Human Pose Estimation in Rare Camera Views with Synthetic Data
• PAENet: A Progressive Attention-Enhanced Network for 3D to 2D Retinal Vessel Segmentation