[論文レビュー] 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個のキーベクトルのみを選択することで、ビジョントランスフォーマーにおける密な自己アテンションを置き換える軽量で微分可能なメカニズムである。これにより、局所的なインダクティブバイアスが強化され、ノイズの多いトークンがフィルタリングされる。この手法は、アーキテクチャの変更なしに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}.

研究の動機と目的

• すべてのトークンを集約する密な自己アテンションによって引き起こされるビジョントランスフォーマーの訓練遅延およびノイズ感受性を解消する。
• 完全結合アテンションの限界、すなわち局所的な画像構造を捉える能力の不足と、ごみ混じりの背景や遮蔽の処理における困難を克服する。
• 畳み込み型インダクティブバイアスの代替として単純だが効果的な手法を提案し、グローバルな長距離依存関係のモデリングを維持しながらインダクティブバイアスを向上させる。
• k-NNアテンションが任意のビジョントランスフォーマーのアーキテクチャに普遍的に適用可能であり、精度と訓練速度の両方を向上させることを示す。
• 理論的および実験的に、k-NNアテンションがアーキテクチャの変更なしにノイズを低減し、収束を高速化することを検証する。

提案手法

• 各クエリトークンに対して、類似度が最も高い上位k個のキーベクトルのみを用いてアテンションスコアを計算する、標準的な密な自己アテンションの代わりにk-NNアテンションを採用する。
• クエリとキーベクトル間の類似度をドット積アテンションで計算し、コサイン類似度またはドット積類似度に基づいてk個の近隣キーベクトルを選択する。
• アテンションマップは、選択されたk個のキーベクトルのみを対象として計算されるため、計算コストが低減され、関係のないまたはノイズの多いトークンがフィルタリングされる。
• k-NNメカニズムは微分可能であり、エンドツーエンドで学習可能であり、アーキテクチャの変更なしに任意のビジョントランスフォーマーに統合可能である。
• 明示的な畳み込み演算がなくても、近接するトークンがより選択されやすいことから、自然に局所的なインダクティブバイアスを獲得する。
• k-NNメカニズムは特徴の類似度に基づいて動的にアグリゲーショングラフを再構成するため、長距離依存関係を捉えつつ、遠く離れていて類似度の低いトークンを抑制できる。

実験結果

リサーチクエスチョン

• RQ1 密な自己アテンションをk-NNアテンションに置き換えることで、ビジョントランスフォーマーの訓練速度とモデルの耐障害性が向上するか？
• RQ2 k-NNアテンションは、背景のごみや遮蔽領域のようなノイズの多いトークンを訓練中に効果的にフィルタリングするか？
• RQ3 標準的な自己アテンションと比較して、k-NNアテンションはどれほど局所的なインダクティブバイアスを保持または強化するか？
• RQ4 k-NNアテンションは、アーキテクチャの変更なしに多様なビジョントランスフォーマーのアーキテクチャに普遍的に適用可能か？
• RQ5 性能と訓練ダイナミクスの観点から、k-NNアテンションは畳み込み型インダクティブバイアスと比較してどのように差をつけるか？

主な発見

• k-NNアテンションは11種類のビジョントランスフォーマーのアーキテクチャすべてで訓練を顕著に高速化し、最終的な精度を損なわずに収束時間を短縮した。
• 背景のごみや遮蔽領域のようなノイズの多いトークンをフィルタリングすることで、一般化性能が向上し、より耐障害性の高い特徴学習が実現した。
• ImageNet-1Kでは、SOTAまたは競争力のある性能を達成し、複数のバックボーンアーキテクチャで一貫した改善が得られた。
• 理論的分析により、k-NNアテンションは外れ値やノイジーサイズの影響を制御することで、強い一般化保証を維持することが示された。
• 実験的結果により、k-NNアテンションはアテンションマップの分散を低減し、スパarsity（スパarsity）を向上させ、より集中的で意味のあるアテンションヘッドを実現した。
• 本手法は普遍的に適用可能である：アーキテクチャの変更なしに任意のビジョントランスフォーマーに組み込み可能であり、異なるモデルやデータセットにおいて一貫して性能向上を示した。