[Paper Review] Polarized Self-Attention: Towards High-quality Pixel-wise Regression

Huajun Liu, Fuqiang Liu+2 more|arXiv (Cornell University)|Jul 2, 2021

Advanced Neural Network Applications59 references152 citations

TL;DR

The paper introduces Polarized Self-Attention (PSA), a plug-in block that preserves high internal resolution in channel and spatial branches and uses a Softmax-Sigmoid non-linearity to better fit pixel-wise regression outputs, boosting performance on pose estimation and segmentation benchmarks.

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.

Motivation & Objective

• Motivate and address pixel-wise regression challenges that require preserving high-resolution features at modest cost.
• Develop an attention block that keeps high internal resolution in both channel and spatial dimensions.
• Design a non-linearity (Softmax-Sigmoid) that matches typical pixel-wise regression outputs (Gaussian heatmaps, binomial masks).
• Demonstrate that PSA provides consistent improvements when plugged into standard baselines and SOTA models.

Proposed method

• Propose PSA with two polarized branches: a channel-only branch and a spatial-only branch, each preserving high-resolution attention computations.
• Maintain high channel resolution (C/2) and full spatial resolution [W, H] in attention while collapsing input along the opposite dimension.
• Use a Softmax at a bottleneck and a Sigmoid tone-mapping (Softmax-Sigmoid composition) to fit 2D Gaussian heatmaps and 2D Binomial segmentation maps.
• Provide parallel (A^ch + A^sp) and sequential (A^sp(A^ch(X)) form) compositions, with marginal difference in performance.
• Insert PSA after the first 3x3 convolution within residual blocks in backbones like ResNet and HRNet, with minimal overhead.
• Evaluate PSA on 2D human pose estimation (MS-COCO) and semantic segmentation (PASCAL VOC2012, Cityscapes) against baselines and state-of-the-art models.

Experimental results

Research questions

• RQ1 Can a high-resolution attention block with polarized filtering improve pixel-wise regression without prohibitive computation?
• RQ2 Does preserving channel and spatial resolutions in attention, along with a distribution-matching non-linearity, yield measurable gains on keypoint heatmaps and segmentation masks?
• RQ3 Are parallel and sequential PSA layouts differently effective, or are they largely similar in performance and cost?
• RQ4 How does PSA perform when applied to strong baselines and SOTA models for 2D pose estimation and segmentation?
• RQ5 What is the trade-off between accuracy gains and computational/memory overhead when inserting PSA into common DCNN heads?

Key findings

• PSA boosts baselines by about 2–4 AP for pose estimation tasks.
• PSA boosts state-of-the-art models by about 1–2 AP on 2D pose estimation benchmarks.
• In semantic segmentation, PSA yields +1.8–2.6 mIoU across tested backbones and achieves new state-of-the-art on Cityscapes under HRNet-OCR+PSA configurations (86.95 mIoU).
• PSA maintains higher internal attention resolution (channel: C/2; space: [W, H]) and uses a Softmax-Sigmoid non-linearity to better fit Gaussian heatmaps and Binomial masks.
• Parallel and sequential PSA layouts show only marginal differences in performance, with similar FLOPs and parameter counts.
• Ablation studies indicate PSA’s channel-only branch (A^ch) often outperforms spatial-only (A^sp), and combinations (PSA_p or PSA_s) outperform individual branches.