Vision Transformers (ViTs) are built on the assumption of treating image patches as ``visual tokens" and learn patch-to-patch attention. The patch embedding based tokenizer has a semantic gap with respect to its counterpart, the textual tokenizer. The patch-to-patch attention suffers from the quadratic complexity issue, and also makes it non-trivial to explain learned ViTs. To address these issues in ViT, this paper proposes to learn Patch-to-Cluster attention (PaCa) in ViT. Queries in our PaCa-ViT starts with patches, while keys and values are directly based on clustering (with a predefined small number of clusters). The clusters are learned end-to-end, leading to better tokenizers and inducing joint clustering-for-attention and attention-for-clustering for better and interpretable models. The quadratic complexity is relaxed to linear complexity. The proposed PaCa module is used in designing efficient and interpretable ViT backbones and semantic segmentation head networks. In experiments, the proposed methods are tested on ImageNet-1k image classification, MS-COCO object detection and instance segmentation and MIT-ADE20k semantic segmentation. Compared with the prior art, it obtains better performance in all the three benchmarks than the SWin and the PVTs by significant margins in ImageNet-1k and MIT-ADE20k. It is also significantly more efficient than PVT models in MS-COCO and MIT-ADE20k due to the linear complexity. The learned clusters are semantically meaningful. Code and model checkpoints are available at https://github.com/iVMCL/PaCaViT.
翻译:视觉变换器(ViTs)建立在将图像块视为“视觉令牌”的假设之上,并学习块到块的注意力机制。基于块嵌入的分词器与其对应文本分词器之间存在语义鸿沟。块到块注意力面临二次复杂度问题,也使得学习到的ViTs难以解释。为解决ViT中的这些问题,本文提出在ViT中学习分块到聚类注意力(PaCa)。所提PaCa-ViT中的查询(Queries)以图像块为起点,而键(Keys)和值(Values)则直接基于聚类(预先定义少量聚类簇)。聚类通过端到端学习实现,从而产生更好的分词器,并诱导联合聚类服务于注意力、注意力服务于聚类,进而构建更优且可解释的模型。二次复杂度被降低为线性复杂度。所提出的PaCa模块用于设计高效且可解释的ViT主干网络和语义分割头网络。在实验中,所提方法在ImageNet-1k图像分类、MS-COCO目标检测和实例分割以及MIT-ADE20k语义分割任务上进行了测试。与现有技术相比,该方法在所有三个基准上均取得更优性能,在ImageNet-1k和MIT-ADE20k上显著优于Swin和PVT模型。由于线性复杂度,其在MS-COCO和MIT-ADE20k上相比PVT模型也显著更高效。学习到的聚类簇具有语义意义。代码和模型检查点已在https://github.com/iVMCL/PaCaViT开源。