[论文解读] Max Pooling with Vision Transformers reconciles class and shape in weakly supervised semantic segmentation

Simone Rossetti, Damiano Zappia+3 more|arXiv (Cornell University)|Oct 31, 2022

Advanced Neural Network Applications被引用 4

一句话总结

本文提出 ViT-PCM，一种基于视觉变换器（Vision Transformers）的新型端到端弱监督语义分割方法，用基于全局最大池化（GMP）的块-类别映射（PCM）机制替代类激活图（CAM）。通过利用 ViT 的局部特征表示，并使用多标签二元交叉熵损失端到端优化块到类的相关性，ViT-PCM 在 Pascal VOC 2012 验证集上实现了 69.3% 的 mIoU，优于参数更少且无需多阶段优化的现有 CAM 基方法。

ABSTRACT

Weakly Supervised Semantic Segmentation (WSSS) research has explored many directions to improve the typical pipeline CNN plus class activation maps (CAM) plus refinements, given the image-class label as the only supervision. Though the gap with the fully supervised methods is reduced, further abating the spread seems unlikely within this framework. On the other hand, WSSS methods based on Vision Transformers (ViT) have not yet explored valid alternatives to CAM. ViT features have been shown to retain a scene layout, and object boundaries in self-supervised learning. To confirm these findings, we prove that the advantages of transformers in self-supervised methods are further strengthened by Global Max Pooling (GMP), which can leverage patch features to negotiate pixel-label probability with class probability. This work proposes a new WSSS method dubbed ViT-PCM (ViT Patch-Class Mapping), not based on CAM. The end-to-end presented network learns with a single optimization process, refined shape and proper localization for segmentation masks. Our model outperforms the state-of-the-art on baseline pseudo-masks (BPM), where we achieve $69.3\%$ mIoU on PascalVOC 2012 $val$ set. We show that our approach has the least set of parameters, though obtaining higher accuracy than all other approaches. In a sentence, quantitative and qualitative results of our method reveal that ViT-PCM is an excellent alternative to CNN-CAM based architectures.

研究动机与目标

• 为解决基于 CAM 的方法在弱监督语义分割（WSSS）中定位能力差和形状估计不佳的局限性。
• 探究视觉变换器（ViT）是否能通过其固有的特征局部性和全局上下文，绕过 CAM 的瓶颈。
• 开发一种简单、端到端的方法，联合优化伪掩码质量和类别定位，无需多阶段优化。
• 证明全局最大池化（GMP）可增强 ViT 将块特征映射到与类别相关区域的能力，从而提升分割精度。
• 在参数最少且无需后处理增强的情况下，实现在 Pascal VOC 2012 上的 SOTA 性能。

提出的方法

• 提出 ViT-PCM，一种基于视觉变换器的架构，用基于全局最大池化（GMP）的块-类别映射（PCM）头替代 CAM，以提取每类最相关的块特征。
• 采用多标签二元交叉熵（BCE）损失端到端训练 PCM 头，直接从图像级标签优化像素级类别相关性。
• 采用共享权重的双分支结构，以增强平移和尺度等变性，提升伪掩码预测的鲁棒性。
• 在推理阶段应用条件随机场（CRF）对基线伪掩码（BPM）进行后处理，但训练阶段不使用 CRF，保持方法端到端且轻量化。
• 通过可学习的上采样层将 ViT 块特征映射到分割掩码，避免使用复杂的优化头或辅助监督。
• 利用自监督预训练（DINO）提升特征质量，随后使用图像级标签进行端到端微调。

实验结果

研究问题

• RQ1 视觉变换器（ViT）在弱监督语义分割中是否能比使用 CAM 的卷积神经网络（CNN）实现更好的定位和形状估计？
• RQ2 全局最大池化（GMP）是否能有效从 ViT 中提取与类别相关的块特征，从而在无需 CAM 的情况下实现准确的伪掩码生成？
• RQ3 仅使用图像级监督，一个简单、端到端的 ViT 基方法是否能超越复杂的、多阶段的 CAM 基方法？
• RQ4 在 mIoU、参数效率和不同主干网络的鲁棒性方面，所提出的 PCM 机制与 CAM 相比如何？
• RQ5 缺乏多阶段优化或辅助监督（如显著性图）是否会影响性能，还是可通过适当的特征映射避免？

主要发现

• ViT-PCM 仅使用图像级标签和端到端训练，就在 Pascal VOC 2012 验证集上实现了 69.3% 的 mIoU，优于所有先前的 SOTA 方法。
• 结合 CRF 后处理后，mIoU 提升至 71.4%，表明其基线伪掩码质量优异。
• ViT-PCM 在端到端基准上创下新 SOTA，验证集 mIoU 达 70.3%，测试集 mIoU 达 70.9%，在某些情况下甚至优于多阶段方法。
• 该模型参数量显著少于竞争方法，在所有 SOTA 方法中参数量最少，如图 5 所示。
• 在 MS-COCO 2014 数据集上，ViT-PCM 在验证集上实现了 45.03% 的 mIoU，证明其在 Pascal VOC 之外也具备良好的泛化能力。
• 消融实验表明，所有测试主干网络中，ViT-PCM 搭载 PCM 均优于 CAM，尤其在 ViT-S/16 上提升最大（mIoU 从 29.3% 提升至 43.3%），证明 PCM 机制的优越性。