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.
Source: §3.1 · pp. 3–4Method
Results
2.6 to 4.3 AP: PSA boosts baseline networks on MS-COCO val2017 with minor computation and parameter overheads.
Source: §4.1 · p. 5Takeaways & Limitations
PSA significantly boosts compared DCNNs, and its parallel and sequential channel-spatial layouts have similar AP, computation, parameters, memory, and inference time.
Source: §4.3 · p. 6, §5 · p. 8PSA's benefits remain unclear in complex DCNN heads combining pixel-wise regression with classification and displacement regression, including instance, anchor-free, and panoptic segmentation.
Source: §5 · p. 8
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.Source: §1 · p. 1
- Encoder-decoder DCNNs reduce spatial resolution for efficiency, but nonlinear pixel appearances and patch shapes make fine-grained semantics difficult to encode.Source: §1 · p. 2
- General-purpose attention hybrids often compromise across classification, detection, and segmentation, helping explain the popularity of channel-only attention blocks.Source: §1 · p. 2
- PSA preserves high internal resolution in channel and spatial attention while collapsing tensors along counterpart dimensions.Source: §1 · p. 2
- PSA combines softmax-sigmoid nonlinear enhancement with polarized attention to fit output distributions in fine-grained regression.Source: §1 · p. 2
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.Source: §2 · p. 2
- PSA pursues the same high-resolution objective from the attention perspective and further boosts these DCNN architectures.Source: §2 · p. 2
- Self-attention captures long-range interactions, and parallel convolution-attention structures address both short- and long-range dependencies.Source: §2 · pp. 2–3
- PSA advances self-attention for pixel-wise regression and may also extend to convolution-augmented attention variants.Source: §2 · p. 3
- Full-tensor non-local attention models pairwise feature similarities but incur very large memory and computational costs, motivating efficient reductions.Source: §2 · p. 3
- PSA addresses fine-grained regression by retaining the highest attention resolution among existing blocks and directly fitting typical output distributions.Source: §2 · p. 3
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.Source: §3.1 · p. 3
- Full-tensor self-attention can highlight features element-wise, but its attention tensor is complex and noise-prone to learn.Source: §3.1 · p. 3
- Non-Local attention uses an HW × HW similarity matrix, creating substantial memory and computational costs while its actual weights remain channel-specific.Source: §3.1 · p. 3
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.Source: §3.1 · p. 4
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.Source: §3.2 · p. 4
Internal Attention Resolution
- PSA preserves the highest attention resolution for both channel (C/2) and spatial ([W, H]) dimensions among compared attention blocks.Source: §3.2 · p. 5
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.Source: §4.1 · p. 5
- 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.Source: §4.1 · p. 6
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.Source: §4.2 · p. 6
- 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.Source: §4.2 · p. 6
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.Source: §4.3 · p. 6
- 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.Source: §4.3 · p. 6
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.Source: §5 · p. 8
- PSA can potentially benefit any computer vision task involving pixel-wise regression.Source: §5 · p. 8
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.Source: §5 · p. 8