[论文解读] KVT: k-NN Attention for Boosting Vision Transformers

Pichao Wang, Xue Wang+5 more|arXiv (Cornell University)|May 28, 2021

Advanced Neural Network Applications参考文献 62被引用 4

一句话总结

本文提出 k-NN 注意力，一种轻量级、可微的机制，通过为每个查询仅选择最相似的前 k 个 key 替代视觉变换器中的密集自注意力，从而增强局部归纳偏置并过滤噪声 token。该方法在不修改网络结构的前提下，加速了 11 种视觉变换器架构的训练，并提升了性能。

ABSTRACT

Convolutional Neural Networks (CNNs) have dominated computer vision for years, due to its ability in capturing locality and translation invariance. Recently, many vision transformer architectures have been proposed and they show promising performance. A key component in vision transformers is the fully-connected self-attention which is more powerful than CNNs in modelling long range dependencies. However, since the current dense self-attention uses all image patches (tokens) to compute attention matrix, it may neglect locality of images patches and involve noisy tokens (e.g., clutter background and occlusion), leading to a slow training process and potential degradation of performance. To address these problems, we propose the $k$-NN attention for boosting vision transformers. Specifically, instead of involving all the tokens for attention matrix calculation, we only select the top-$k$ similar tokens from the keys for each query to compute the attention map. The proposed $k$-NN attention naturally inherits the local bias of CNNs without introducing convolutional operations, as nearby tokens tend to be more similar than others. In addition, the $k$-NN attention allows for the exploration of long range correlation and at the same time filters out irrelevant tokens by choosing the most similar tokens from the entire image. Despite its simplicity, we verify, both theoretically and empirically, that $k$-NN attention is powerful in speeding up training and distilling noise from input tokens. Extensive experiments are conducted by using 11 different vision transformer architectures to verify that the proposed $k$-NN attention can work with any existing transformer architectures to improve its prediction performance. The codes are available at \url{https://github.com/damo-cv/KVT}.

研究动机与目标

• 解决由于密集自注意力聚合所有 token 而导致的视觉变换器训练缓慢和对噪声敏感的问题。
• 克服全连接注意力在捕捉局部图像结构以及处理杂乱背景或遮挡时的局限性。
• 提出一种简单但有效的卷积归纳偏置替代方案，在保持全局长程建模能力的同时提升归纳偏置。
• 证明 k-NN 注意力可普遍应用于任意视觉变换器架构，以提升准确率和训练速度。
• 从理论和实证两方面验证 k-NN 注意力在不修改架构的前提下，可减少噪声并加快收敛速度。

提出的方法

• 用 k-NN 注意力替代标准的密集自注意力机制，即对每个查询 token，仅使用最相似的前 k 个 key token 计算注意力得分。
• 通过点积注意力计算查询与 key 向量之间的相似度，然后基于余弦相似度或点积相似度选择最接近的 k 个 key。
• 仅在选定的 k 个 key 上计算注意力图，从而降低计算成本并过滤无关或噪声 token。
• k-NN 机制是可微的且端到端可训练，可无缝集成到任意视觉变换器中而无需架构修改。
• 该方法天然继承局部归纳偏置，因为邻近 token 更可能被选中，即使未显式引入卷积操作。
• k-NN 机制根据特征相似度动态调整聚合图，既允许长程依赖，又能抑制远距离且不相似的 token。

实验结果

研究问题

• RQ1 用 k-NN 注意力替代密集自注意力是否能提升视觉变换器的训练速度与模型鲁棒性？
• RQ2 k-NN 注意力在训练过程中是否能有效过滤背景杂乱和遮挡等噪声 token？
• RQ3 与标准自注意力相比，k-NN 注意力在多大程度上保留或增强了局部归纳偏置？
• RQ4 k-NN 注意力是否可普遍应用于多种视觉变换器架构而无需架构修改？
• RQ5 在性能与训练动态方面，k-NN 注意力与卷积归纳偏置相比表现如何？

主要发现

• k-NN 注意力在 11 种视觉变换器架构中显著加速了训练，缩短了收敛时间，且未牺牲最终准确率。
• 该方法通过过滤背景杂乱和遮挡区域等噪声 token 改进了泛化能力，实现了更鲁棒的特征学习。
• k-NN 注意力在 ImageNet-1K 上达到或超越了当前最先进性能，且在多种主干网络架构上均表现出一致的性能提升。
• 理论分析表明，k-NN 注意力通过控制异常值和噪声特征的影响，维持了强泛化保证。
• 实证结果证实，k-NN 注意力降低了注意力图的方差并提升了注意力图的稀疏性，使注意力头更加聚焦且有意义。
• 该方法具有普适性：可无缝插入任意视觉变换器而无需架构修改，且在不同模型与数据集上均能持续提升性能。