kMaX-DeepLab：k-均值掩码变换器

变压器在视觉任务中的兴起不仅推动了网络骨干设计的进步，还开启了一个全新的篇章，实现了端到端的图像识别（例如，目标检测和全景分割）。变压器架构起源于自然语言处理（NLP），由自注意力机制和交叉注意力机制组成，能够有效地学习序列中元素之间的长距离交互。然而，我们观察到大多数现有的基于变压器的视觉模型简单地借鉴了NLP的思想，忽视了语言和图像之间的重要差异，特别是空间展平像素特征的极长序列长度。这随后阻碍了像素特征与对象查询之间的交叉注意力学习。在本文中，我们重新思考了像素与对象查询之间的关系，并提出将交叉注意力学习重新表述为一个聚类过程。受传统k均值聚类算法的启发，我们开发了一种用于分割任务的k均值掩码变换器（kMaX-DeepLab），该方法不仅提升了现有技术水平，还具有简洁优雅的设计。因此，我们的kMaX-DeepLab在COCO验证集上取得了58.0%的PQ，在Cityscapes验证集上取得了68.4%的PQ、44.0%的AP和83.5%的mIoU，在ADE20K验证集上取得了50.9%的PQ和55.2%的mIoU，这些成绩均未使用测试时增强或外部数据集。我们希望这项工作能够为设计专门针对视觉任务的变压器提供一些启示。TensorFlow代码和模型可在https://github.com/google-research/deeplab2 获取，PyTorch版本的实现也可在https://github.com/bytedance/kmax-deeplab 获取。