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

A DCNN for pixel-wise regression learns a weighted combination of features along two dimensions: (1) channel-specific weighting to estimate the class-specific output scores; (2) spatial-specific weighting to detect pixels of the same semantics. The self-attention mechanism applied to the DCNN is expected to further highlight features for both above goals.

Ideally, with a full-tensor self-attention $\mathbf{Z}=\mathbf{A}(\mathbf{X})\odot\mathbf{X}$ ($\mathbf{A}(\mathbf{X})\in\Re^{C\times H\times W}$), the highlighting could potentially be achieved at the element-wise granularity ($C\times H\times W$ elements). However, the attention tensor $\mathbf{A}$ is very complex and noise-prone to learn directly. In the Non-Local self-attention block [47], $\mathbf{A}$ is calculated as,

$$\mathbf{A}=\mathbf{W}_{z}(F_{sm}(\mathbf{X}^{T}\mathbf{W}_{k}^{T}\mathbf{W}_{q}\mathbf{X})\mathbf{W}_{v}\mathbf{X}).$$

(1)

There are four ($1\times 1$) convolution kernels, i.e., $\mathbf{W}_{z}$,$\mathbf{W}_{k}$, $\mathbf{W}_{q}$, and $\mathbf{W}_{v}$, that learns the linear combination of spatial features among different channels. Within the same channels, the $HW\times HW$ outer-product between $\mathbf{W}_{k}\mathbf{X}$ and $\mathbf{W}_{q}\mathbf{X}$ activates any features at different spatial locations that have a similar intensity. The joint activation mechanism of spatial features is very likely to highlight the spatial noise. The only actual weights, $\mathbf{W}$s, are channel-specific instead of spatial-specific, making the Non-Local attention exceptionally redundant at the huge memory-consumption of the $HW\times HW$ matrix. For efficient computation, reduction of NL leads to many possibilities: Low rank approximation of $\mathbf{A}$ (EA), Channel-only self-attention $\mathbf{A}^{ch}\in\Re^{C\times 1\times 1}$ that highlight the same global context for all pixels(GC [3] and SE [19] ), Spatial-only self-attention $\mathbf{A}^{sp}\in\Re^{1\times W\times H}$ not powerful enough to be recognized as a standalone model, Channel-spatial composition $\mathbf{A}^{sp}$, where the parallel composition: $\mathbf{Z}=\mathbf{A}^{ch}\odot^{ch}\mathbf{X}+\mathbf{A}^{sp}\odot^{sp}\mathbf{X}$ and the sequential composition: $\mathbf{Z}=\mathbf{A}^{ch}\odot^{ch}(\mathbf{A}^{sp}\odot^{sp}\mathbf{X})$ introduce different order of non-linearity. Different conclusions were empirically drawn, such as CBAM [48] (sequential$>$parallel) and DA [14] (parallel$>$sequential), which partially indicates that the intended non-linearity of the tasks are not fully modeled within the attention blocks.

These issues are typical examples of general attention design that does not target the pixel-wise regression problem. With the help of Table 1, we re-visit critical design aspects of existing attention blocks and raise challenges on how to achieve both channel-specific and spatial-specific weighting for pixel-wise regression. (All the attention blocks are compared with their top-performance configurations.)

Internal Attention Resolution. Recall that most pixel-wise regression DCNNs use the same backbone networks, e.g., ResNet, as the classification (i.e., image recognition) and coordinate regression(i.e. bbox detection, instance segmentation) tasks. For robustness and computational efficiency, these backbones produce low-resolution features, for instance $1\times 1\times 512$ for the classification and $[W/r,H/r]$ for bbox detection, where $r$ is the longest side pixels of the smallest object bounding box. Pixel-wise regression cannot afford such loss of resolution, especially because the highly non-linearity in object edges and body parts are very difficult to encode in low-resolution features [4, 44, 40].

Using these backbones in pixel-wise regression, self-attention blocks are expected to preserve high-resolution semantics in attention computation. However, in Table 1, all the reductions of NL reach their top performance at a lower internal resolution. Since their performance metrics are far from perfect, the natural question to ask is: are there better non-linearity that could leverages higher resolution information in attention computation?

Output Distribution/Non-linearity. In DCNNs for pixel-wise regression, outputs are usually encoded as 3D tensors. For instance, the 2D keypoint coordinates are encoded as a stack of 2D Gaussian maps $[\#keypoint\_type\times W\times H]$. The pixel-wise class indices are encoded as a stack of binary maps $[\#semantic\_classes\times W\times H]$ which follows the Binormial distribution. Non-linearity that directly fits the distribution upon linear transformations (such as convolution) could potentially alleviate the learning burden of DCNNs. The natural nonlinear functions to fit the above distributions are SoftMax for 2D Gaussian maps, and Sigmoid for 2D Binormial Distribution. However, none of the existing attention blocks in Table 1 contains such a combination of nonlinear functions.

Our solution to the above challenges is to conduct “polarized filtering” in attention computation. A self-attention block operates on an input tensor $\mathbf{X}$ to highlight or suppress features, which is very much like optical lenses filtering the light. In photography, there are always random lights in transverse directions that produce glares/reflections. Polarized filtering, by only allowing the light pass orthogonal to the transverse direction, can potentially improve the contrast of the photo. Due to the loss of total intensity, the light after filtering usually has a small dynamic range, therefore needs a additional boost, e.g. by High Dynamic Range (HDR), to recover the details of the original scene.

We borrow the key factors of photography, and propose the Polarized Self-Attention (PSA) mechanism: (1) Filtering: completely collapse features in one direction while preserving high-resolution in its orthogonal direction; (2) HDR: increase the dynamic range of attention by Softmax normalization at the bottleneck tensor (smallest feature tensor in attention block), followed by tone-mapping with the Sigmoid function. Formally, we instantiate the PSA mechanism as a PSA block below (also see diagram in Figure 2):

Channel-only branch $A^{ch}(\mathbf{X})\in\Re^{C\times 1\times 1}$:

$$\mathbf{A}^{ch}(\mathbf{X})=F_{SG}\Big{[}\mathbf{W}_{z|\theta_{1}}\Big{(}(\sigma_{1}(\mathbf{W}_{v}(\mathbf{X}))\times F_{SM}(\sigma_{2}(\mathbf{W}_{q}(\mathbf{X})))\Big{)}\Big{]},$$

(2)

where $\mathbf{W}_{q},\mathbf{W}_{v}$ and $\mathbf{W}_{z}$ are $1\times 1$ convolution layers respectively, $\sigma_{1}$ and $\sigma_{2}$ are two tensor reshape operators, and $F_{SM}(\cdot)$ is a SoftMax operator and ”$\times$” is the matrix dot-product operation $F_{SM}(\mathbf{X})=\sum_{j=1}^{N_{p}}{\frac{e^{x_{j}}}{\sum_{m=1}^{N_{p}}{e^{x_{m}}}}{x_{j}}}$. The internal number of channels, between $\mathbf{W}_{v}|\mathbf{W}_{q}$ and $\mathbf{W}_{z}$, is $C/2$. The output of channel-only branch is $\mathbf{Z}^{ch}=\mathbf{A}^{ch}(\mathbf{X})\odot^{ch}\mathbf{X}\in\Re^{C\times H\times W}$, where $\odot^{ch}$ is a channel-wise multiplication operator.

Spatial-only branch $A^{sp}(\mathbf{X})\in\Re^{1\times H\times W}$:

$$\mathbf{A}^{sp}(\mathbf{X})=F_{SG}\Big{[}\sigma_{3}\Big{(}F_{SM}(\sigma_{1}(F_{GP}(\mathbf{W}_{q}(\mathbf{X}))))\times\sigma_{2}(\mathbf{W}_{v}(\mathbf{X}))\Big{)}\Big{]},$$

(3)

where $\mathbf{W}_{q}$ and $\mathbf{W}_{v}$ are standard $1\times 1$ convolution layers respectively, $\theta_{2}$ is an intermediate parameter for these channel convolutions, and $\sigma_{1}$, $\sigma_{2}$ and $\sigma_{3}$ are three tensor reshape operators, and $F_{SM}(\cdot)$ is the SoftMax operator. $F_{GP}(\cdot)$ is a global pooling operator $F_{GP}(\mathbf{X})=\frac{1}{H\times W}\sum_{i=1}^{H}\sum_{j=1}^{W}X(:,i,j)$, and $\times$ is the matrix dot-product operation. The output of spatial-only branch is $\mathbf{Z}^{sp}=\mathbf{A}^{sp}(\mathbf{X})\odot^{sp}\mathbf{X}\in\Re^{C\times H\times W}$, where $\odot^{sp}$ is a spatial-wise multiplication operator.

Composition: The outputs of above two branches are composed either under the parallel layout

	$\displaystyle PSA_{p}(\mathbf{X})$	$\displaystyle=$	$\displaystyle\mathbf{Z}^{ch}+\mathbf{Z}^{sp}$
		$\displaystyle=$	$\displaystyle\mathbf{A}^{ch}(\mathbf{X})\odot^{ch}\mathbf{X}+\mathbf{A}^{sp}(\mathbf{X})\odot^{sp}\mathbf{X},$

or under the sequential layout

	$\displaystyle PSA_{s}(\mathbf{X})$	$\displaystyle=$	$\displaystyle\mathbf{Z}^{sp}(\mathbf{Z}^{ch})$
		$\displaystyle=$	$\displaystyle\mathbf{A}^{sp}(\mathbf{A}^{ch}(\mathbf{X})\odot^{ch}\mathbf{X})\odot^{sp}\mathbf{A}^{ch}(\mathbf{X})\odot^{ch}\mathbf{X}.$

where ”+” is the element-wise addition operator.

Relation of PSA to other Self-Attentions: We add PSA to Table 1 and make the following observations:

• •

  Internal Resolution vs Complexity: Comparing to existing attention blocks under their top configuration, PSA preserves the highest attention resolution for both the channel ($C/2$)³³3$C/2$ is the smallest channel number when PSA produces the best metrics, and is used throughout our experiments. and spatial ($[W,H]$) dimension.

  Moreover, in our channel-only attention, the Softmax re-weighting is fused with squeeze-excitation leveraging Softmax as the nonlinear activation at the bottleneck tensor of size $C/2\times W\times H$. The channel numbers $C$-$C/2$-$C$ follow a squeeze-excitation pattern that benefited both GC and SE blocks. Our design conducts higher-resolution squeeze-and-excitation while at comparable computation complexity of the GC block.

  Our spatial-only attention not only keeps the full $[W,H]$ spatial resolution, but also internally keeps $2\times C\times C/2$ learnable parameters in $W_{q}$ and $W_{v}$ for the nonlinear Softmax re-weighting, which is more powerful structure than existing blocks. For instance, the spatial-only attention in CBAM is parameterized by a $7\times 7\times 2$ convolution (a linear operator), and EA learns $C\times d_{k}+C\times d_{v}$ parameters for linear re-weighting ($d_{k},d_{v}\ll C$ ).

• •

  Output Distribution/Non-linearity. Both the PSA channel-only and spatial-only branches use a Softmax-Sigmoid composition. Considering the Softmax-Sigmoid composition as a probability distribution function, both the multi-mode Gaussian maps (keypoint heatmaps) and the piece-wise Binomial maps (segmentation masks) can be approximated upon linear transformations, i.e. $1\times 1$ convolutions in PSA. We therefore expect the non-linearity could fully leverage the high resolution information preserved within in PSA attention branches.

We first add PSA blocks to standard baseline networks of the following tasks.

Top-Down 2D Human Pose Estimation: Among the DCNN approaches for 2D human pose estimation, the top-down approaches generally dominate the top metrics. This top-down pipeline consists of a person bounding box detector and a keypoint heatmap regressor. Specifically, we use the pipelines in [51] and [40] as our baselines. An input image is first processed by a human detector [51] of $56.4$AP (Average Precision) on MS-COCO val2017 dataset [28]. Then all the detected human image patches are cropped from the input image and resized to $384\times 288$. Finally, the $384\times 288$ image patches are used for keypoint heatmap regression by a single person pose estimator. The output heatmap size is $96\times 72$.

We add PSA on Simple-Baseline [51] with the Resnet50/152 backbones and HRnet [40] with the HRnet-w32/w48 backbones. The results on MS-COCO val2017 are shown in Table 2. PSA boosts all the baseline networks by $2.6$ to $4.3$ AP with minor overheads of computation (Flops) and the number of parameters(mPara). Even without ImageNet pre-training, PSA with “Res50” backbone gets $76.5$ AP, which is not only $4.3$ better than Simple-Baseline with Resnet50 backbone, but also better than Simple-Baseline even with Resnet152 backbone. A similar benefit is also observed on PSA with HRNet-W32 backbone outperforms the baseline with “HR-w48” backbone. This giant performance gains of PAS and the small overheads make PSA+HRNet-W32 the most cost-effective model among all models in Table 2.

Semantic Segmentation. This task maps an input image to a stack of segmentation masks, one output mask for one semantic class. In Table 3, we compare PSA with the DeepLabV3Plus [4] baseline on the Pascal VOC2012 Aug [12] (21 classes, input image size $513\times 513$, output mask size $513\times 513$). PSA boosts all the baseline networks by $1.8$ to $2.6$mIoU(mean Intersection over Union) with minor overheads of computation (Flops) and the number of parameters (mPara). PSA with “Res50” backbone got $79.0$ mIoU, which is not only $1.8$ better than the DeepLabV3Plus with the Resnet50 backbone, but also better than DeepLabV3Plus even with Resnet101.

We then apply PSA to the current state-of-the-arts of above tasks. Top-down 2D Human Pose Estimation. To our knowledge, the current state-of-the-art results by single models were achieved by UDP-HRnet with 65.1mAP bbox detector on the MS-COCO keypoint testdev set. In Table 4, we add PSA to the UDP-Pose with HRnet-W48 backbone and achieve a new state-of-the-art AP of $79.5$. PSA boosts UDP-Pose (baseline) by $1.7$ points (see Figure 3 (a) for their qualitative comparison).

Note that there is only a subtle metric difference between the parallel (p) and sequential(s) layout of PSA. We believe this partially validate that our design of the channel-only and spatial-only attention blocks has exhausted the representation power along the channel and spatial dimension.

Semantic Segmentation. To our knowledge, the current state-of-the-art results by single models were produced by HRNet-OCR(MA) [41] on the Cityscapes validation set  [9](19 classes, input image size $1024\times 2048$, output mask size $1024\times 2048$). In Table 5, we add PSA to the basic configuration of HRNet-OCR and achieve the new state-of-the-arts mIoU of $86.95$. PSA boosts HRNet-OCR (strong baseline) by $2$ points(see Figure 3 (b) for their qualitative comparison). Again that there is only a subtle metric difference between the PSA results under the parallel(p) layout and the sequential(s) layout.

In Table 6, we conduct an ablation study of PSA configurations on Simple-Baseline(Resnet50) [51] and compare PSAs with other related self-attention methods. All the overheads, such as Flops, mPara, inference GPU memory(”Mem.”), and inference time (”Time”)) are inference costs of one sample. To reduce the randomness in CUDA and Pytorch scheduling, we ran inference on MS-COCO val2017 using 4 TITAN RTX GPUs, batchsize 128 (batchsize 32/GPU), and averaged over the number of samples.

From the results of ”PSA ablation” in Table 6, we observe that (1) the channel-only block ($\mathbf{A}^{ch}$) outperform spacial-only attention ($\mathbf{A}^{sp}$), but can be further boosted by their parallel ([$\mathbf{A}^{ch}|\mathbf{A}^{sp}$]) or sequential ($\mathbf{A}^{sp}(\mathbf{A}^{ch})$) compositions; (2) The parallel ([$\mathbf{A}^{ch}|\mathbf{A}^{sp}$]) or sequential ($\mathbf{A}^{sp}(\mathbf{A}^{ch})$) compositions has similar AP, Flops, mPara, inference memory(Mem.), and inference (Time.).

From the results of ”related self-attention methods”, we observe that (1) the NL block costs the most memory while produces the least boost ($2.3$AP) over the baseline, indicating that NL is highly redundant. (2) The channel-only attention GC is better than SE since it includes SE. GC is even better than channel+spatial attention CBAM because the inner-product-based attention mechanism in GC is more powerful than the convolution/MLP-based CBAM. (3) PSA $\mathbf{A}^{ch}$ is the best channel-only attention block over GC and SE. We believe PSA benefits from its highest channel resolution ($C/2$) and its output design. (4) The channel+spatial attention CBAM with a relatively early design is still better than the channel-only attention SE. (5) Under the same sequential layout of spatial and channel attention, PSA is significantly better than CBAM. Finally, (6) At similar overheads, both the parallel and sequential PSAs are better than the compared blocks.