CVPR 2022 | 从自注意力中学习语义Affinity，用于端到端弱监督语义分割

科技

2022-06-18 15:06

©PaperWeekly 原创 · 作者 | 邵琪钧
研究方向 | 计算机视觉、深度学习

论文标题：

Learning Affinity from Attention: End-to-End Weakly-Supervised SemanticSegmentation with Transformers

论文地址：

https://arxiv.org/abs/2203.02664

代码地址：

https://github.com/rulixiang/afa

推荐一个自己的深度学习项目：主要用于论文复现，包含图像识别、目标检测、图像分割、自监督等领域。每个项目可以独立运行，并且有对应论文解读。欢迎 star：

https://github.com/KKKSQJ/DeepLearning

前言

本文是一篇基于弱监督，并且使用 Transformer 来进行端到端的语义分割算法。目前基于 image-level 的弱监督语义分割基础都是建立一个良好的类激活图（CAM），而后利用像素间的关系进行建模（pixel-level 的 affinity），从而在 CAM 基础上生成伪标签，再对伪标签进行细化，作为语义分割模型的监督信息。目前的方法主要基于卷积神经网络（CNN），由于 CNN 无法很好的挖掘全局信息，所以，往往导致对象区域不完整。如下图（a）所示。

为了解决这个问题，本文引入自然集成全局信息的 Transformer，为端到端的弱监督语义分割（WSSS）生成完整的初始伪标签。基于 Transformer 的自注意力与语义 affinity 之间的内在一致性，本文提出一个从注意力中学习 affinity 的模块此外（affinity from Attention，AFA）。为了有效的获得可靠的 affinity 标签来监督 AFA 并保证伪标签的局部一致性，作者设计了一个像素自适应细化模块（PAR），该模块融合低级图像外观信息来细化标签。实验证明，该方法在 PASCAL VOC 2012 和 MS COCO 2014 数据集上分别实现了 66.0% 和 38.9% 的 mIoU，显著优于近期的端到端方法和几个多阶段算法。

相关介绍

弱监督语义分割（WSSS）：目前的基于 image-level 的弱监督语义分割基础都是建立于一个良好的类激活图（CAM），而后利用像素间的关系进行建模（pixel-level 的 affinity），从而在一个 CAM 的基础上得到一些良好的结果，进而为语义分割生成伪标签。

先前的一些方法：

1. 双阶段方法：第一阶段训练一个分类网络，利用 CAM 来产生初始伪标签。伪标签经过细化之后，作为第二阶段的监督信息，去训练一个独立的分割模型。（该方法流程复杂化，并且效率低下）

2. 端到端方法：目前大多数端到端的方法基于 CNN，利用 CNN 来产生 CAM，进而生成伪标签。然而 CNN 无法很好的获取全局信息，从而导致无法获取完整的对象区域，从而影响生成的伪标签的质量。

视觉中的 Transformer：Vision Transformer 在论文 VIt 中被提出，将 Transformer 应用到图像中，并取得了显著的成功。由于 Transformer 的自注意力机制，它能很好的获取图像的全局信息。

动机&贡献

3.1 动机

处理弱监督语义分割大致分为两类方法：双阶段方法和端到端方法。其中，双阶段方法流程复杂，效率低下。有必要提出一个端到端的方法。

端到端方法大都基于 CNN，但是 CNN 无法很好的获取全局信息，从而导致无法获取完整的目标信息（如下图（a）所示），从而导致生成的伪标签效果不佳。

Transforerm 的固有优点是可以捕获全局信息，但是，Transformer 自然（不加约束）产生的 CAM 是粗糙的（如上图（c）所示），需要对其进一步细化。

3.2 贡献

本文提出一个端到端的基于 Transformer 的弱监督语义分割框架。据我们所知，这是第一个在 WSSS 领域探索 transformer 的工作。

作者利用 Transformer 的固有优点，设计了affinity from attention 模块（AFA），它旨在从 MHSA 中学习可靠的语义 affinity，并利用学习到的 affinity 去修正初始伪标签。

本文提出一个高效的像素自适应细化模块（PAR），该模块融合了局部像素的 RGB 和位置信息，用于标签细化。

方法

4.1 Transformer Backbone

如图 3 所示，该框架使用 Transformer 作为 backbone。输入一张图片，首先划分为个 patchs，然后经过线性层得到相同数量的 patch tokens。在每一个 Transformer block 中，多头注意力（multi-head self-attention (MHSA) ）将会捕获全局特征。具体来说，对于第个 head，patch tokens 经过 MLP 被映射成 queries ，keys ，values ，其中表示特征维度。在的基础之上，自注意力矩阵和输出的计算如公式 1 所示：

最终每个 Transformer block 的输出是将送入前向传播网络（feed-forward layers (FFN)）得到。

例如：。其中表示 concatenation，由 Layer Norm 和 MLP 组成。

通过堆叠多个 Transformer block，backbone 产生 feature maps。

4.2 CAM Genaration

使用类激活图（class activation maps (CAM) ）作为初始伪标签。

对于从 backbone 提取的 feature maps ，对于一个类别，它的类激活图将会根据它的第 c 个类别权重（由分类层产生类别权重）和 feature maps中第 c 个通道的特征相乘生成。具体如公式 2 所示：

其中，RELU 函数用于去除负激活。随后，经过 Min-Max 归一化到 [0,1] 区间。用 backbone score 来区域前景和背景区域。

4.3 Affinity from Attention

如上图所示，Transformer 中的 MHSA 和语义 affinity 之间存在一致性。这启示我们使用 MHSA 去发现目标区域。但是，在训练过程中，不受额外约束的自注意矩阵学习到的 affinity 是粗糙并且不精确的。这意味着直接应用 MHSA 作为 affinity 去细化伪标签是不行的，如图 2（c）所示。

因此，提出 Affinity from Attention module（AFA）去解决这个问题。

假设 Transformer block 中的 MHSA 用表示。其中，是 flattened spatial size。是 attention heads 的数量。在 AFA 模块中，直接使用 MLP 来编码 MHSA 以此来产生语义 affinity。

值得注意的是：自注意力机制是一个定向图模型，而 affinity 矩阵应该是对称，因为共享相同语义的节点应该是平等的。

为了执行上述变换，作者将 S 和 S 的转置相加。如公式 3 所示：

表示预测的语义 affinity。

Pseudo Affinity Label Generation：为了学到良好的语义 affinity ，一个关键的步骤是获取一个可信赖的伪 affinity 标签作为监督信息。

如上图所示，Transformer 提取的特征通过 CAM 生成初始伪标签，初始伪标签经过 PAR 模块进行细化得到细化后的伪标签（refined pseudo labels）。（PAR 模块再下一节进行介绍）

作者使用两个 background scores 去过滤 refined pseudo label，得到可靠的前景，背景以及不确定区域。

给定 CAM ，伪标签的生成如公式 4 所示：

其中，0 和 255 代表背景和 ignore 区域的索引。表示提取的类别索引。

随后，pseudo affinity label 由生成。具体来说，对于，如果像素和像素共享相同语义，则设置它们的 affinity 为 positive。反之，设置 affinity 为 negative。如果像素来自于 ignore 区域，则它们的 affinity 也被忽略。除此之外，只考虑同一个局部窗口下的像素和像素，并且忽视远处像素对的 affinity。

Affinity Loss：生成的伪 affinity 标签去监督网络预测的 affinity 。affinity loss 如公式 5 所示：

其中，分别表示在进行采样的 positive 和 negative 集合。分别表示的数量。

该损失函数强迫网络从 MHSA 中学习高置信度的语义 affinity 依赖，并且有益于自注意去发现完整的目标区域。

Propagation with Affinity：学习到的语义 affinity 用于修正初始的 CAM，作者利用 random walk 完成这一步骤。具体来说，对于学习到的语义 affinity 矩阵，其对应的语义变换矩阵 T 如公式 6 所示：

其中是一个超参数，用于忽略中平凡的 affinity 值。是一个对角矩阵，用于将归一化。

对于初始 CAM ，rankdom walk 算法如公式 7 所示：

其中，表示将矢量化。这种传播过程，扩散了具有高 affinity 的语义区域，抑制了错误激活的区域，从而使激活图更好地与语义边界对齐。

4.4 Pixel-Adaptive Refinement

如图 3 所示，伪 affinity 标签由初始伪标签生成。但是初始伪标签是粗糙的并且局部不一致。比如：具有相似图像外观的相邻像素不会共享同一个语义。

为了确保局部的一致性，一些方法使用 dense CRF 去细化初始伪标签，但是 dense CRF 不是一个端到端的好选择，因为它的训练效率低下。

本文采用像素自适应卷积（ pixel-adaptive convolution）来提取局部 RGB 信息进行细化，即集合 RGB 和空间信息定义低级的 pairwise affinity 并且构建像素自适应模块（ Pixel-Adaptive Refinement module (PAR).）

具体来说，给定一张图片，对于像素位置 (i,j),(k,l)，其 RGB 和空间对的 kernel 如公式 8 所定义：

其中表示 RGB 信息。表示像素 (i,j) 的空间位置。表示 RGB 的标准差。表示位置的标准差。

然后将经过 softmax 进行归一化，再进行加权，得到作为 PAR 模块的 affinity kernel。如公式 9 所示。

通过这个 affinity kernel，我们同时细化初始的 CAM 以及传播的 CAM。这个细化操作在多个迭代过程被构建。对于 CAM ，在第次迭代，有：

与 CRF 中计算全局像素的 low-level affinity 不同，作者根据像素的 8 个最近邻像素进行计算。该过程可以使用的像素自适应卷积实现，从而能够高效的插入到端到端的训练框架中。同时为了扩大感受野，使用多个空洞卷积进行多个最近邻像素信息的提取。得到像素之间的 low-level affinity 之后，通过多次迭代进行伪标签的修正。

4.5 Network Training

如图 3 所示，训练损失函数包括分类损失、分割损失、正则化损失和前文介绍的 affinity 损失。其中分类损失采用常见的多分类软间隔损失，分割损失采用常见的交叉熵损失：

实验

数据集：

PASCAL VOC2012
MS COCO2014

实验细节：

Transformer backbone：Mix Transformer(MiT)（Segformer）
优化器：AdamW
学习率：
学习率衰减策略：polynomial
数据增强：随机缩放 [0.5,2.0]、随机翻转、随机裁剪 [512,512]
batch size：8
迭代次数：20,000
: (0.55,0.35)
在PAR，空洞卷积的空洞率：[1,2,4,8,12,24]
: (0.3,0.3,0.01)

5.1 消融实验

不同 tok-k 池化对 CAM 的影响。

top-k 池化指的是去每一个 feature map 的 top k% 的值进行池化操作。gmp：global max-pooling（top-100%），gap：global average-pooling（top-1）。对于 transformer based 的分类网络来说，使用 gmp 可以获得更好的性能。

不同模块带来的性能提升