Multi-class Token Transformer for Weakly Supervised Semantic Segmentation

code：https://github.com/xulianuwa/MCTformer

摘要

本文提出了一种基于Transformer的新框架，用于学习类别特定的对象定位图，并将其作为弱监督语义分割（WSSS）的伪标签。受到标准视觉Transformer中单类别token的关注区域可用于生成类别无关定位图的启发，我们探讨了Transformer模型是否也能够通过学习多个类别token来有效捕捉类别特定的注意力，从而实现更具区分性的对象定位。为此，我们提出了一种多类别token Transformer，称为MCTformer，它使用多个类别token来学习类别token与patch token之间的交互。所提出的MCTformer能够成功生成类别区分度高的对象定位图，这些图是通过不同类别token与patch之间的注意力关系得到的。我们还提出使用patch间的成对关联性（从patch到patch的Transformer注意力中提取）进一步优化定位图。此外，所提出的框架与类激活映射（CAM）方法完全互补，在PASCAL VOC和MS COCO数据集上显著提升了WSSS的效果。这些结果强调了类别token在WSSS中的重要性。

• 在图像处理中，图像会被切分成小块（通常是固定大小的区域），这些小块称为“patch”。
• 每个“patch”会被转换为一个特征向量，这个特征向量就被称为一个“token”。
• 在本文中，除了图像patch生成的token，还有用于表示类别信息的“class token”。这些类别token用于捕捉不同类别的注意力信息，从而帮助模型生成类别特定的定位图。

图一：(a)在之前的视觉变形器[10]中，仅使用一个类token(红色方块)来聚合patch  token(蓝色方块)的信息。学习到的与类token对应的patch关注生成与类无关的定位图。(b)相比之下，所提出的MCTformer使用多个类token来学习类token与patch  token之间的交互。学习到的不同类token的类到patch注意可以生成特定类的对象定位映射。 

1.引言

从Transformer中提取类别特定的注意力是一项具有挑战性的任务。我们认为，现有的基于Transformer的方法存在一个共同的问题，即仅使用一个类别token，这使得在单张图像上精确定位不同的对象变得困难。这有两个主要原因。首先，单一类别token的设计不可避免地会捕捉到来自其他类别对象和背景的上下文信息。换句话说，由于只考虑了一个类别token，模型在学习时会同时捕捉到类别特定和通用的表示，从而导致对象定位的区分度较低且含有噪声。其次，模型仅使用一个类别token来与patch tokens进行交互，以学习数据集中多个不同类别对象的表示。因此，模型的容量不足以实现目标的区分性定位性能。

为了解决这些问题，一个直接的想法是利用多个类别token，这些token负责学习不同类别对象的表示。为此，我们提出了一种多类别token Transformer（MCTformer），其中使用了多个类别特定的token来提取类别特定的Transformer注意力。仅仅通过增加ViT中的类别token数量并不能实现类别特定的目标，因为这些类别token仍然没有特定的含义。为了确保每个类别token能够有效地学习特定类别对象的高级区分表示，我们提出了一种类别感知的训练策略用于多个类别token。具体来说，我们在Transformer编码器输出的类别token上沿着嵌入维度应用平均池化，以生成类别分数，这些分数直接由真实的类别标签进行监督。通过这种设计，成功建立了每个类别token与相应类别标签之间的一对一强连接。通过这种设计，最大的优势在于，不同类别的类别到patch的注意力可以直接作为类别特定的定位图使用。

值得注意的是，在训练过程中学习到的patch-to-patch注意力作为附带产物，无需额外计算就可以作为一种patch级别的成对关联性。这种关联性可以用来进一步优化类别特定的Transformer注意力图，显著提升定位性能。此外，我们还展示了所提出的Transformer框架在应用于patch tokens时，与CAM（类激活映射）方法完全互补（通过同时基于类别token和patch token的表示进行分类学习）。这使得类别token和patch token之间的一致性大大提高，从而显著增强它们所生成的对象定位图的区分能力。

总结来说，我们的主要贡献有三点：

1. 我们提出了利用类别特定的Transformer注意力进行弱监督语义分割。
2. 我们提出了一个有效的Transformer框架，其中包括一种新颖的多类别token Transformer（MCTformer），结合类别感知的训练策略，通过不同类别token的类到patch注意力学习类别特定的定位图。
3. 我们提出使用patch-to-patch Transformer注意力作为patch级别的成对关联性，可以显著优化类别特定的Transformer注意力。此外，所提出的MCTformer能够与CAM机制完全互补，生成高质量的对象定位图。

该方法可以为弱监督语义分割（WSSS）生成高质量的类别特定多标签定位图，并在PASCAL VOC测试集（mIoU为71.6%）和MS COCO数据集（mIoU为42.0%）上取得了新的最先进结果。

2.相关工作

图二：所提出的多类别token Transformer（MCTformer-V1）的概述如下：首先，将输入的RGB图像分割并转换为一系列patch tokens。我们提出学习C个额外的类别token，其中C是类别的数量。将这C个类别token与patch tokens拼接在一起，并加入位置嵌入（Position Embeddings，PE），然后通过连续的L个Transformer编码层。最后，输出的C个类别token通过平均池化生成类别分数。我们从最后K层和多个头部中聚合Transformer的注意力，以生成最终的注意力图，从中可以分别提取类别特定的对象定位图和基于类别到patch以及patch到patch注意力的patch级别成对关联性图。这个patch级别的成对关联性可以用于优化类别特定的Transformer注意力图，从而生成更精确的对象定位图。 

我们提出了一个新的纯基于transformer的框架(MCTformer-V1)来利用来自transformer注意力的类特定对象定位映射。MCTformer-V1的总体体系结构如图2所示。首先将输入的RGB图像分割成不重叠的小块，然后将其转换为一系列小块标记。与仅使用一个类token的传统transformer不同，我们建议使用多个类token。这些类token与补丁token相连接，嵌入位置信息，形成transformer编码器的输入token。在transformer编码器中使用了几个transformer块来提取patch tokens and class tokens. 的特征，我们对最后一层的输出类token应用平均池化来生成类分数，而不是像传统transformer那样使用多层感知(MLP)来进行分类预测

在训练时，为了确保不同的类token可以学习不同的类特定表示，我们采用了3.2节中详细介绍的类感知训练策略。分类损失是在由类token直接产生的类分数和真类标签之间计算的。这样就可以在每个类token和相应的类标签之间建立牢固的连接。在测试时，我们可以从transformer中的类到patch注意中提取特定于类的位置映射。我们进一步聚合来自多层的注意力图，以利用从不同transformer层学习到的互补信息。此外，可以从patch到patch的关注中提取patch级的成对亲和力，以进一步细化类到patch的关注，从而显著改进类特定的定位图。使用特定于类的定位映射作为种子来生成伪标签以监督分割模型。

3. Multi-class Token Transformer

3.2. Class-Specific Transformer Attention Learning

特定于类的多类token注意。我们使用标准的自关注层来捕获token之间的长期依赖关系。更具体地说，我们首先将输入token序列归一化，并通过线性层将其转换为Q∈R ^(C+M)×D, K∈R^ (C+M)×D和V∈R^ (C+M)×D的三元组[10]。我们采用缩放点积注意[37]机制来计算查询和键之间的注意值。每个输出标记是所有tokens的加权和，使用注意力值作为权重，表示为:

 

其中，我们可以得到一个token到令token的注意力映射At2t∈R^ (C+M)×(C+M)，且At2t = softmax(QK^T/√D)。

从全局成对关注图At2t中，我们可以提取类对patch Ac2p∈R^C×M的关注，即类对patch的关注，其中Ac2p = At2t[1: C, C + 1: C + M]，如图2中黄点矩阵所示.每一行表示一个特定类对所有补丁的关注得分。利用这些注意力向量，加上所有补丁的原始空间位置，可以生成与C类相关的定位地图。我们可以从每个转换器编码层提取与类相关的本地化映射。考虑到更高的层学习更高级的判别表示(而更早的层捕获更一般和低级的视觉信息)，我们建议融合最后K转换编码层的类到补丁关注，以探索生成的对象定位图的精度和召回率之间的良好权衡。这个过程的公式是:

其中，A^l mct是从提出的mctform - v1的第L个transformer编码层提取的特定于类的transformer注意力。将融合后的Amct沿两个空间维度进一步采用最小-最大归一化方法进行归一化，生成最终的类特定对象定位Amct∈R^ C×N×N。关于如何选择K的详细结果见图6。

特定于类的注意细化。在之前的研究中[1,38,44]经常使用成对亲和来细化目标定位图。它通常需要额外的网络或额外的层来学习亲和映射。相反，我们从所提出的MCTformer的patch- to-patch关注中提取成对亲和映射，而不需要额外的计算和监督。这是通过提取patch-to-patch的关注点Ap2p∈RM×M来实现的，其中Ap2p = At2t[C + 1: C + M, C + 1: C + M]，如图2中蓝点矩阵所示。patch-to-patch的关注被重塑为一个4D张量:Ap2p∈R N×N×N×N。提取的关联用于进一步细化特定于类的transformer注意。这个过程的公式是:

 其中Amct ref∈R C×N×N是细化后的特定于类的定位图。如表5和图5所示，使用patch级成对亲和度进行细化可以得到更好的对象定位图，并提高了外观的连续性

我们需要确保不同的类tokens可以学习到不同的分类判别信息。为此，我们对输出类tokens使用平均池化产生类分数

图3：概述MCTformer-V2。我们在MCTformer-V1中引入了CAM模块。具体来说，CAM模块由卷积层和全局平均池化(GAP)层组成。它接受来自最后一个transformer编码层的重塑输出补丁标记作为输入，并输出类分数。对于MCTformer-V1，我们也使用输出类令牌来生成类分数。因此，整个模型通过分别应用于两种类型的分类预测的两种分类损失来优化。

在推理时，我们融合了类特定的transformer注意(MCT注意)和PatchCAM映射。通过从patch-to-patch transformer关注中提取的patch亲和度进一步细化结果，生成最终的目标定位图。