Source-linked AI summary
Polarized Self-Attention: Towards High-quality Pixel-wise Regression
Huajun Liu, Fuqiang Liu, Xinyi Fan, Dong Huang
TL;DR
Pixel-wise regression requires high-resolution representations for nonlinear semantics, while existing self-attention can be redundant, memory-intensive, or noise-sensitive. PSA preserves high resolution in channel-only and spatial-only attention branches and adds nonlinear enhancement, boosting baseline networks by 2.6 to 4.3 AP with minor overheads.
Problem
Pixel-wise regression is difficult because nonlinear object edges and body parts are hard to encode in low-resolution features, while Non-Local attention uses a costly HW × HW matrix and can highlight spatial noise.
Method
PSA preserves high internal resolution in polarized channel-only and spatial-only attention branches and incorporates nonlinear composition to leverage the retained information.
Results
2.6 to 4.3 AP: PSA boosts baseline networks on MS-COCO val2017 with minor computation and parameter overheads.
Takeaways & Limitations
PSA significantly boosts compared DCNNs, and its parallel and sequential channel-spatial layouts have similar AP, computation, parameters, memory, and inference time.
Takeaways & Limitations
PSA's benefits remain unclear in complex DCNN heads combining pixel-wise regression with classification and displacement regression, including instance, anchor-free, and panoptic segmentation.
Abstract
from arXiv · showhide
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.
1. Introduction
Pixel-wise regression maps image pixels to semantic scores but is difficult because fine-grained structures are highly nonlinear and standard attention designs trade resolution, complexity, and generality. PSA addresses this with high-resolution polarized attention and distribution-matched nonlinear enhancement.
- Pixel-wise regression maps pixels sharing semantics to common scores, with keypoint heatmaps and segmentation masks as typical tasks.
- Encoder-decoder DCNNs reduce spatial resolution for efficiency, but nonlinear pixel appearances and patch shapes make fine-grained semantics difficult to encode.
- General-purpose attention hybrids often compromise across classification, detection, and segmentation, helping explain the popularity of channel-only attention blocks.
- PSA preserves high internal resolution in channel and spatial attention while collapsing tensors along counterpart dimensions.
- PSA combines softmax-sigmoid nonlinear enhancement with polarized attention to fit output distributions in fine-grained regression.
2. Related Work
Prior work improves pixel-wise vision through higher-resolution DCNNs and attention mechanisms, while full-tensor attention remains costly and simplified variants reduce attention resolution. PSA extends these efforts by targeting high-resolution attention specifically for pixel-wise regression.
- Recent pose-estimation methods pursue higher resolution through deconvolution, parallel multiresolution branches, and sub-pixel keypoint estimation.
- PSA pursues the same high-resolution objective from the attention perspective and further boosts these DCNN architectures.
- Self-attention captures long-range interactions, and parallel convolution-attention structures address both short- and long-range dependencies.
- PSA advances self-attention for pixel-wise regression and may also extend to convolution-augmented attention variants.
- Full-tensor non-local attention models pairwise feature similarities but incur very large memory and computational costs, motivating efficient reductions.
- PSA addresses fine-grained regression by retaining the highest attention resolution among existing blocks and directly fitting typical output distributions.
3. Our Method
PSA targets pixel-wise regression by preserving high-resolution channel and spatial attention while fitting nonlinearities to fine-grained output distributions. Its polarized filtering collapses features along one dimension, and its branches can be composed in parallel or sequential layouts.
- Pixel-wise regression requires channel-specific weighting for class scores and spatial-specific weighting for pixels sharing semantics.
- Full-tensor self-attention can highlight features element-wise, but its attention tensor is complex and noise-prone to learn.
- Non-Local attention uses an HW × HW similarity matrix, creating substantial memory and computational costs while its actual weights remain channel-specific.
- Internal Attention Resolution: Pixel-wise regression cannot afford the low-resolution features produced by standard backbones because object edges and body parts are difficult to encode at low resolution.
- Polarized Self-Attention: PSA applies polarized filtering by collapsing features in one direction while preserving high resolution orthogonally, then uses Softmax normalization and Sigmoid tone-mapping for enhancement.
- Internal Attention Resolution: PSA preserves the highest attention resolution for both channel (C/2) and spatial ([W, H]) dimensions among compared attention blocks.
4. Experiments
Experiments show that PSA improves standard and state-of-the-art pose-estimation and segmentation systems with minor computational overhead, while parallel and sequential layouts perform similarly.
- PSA vs. Baselines: PSA boosts Simple-Baseline and HRNet pose-estimation networks by 2.6 to 4.3 AP with minor computation and parameter overhead.Without ImageNet pre-training, PSA with Res50 reaches 76.5 AP and outperforms Simple-Baseline with Resnet152.
- PSA vs. Baselines: PSA boosts DeepLabV3Plus semantic-segmentation baselines by 1.8 to 2.6 mIoU with minor computation and parameter overhead.PSA with Res50 reaches 79.0 mIoU, exceeding DeepLabV3Plus with Resnet50 and Resnet101 comparisons described in the passage.
- Comparing with State-of-the-Arts: PSA reaches 79.5 AP on MS-COCO keypoint testdev when added to UDP-Pose with an HRnet-W48 backbone, improving the baseline by 1.7 points.The result is reported as a new state-of-the-art single-model pose-estimation result.
- Comparing with State-of-the-Arts: PSA reaches 86.95 mIoU on the Cityscapes validation set with HRNet-OCR, improving the strong baseline by 2 points.Parallel and sequential PSA layouts show only a subtle metric difference in this experiment.
- Ablation Study: Channel-only PSA outperforms spatial-only attention, while parallel and sequential compositions provide similar AP, Flops, parameters, memory, and inference time.The ablation evaluates configurations on Simple-Baseline with a Resnet50 backbone.
- Ablation Study: At similar overheads, both parallel and sequential PSAs outperform the compared attention blocks, while Non-Local attention costs the most memory and gives a 2.3 AP boost.The comparison is conducted in the related self-attention methods ablation.
5. Conclusion and Future Work
PSA targets high-quality pixel-wise regression by preserving high internal resolution in channel-only and spatial-only attention branches and using nonlinear composition. Its benefit in complex DCNN heads that combine pixel-wise regression with classification or displacement regression remains unclear.
- Conclusion: PSA significantly boosts the compared DCNNs through high-resolution polarized attention branches and nonlinear composition.The two designs preserve high internal resolution and leverage the information retained in those branches.
- Conclusion: PSA can potentially benefit any computer vision task involving pixel-wise regression.
- Future Work: PSA's benefits remain unclear for pixel-wise regression embedded in complex DCNN heads.The cited examples include instance segmentation, anchor-free object detection, and panoptic segmentation.