从聚类角度重新审视 Mask Transformer

1726673021914.jpg

全景分割是一个计算机视觉问题,是许多实际应用的核心任务。由于其复杂性,先前的工作通常将全景分割分为语义分割(为图像中的每个像素分配语义标签,例如“人”和“天空”)和实例分割(识别和分割图像中的可数对象,例如“行人”和“汽车”),并进一步将其划分为多个子任务。每个子任务单独处理,并使用额外的模块来合并每个子任务阶段的结果。这个过程不仅复杂,而且在处理子任务和组合不同子任务阶段的结果时会引入许多手工设计的先验。

最近,受到Transformer和DETR的启发, MaX-DeepLab提出了一种使用mask transformers(Transformer架构的扩展,用于生成分割mask)进行全景分割的端到端解决方案。该解决方案采用像素路径(由卷积神经网络或视觉Transformer组成)提取像素特征,记忆路径(由Transformer解码器模块组成)提取记忆特征,以及双路径Transformer进行像素特征和记忆特征的交互。然而,利用交叉注意力的双路径Transformer最初是为语言任务设计的,输入序列由几十到几百个单词组成。然而,当涉及到视觉任务,特别是分割问题时,输入序列由数万个像素组成,这不仅意味着输入规模大得多,而且与语言单词相比,它代表了较低级别的嵌入。

在CVPR 2022上发表的 “ CMT-DeepLab:用于全景分割的聚类掩模变换器”和即将在ECCV 2022上发表的“ kMaX-DeepLab:k 均值掩模变换器”中,我们建议从聚类的角度重新解释和重新设计交叉注意力(即将具有相同语义标签的像素分组在一起),以更好地适应视觉任务。 CMT-DeepLab 建立在之前最先进的方法 MaX-DeepLab 之上,采用像素聚类方法执行交叉注意力,从而得到更密集、更合理的注意力图。 kMaX-DeepLab 进一步重新设计了交叉注意力,使其更像k 均值聚类算法,并对激活函数进行了简单的更改。我们证明 CMT-DeepLab 实现了显著的性能提升,而 kMaX-DeepLab 不仅简化了修改,而且还在不增加测试时间的情况下大幅提升了最新水平。我们还很高兴地宣布在 DeepLab2 库中开源发布我们性能最佳的分割模型 kMaX- DeepLab。

概述

我们并不打算将交叉注意力直接应用于视觉任务,而是从聚类的角度对其进行重新解释。具体来说,我们注意到 mask Transformer 对象查询可以被视为聚类中心(旨在将具有相同语义标签的像素分组),交叉注意力的过程类似于 k-means 聚类算法,它采用以下迭代过程:(1)将像素分配给聚类中心,其中多个像素可以分配给单个聚类中心,某些聚类中心可能没有分配的像素;(2)通过平均分配给同一聚类中心的像素来更新聚类中心(如果没有像素分配给聚类中心,则不会更新聚类中心)。

在 CMT-DeepLab 和 kMaX-DeepLab 中,我们从聚类的角度重新表述了交叉注意力,它由迭代聚类分配和聚类更新步骤组成。

鉴于 k-means 聚类算法的流行,我们在 CMT-DeepLab 中重新设计了交叉注意力,以便将空间方面的softmax运算(即沿图像空间分辨率应用的 softmax 运算)实际上将聚类中心分配给像素,改为沿聚类中心应用。在 kMaX-DeepLab 中,我们进一步将空间方面的 softmax 简化为聚类方面的argmax(即沿聚类中心应用 argmax 运算)。我们注意到 argmax 运算与 k-means 聚类算法中使用的硬分配(即,一个像素仅分配给一个聚类)相同。

从聚类角度重新表述掩码转换器的交叉注意机制可显著提高分割性能,并简化复杂的掩码转换器流程,使其更易于解释。首先,使用编码器-解码器结构从输入图像中提取像素特征。然后,使用一组聚类中心对像素进行分组,并根据聚类分配进一步更新这些像素。最后,迭代执行聚类分配和更新步骤,最后的分配直接用作分割预测。

为了将典型的 mask Transformer 解码器(由交叉注意力、多头自注意力和前馈网络组成)转换为我们提出的 k-means 交叉注意力,我们只需将空间方面的 softmax 替换为聚类方面的 argmax。

我们提出的 kMaX-DeepLab 的元架构由三个组件组成:像素编码器、增强像素解码器和 kMaX 解码器。像素编码器是任何网络主干,用于提取图像特征。增强像素解码器包括用于增强像素特征的变换器编码器和用于生成更高分辨率特征的上采样层。一系列 kMaX 解码器将聚类中心转换为 (1) 掩码嵌入向量,该向量与像素特征相乘以生成预测掩码,以及 (2) 每个掩码的类预测。

kMaX-DeepLab 的元架构。

结果

我们使用全景质量(PQ) 指标在两个最具挑战性的全景分割数据集COCO和Cityscapes上对 CMT-DeepLab 和 kMaX-DeepLab 进行了评估,并与 MaX-DeepLab 和其他最先进的方法进行了比较。CMT-DeepLab 实现了显着的性能提升,而 kMaX-DeepLab 不仅简化了修改,而且还大幅推动了最先进的技术,在 COCO 验证集上的 PQ 为 58.0%,在 Cityscapes 验证集上的 PQ 为 68.4%,mask 平均精度(mask AP) 为 44.0%,平均交并比(mIoU) 为 83.5%,并且无需测试时间增强或使用外部数据集。

方法 帕金森病

MaX-DeepLab 51.1%(-6.9%)

面膜成型机 52.7%(-5.3%)

K网 54.6%(-3.4%)

CMT-DeepLab 55.3%(-2.7%)

kMaX-DeepLab 58.0%

在 COCO 验证集上进行比较。

方法 帕金森病 AP掩码 平均

Panoptic-DeepLab 63.0%(-5.4%) 35.3%(-8.7%) 80.5%(-3.0%)

Axial-DeepLab 64.4%(-4.0%) 36.7%(-7.3%) 80.6%(-2.9%)

宽带网 66.4%(-2.0%) 40.1%(-3.9%) 82.2%(-1.3%)

kMaX-DeepLab 68.4% 44.0% 83.5%

与 Cityscapes val 集进行比较。

kMaX-DeepLab 从聚类角度进行设计,不仅性能更高,而且注意力图的可视化也更合理,便于理解其工作机制。在下面的示例中,kMaX-DeepLab 迭代执行聚类分配和更新,从而逐步提高掩码质量。

kMaX-DeepLab 的注意力图可以直接可视化为全景分割,这为模型工作机制提供了更好的合理性(图片来源:coco_url和许可证)。

结论

我们展示了一种更好地为视觉任务设计掩码转换器的方法。通过简单的修改,CMT-DeepLab 和 kMaX-DeepLab 重新制定了交叉注意力机制,使其更像聚类算法。因此,所提出的模型在具有挑战性的 COCO 和 Cityscapes 数据集上实现了最先进的性能。我们希望DeepLab2 库中 kMaX-DeepLab 的开源版本将促进未来设计特定于视觉的转换器架构的研究。

致谢

我们感谢 Huiyu Wang、Dahun Kim、Siyuan Qiao、Maxwell Collins、Yukun Zhu、Florian Schroff、Hartwig Adam 和 Alan Yuille 的宝贵讨论和支持。

版权声明

本文仅代表作者观点,不代表本站立场。
本文系作者授权发表,未经许可,不得转载。

评论