ECCV 2022 谷歌新作| 基于k-means的Transformer
-
论文标题：k-means Mask Transformer
论文地址：https://arxiv.org/pdf/2207.04044.pdf
代码地址：https://github.com/google-research/deeplab2
摘要
- 引出主题： 视觉任务中transformer的兴起不仅推进了网络主干设计，还开启了实现端到端图像识别（例如，目标检测和全景分割）的崭新一页。
- 现存问题： 然而，大多数现有的基于transformer的视觉模型只是借用了自然语言处理的思想，忽略了语言和图像之间的关键差异，特别是空间flattened像素特征的超大序列长度。这随后阻碍了像素特征和对象查询之间交叉关注的学习。
- 解决方法： 重新思考了像素和对象查询之间的关系，并提出将交叉注意力学习重新表述为一个聚类过程。受传统k-means聚类算法的启发，为分割任务开发了k-means掩码transformer（kMaX DeepLab），它不仅提高了最先进的技术水平，而且设计简单优雅。
- 是眼结果：kMaX DeepLab在COCO val集上实现58.0%PQ，在Cityscapes val集上实现了68.4%PQ，44.0%AP和83.5%mIoU。
算法
Mask-Transformer-Based Segmentation Framework
transformer已被有效地部署到细分任务。不失一般性，在下面的问题表述中考虑全景分割，它可以很容易地推广到其他分割任务中。
**问题陈述。**全景分割的目的是分割图像I∈RH×W×3转换为一组具有相关语义标签的非重叠掩码：
K 个GT掩码mi∈{0,1}H×W彼此不重叠。并且ci表示掩模mi的类标签。
从DETR和MaX DeepLab开始，全景分割方法转向了一种新的端到端范例。
其中ˆpi（c）表示对应掩码的语义类预测置信度，包括“thing”类、“stuff”类和void类∅。
基于N个对象查询来预测N个掩码，该查询通过transformer解码器（由自注意力和交叉注意力模块组成）从像素特征聚集信息。
由多个transformer解码器更新的对象查询被用作掩码嵌入向量，该向量将与像素特征相乘以产生由N个掩模组成的最终预测Z∈RHW×N。即
其中F∈RHW×D和C∈RN×D分别指最终像素特征和对象查询。D是像素特征和对象查询的通道维度。
Relationship between Cross-Attention and k-means Clustering
尽管基于transformer的分割框架成功地将对象查询和掩码预测以端到端的方式连接起来，但问题的关键在于如何将对象查询从可学习的嵌入(随机初始化)转换为有意义的掩码嵌入向量。
交叉注意力。交叉注意模块用于聚集附加像素特征以更新对象查询。正式地说，有
其中C∈RN×D表示具有D个通道的N个对象查询，而ˆC则表示更新的对象查询。使用下标HW表示空间维度上softmax的轴，上标p和c分别表示从像素特征和对象查询投影的特征。Qc∈RN×D，Kp∈RHW×D，Vp∈RHW×D表示查询、键和值的线性投影特征。
如等式（4）所示，在更新对象查询时，将softmax函数应用于图像分辨率（HW），该分辨率通常在数千像素范围内。鉴于像素数量巨大，可能需要多次训练迭代才能学习注意力图，注意力图从一开始的均匀分布开始（因为查询是随机初始化的）。在训练的早期阶段，每个对象查询都很难识别丰富像素中最可区分的特征。这种行为与将transformer应用于自然语言处理任务（例如神经机器翻译非常不同，其中输入和输出序列具有相似的短长度。视觉任务，尤其是分割问题，为有效学习交叉注意力提出了另一个挑战。
讨论。与交叉注意类似，自注意力需要执行沿图像分辨率运行的softmax功能。因此，学习自我注意力的注意力图也可能需要多次训练迭代。有效的替代方法，如轴向注意力或局部注意力，通常应用高分辨率特征图，从而减轻问题，而交叉注意的解决方案仍然是研究问题。
K-means聚类算法。等式(4)交叉注意力计算目标查询与像素之间的亲和力，通过空间分辨率softmax(沿图像分辨率运算)转换为注意图。然后使用注意力映射检索(并相应加权)附属像素特征，以更新对象查询。令人惊讶的是，观察到整个过程实际上类似于经典的 K-means算法 ，其工作原理如下:
其中C∈RN×D，P∈RHW×D和A∈RN×HW分别代表聚类中心、像素特征和聚类分配。
通过比较等式（4）、等式（5）和等式（6），注意到k-means聚类算法是无参数的，因此查询、键和值不需要线性投影。群集中心上的更新不是以残余链接方式进行的。最重要的是，k-means在将亲和力转换为注意力图（即，检索和更新特征的权重）时，采用了基于聚类的argmax（即，沿着聚类维度操作的argmax），而不是基于空间的softmax。
这一观察激励重新表述视觉问题中的交叉注意，特别是图像分割。从聚类的角度来看，图像分割等价于将像素分组到不同的聚类中，其中每个聚类对应于预测的掩码。然而，交叉注意力机制也试图将像素分组到不同的对象查询，相反，它采用了与聚类argmax不同的空间softmax操作。
鉴于k-means的成功，假设在像素聚类方面，基于聚类的argmax比基于空间的softmax更适合操作，因为基于聚类的argmax执行硬分配，并有效地将操作目标从数千像素（HW）减少到仅几个聚类中心（N），这加快了训练收敛并导致更好的性能。
k-means Mask Transformer
在此，首先介绍所提出的k-means掩模变换器的关键组件，即k-means交叉注意。然后，介绍了它的元体系结构和模型实例化。
k-means交叉注意。所提出的k-means交叉注意以类似于k-means聚类的方式重新模拟交叉注意：
比较等式(4)及等式(7)，将空间方向的 softmax 改为集群方向的 argmax。如图1所示，用这样一个简单而有效的改变，可以将典型的transformer解码器转换为kMaX解码器。
与原始交叉注意不同，提出的k-means交叉注意采用了不同的操作（即，分簇argmax）来计算注意图，并且不需要多头机制。然而，作为聚集像素特征以进行聚类中心更新的硬分配，分簇argmax不是可微操作，在训练过程中带来了挑战。作者探索了几种方法（例如，Gumbel Softmax[48]），并发现一种简单的深度监管方案是最有效的。该公式允许对每个kMaX解码器添加深度监控，以便训练k均值交叉注意模块中的参数。
**Meta architecture。**图2显示了kMaX深度结构的元架构，它包含三个主要组件：像素编码器、增强像素解码器和kMaX解码器。
- 像素编码器通过CNN或transformer主干提取像素特征
- 增强像素解码器负责恢复特征图分辨率，并通过变换器编码器[89]或轴向注意增强像素特征。
- kMaX解码器从k-means聚类的角度将对象查询（即聚类中心）转换为掩码嵌入向量。
**Model instantiation。**将整个模型分为两条路径：像素路径和聚类路径，这两条路径负责提取像素特征以及集群中心。图3详细描述了kMaX DeepLab实例化，其中包含两个示例主干。
Pixel path。像素路径由像素编码器和增强像素解码器组成。像素编码器是ImageNet预训练主干，如ResNet、MaX-S（即，具有轴向注意力的ResNet-50）和ConvNeXt。增强像素解码器由几个轴向注意块和瓶颈块组成。
Cluster path。聚类路径包含总共六个kMaX解码器，它们均匀分布在不同空间分辨率的特征图中。具体而言，分别为输出步长32、16和8的像素特征部署了两个kMaX解码器。
**Loss functions。**训练损失函数主要遵循MaX DeepLab的设置。采用相同的PQ样式损失、辅助语义损失、mask-id交叉熵损失和像素级实例区分损失。
实验
