[论文解读] Polarized Self-Attention: Towards High-quality Pixel-wise Regression

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

Advanced Neural Network Applications参考文献 59被引用 152

一句话总结

该论文提出极化自注意力（PSA），一个插件块，在通道和空间分支中保持高内部分辨率，并使用 Softmax-Sigmoid 非线性来更好地拟合像素级回归输出，从而在姿态估计和分割基准上提升性能。

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.

研究动机与目标

• 激发并解决需要在适度成本下保持高分辨率特征的像素级回归挑战。
• 开发一个在通道和空间维度上都保持高内部分辨率的注意力块。
• 设计一个匹配典型像素级回归输出（高斯热图、二值掩模）的非线性（Softmax-Sigmoid）。
• 证明在把 PSA 插入标准基线和最先进模型时能够带来稳定的改进。

提出的方法

• 提出带有两个极化分支的 PSA：一个仅通道分支和一个仅空间分支，每个分支都保持高分辨率的注意力计算。
• 在注意力中维持高通道分辨率(C/2)和完整的空间分辨率[W, H]，同时在相反维度对输入进行压缩。
• 在瓶颈处使用 Softmax 的同时采用 Sigmoid 色调映射（Softmax-Sigmoid 组合）以拟合二维高斯热图和二维二项分割映射。
• 提供并行形式（A^ch + A^sp）与顺序形式（A^sp(A^ch(X))）的组合，性能差异很小。
• 在ResNet、HRNet等骨干网络的残差块内的第一个3x3卷积之后插入 PSA，开销极小。
• 在二维人体姿态估计（MS-COCO）和语义分割（PASCAL VOC2012、Cityscapes）上评估 PSA 相对于基线和最先进模型的表现。

实验结果

研究问题

• RQ1 一个具有极化滤波的高分辨率注意力块是否能够在不过度增加计算的情况下改进像素级回归？
• RQ2 在注意力中保持通道和空间分辨率，并结合分布匹配的非线性，是否能在关键点热图和分割掩模上获得可测量的增益？
• RQ3 并行和顺序的 PSA 布局在效果或成本上是否存在显著差异，还是差异不大？
• RQ4 将 PSA 应用于强基线和 SOTA 模型时，在二维姿态估计与分割任务上表现如何？
• RQ5 将 PSA 插入常见的 DCNN 头部时，准确性增益与计算/内存开销之间的权衡如何？

主要发现

• PSA 通过大约 2–4 个 AP 提升姿态估计任务的基线性能。
• PSA 通过大约 1–2 个 AP 提升二维姿态估计基准的最先进模型。
• 在语义分割中，PSA 在测试的骨干网络上实现了 +1.8–2.6 mIoU 的提升，并在 HRNet-OCR+PSA 配置下达到城市景观数据集的新状态（86.95 mIoU）。
• PSA 保持更高的内部注意力分辨率（通道：C/2；空间：[W, H]）并使用 Softmax-Sigmoid 非线性来更好地拟合高斯热图和二项分割映射。
• 并行和顺序的 PSA 布局在性能上仅呈现边际差异， FLOPs 与参数量相近。
• 消融研究表明 PSA 的通道分支（A^ch）通常优于空间分支（A^sp），而组合（PSA_p 或 PSA_s）优于单独的分支。