KVT：用于提升视觉 Transformer 的 k-NN 注意力

卷积神经网络（CNNs）因其在捕获局部性与平移不变性方面的能力，多年来主导了计算机视觉。近来，许多视觉 transformer 架构被提出，并展现出有前景的性能。视觉 transformer 中的一个关键组件是全连接自注意力，其在建模长程依赖上比 CNNs 更强大。然而，由于当前的稠密自注意力使用所有图像块（tokens）来计算注意力矩阵，它可能忽略图像块的局部性并引入噪声 tokens（例如杂乱背景与遮挡），导致训练过程缓慢以及潜在的性能退化。为解决这些问题，我们提出 $k$-NN 注意力以提升视觉 transformer。具体而言，我们不为注意力矩阵计算涉及所有 tokens，而是对每个查询仅从键中选取最相似的 top-$k$ 个 tokens 来计算注意力图。所提出的 $k$-NN 注意力天然继承了 CNNs 的局部偏置而无需引入卷积操作，因为邻近 tokens 往往比其他 tokens 更相似。此外，$k$-NN 注意力允许探索长程关联，同时通过从整幅图像中选择最相似的 tokens 来过滤掉不相关的 tokens。尽管简单，我们理论上与经验上均验证了 $k$-NN 注意力在加速训练与从输入 tokens 中蒸馏噪声方面是有效的。通过使用 11 种不同的视觉 transformer 架构进行了大量实验，以验证所提出的 $k$-NN 注意力可与任何现有 transformer 架构配合工作以提升其预测性能。代码可在 \url{https://github.com/damo-cv/KVT} 获取。