Polarized Self-Attention: Towards High-quality Pixel-wise Regression (2107.00782v2)

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.

• The paper introduces the PSA block that employs polarized filtering and enhancement to improve pixel-wise regression tasks.
• It refines channel and spatial computations, reducing noise and efficiently modeling long-range image dependencies in DCNNs.
• Performance results show 2-4 point improvements on baseline and state-of-the-art models with only marginal extra computation.

Polarized Self-Attention: Advancements in Pixel-wise Regression for Computer Vision

The paper introduces the Polarized Self-Attention (PSA) block, an innovative approach designed to improve pixel-wise regression tasks, which are critical in fine-grained computer vision contexts like keypoint estimation and semantic segmentation. Pixel-wise regression problems demand high computing efficiency while modeling long-range dependencies across high-resolution images. This paper argues for a novel self-attention strategy aimed at addressing challenges inherent in existing attention mechanisms within deep convolutional neural networks (DCNNs).

The core contribution of the PSA block lies in its twofold design: polarized filtering and enhancement. Polarized filtering maintains high internal resolution in channel and spatial computations while collapsing input tensors along their counterpart dimensions. This strategy counters the high complexity and noise susceptibility typically associated with element-specific attention methods, e.g., Nonlocal blocks. This focused attention mechanism allows PSA to exhaust the representation capacity within its channel-only and spatial-only branches.

Furthermore, the PSA block enhances model outputs through a non-linearity that naturally aligns with fine-grained regression outputs. It adapts to fit the 2D Gaussian distribution, a common requirement for keypoint heatmaps, and the binomial distribution typical for binary segmentation masks. This approach helps to maintain the fidelity of high-resolution inputs and outputs, ensuring that critical details are effectively captured and modeled.

Performance evaluations reveal that PSA adds significant improvements to standard baselines and even current state-of-the-art methods, demonstrating boosts of 2-4 points for standard baselines and 1-2 points for state-of-the-art models in tasks such as 2D pose estimation and semantic segmentation. These improvements are particularly notable given that they come with only marginal increases in computational and memory overhead compared to vanilla DCNNs, showcasing the efficiency of PSA’s design.

The implications of this research are substantial, indicating potential shifts in how pixel-wise regression is approached within DCNNs. PSA’s ability to boost performance without extra pre-training in some comparisons, like with the ResNet50 baseline, highlights its robustness and potential applications beyond initial evaluations.

Looking forward, the use of PSA in broader contexts presents an intriguing direction for future research. While the paper primarily evaluates PSA in the field of pixel-wise regression, its adaptability and impact suggest that it might also enhance tasks involving more complex DCNN heads, such as instance segmentation or anchor-free object detection. By effectively leveraging PSA’s high-resolution capacity and innovative non-linearity, future developments might further advance pixel-wise regression's integration with classification and coordinate regression problems within diverse computer vision applications.