[论文解读] Learning Affinity from Attention: End-to-End Weakly-Supervised Semantic Segmentation with Transformers

Lixiang Ru, Yibing Zhan+2 more|arXiv (Cornell University)|Mar 5, 2022

Advanced Neural Network Applications被引用 15

一句话总结

本文提出了一种新颖的端到端弱监督语义分割框架，采用Transformer模型，引入注意力中的亲和力（AFA）模块，从多头自注意力特征中学习可靠的语义亲和力，并通过随机游走传播机制优化伪标签。该方法在PASCAL VOC 2012上达到66.0%的mIoU，在MS COCO 2014上达到38.9%，在端到端弱监督语义分割任务中创下新的最先进性能纪录。

ABSTRACT

Weakly-supervised semantic segmentation (WSSS) with image-level labels is an important and challenging task. Due to the high training efficiency, end-to-end solutions for WSSS have received increasing attention from the community. However, current methods are mainly based on convolutional neural networks and fail to explore the global information properly, thus usually resulting in incomplete object regions. In this paper, to address the aforementioned problem, we introduce Transformers, which naturally integrate global information, to generate more integral initial pseudo labels for end-to-end WSSS. Motivated by the inherent consistency between the self-attention in Transformers and the semantic affinity, we propose an Affinity from Attention (AFA) module to learn semantic affinity from the multi-head self-attention (MHSA) in Transformers. The learned affinity is then leveraged to refine the initial pseudo labels for segmentation. In addition, to efficiently derive reliable affinity labels for supervising AFA and ensure the local consistency of pseudo labels, we devise a Pixel-Adaptive Refinement module that incorporates low-level image appearance information to refine the pseudo labels. We perform extensive experiments and our method achieves 66.0% and 38.9% mIoU on the PASCAL VOC 2012 and MS COCO 2014 datasets, respectively, significantly outperforming recent end-to-end methods and several multi-stage competitors. Code is available at https://github.com/rulixiang/afa.

研究动机与目标

• 为解决基于CNN的端到端弱监督语义分割方法在捕捉全局上下文信息方面存在的局限性，以及由此导致的物体区域不完整问题。
• 利用Transformer的全局建模能力，生成更完整的初始伪标签，以提升弱监督语义分割的性能。
• 利用多头自注意力与语义亲和力之间的内在一致性，学习可靠的亲和力图，用于伪标签的优化。
• 通过像素自适应的细化模块，融合低层次图像外观特征，确保伪标签的局部一致性和可靠性。
• 建立一种简单、端到端的训练流程，避免多阶段框架的复杂性。

提出的方法

• 提出一种基于Transformer的端到端弱监督语义分割框架，仅使用图像级标签，用视觉Transformer主干网络替代CNN，生成初始类激活图（CAMs）作为伪标签。
• 引入注意力中的亲和力（AFA）模块，从多头自注意力（MHSA）特征中学习可靠的语义亲和力，利用伪标签作为监督信号，优化基于注意力的亲和力预测。
• 采用基于学习到的亲和力的随机游走传播机制，对初始伪标签进行优化，提升物体区域的完整性并抑制误激活。
• 设计像素自适应细化（PAR）模块，融合局部RGB特征与位置编码，对伪标签进行细化，提升其与低层次图像结构的一致性。
• 利用AFA模块输出的亲和力图引导标签传播，确保全局一致性的同时，通过PAR模块保持局部细节。
• 采用端到端方式训练整个模型，无需独立的训练阶段，简化了训练流程。

Figure 1: (a) Image and the query points (denote with ” $\bigstar$ ”) to visualize the attention and affinity maps; (b) the self-attention maps in Transformer blocks only capture coarse semantic-level affinity relations; (c) the learned reliable semantic affinity from self-attention with our propose

实验结果

研究问题

• RQ1 Transformer是否能通过更有效地建模全局上下文，提升端到端弱监督语义分割中初始伪标签的质量，相较于CNN更具优势？
• RQ2 Transformer中固有的自注意力机制是否可被有效利用，以学习可靠的语义亲和力，用于伪标签的优化？
• RQ3 如何对基于注意力的亲和力学习进行监督，以确保其生成准确且可靠的亲和力图，用于标签传播？
• RQ4 融合低层次图像外观特征是否能提升优化后伪标签的局部一致性和准确性？
• RQ5 采用基于注意力的亲和力学习的端到端训练策略，是否在标准基准上优于现有的多阶段和端到端弱监督语义分割方法？

主要发现

• 所提方法在PASCAL VOC 2012验证集上达到66.0%的mIoU，显著优于近期的端到端方法，也超越了多个多阶段方法。
• 在更具挑战性的MS COCO 2014数据集上，该方法在验证集上达到38.9%的mIoU，超越了大多数近期多阶段方法，并在端到端弱监督语义分割任务中创下新的SOTA纪录。
• 在VOC 2012上，该模型性能达到全监督Segformer基线模型的83.8%，表明在弱监督设置下仍具备极强的效率与有效性。
• 消融实验表明，AFA与PAR模块均对性能提升有显著贡献：AFA模块提升了亲和力学习质量，PAR模块增强了伪标签的局部一致性。
• 定性结果表明，与1Stage和AA&LR等基线方法相比，该方法生成的分割图具有更精确的边界对齐和更完整的物体区域。
• 端到端训练策略高效且有效，避免了以往弱监督语义分割方法中常见的复杂多阶段流程。

Figure 2: CAM generated with (b) Transformers activates more integral regions than (a) CNN. Refining CAM with (c) coarse MHSA doesn’t work well, while (d) the learned affinity could remarkably improve the generated CAM.