[論文レビュー] On the Integration of Self-Attention and Convolution

Xuran Pan, Chunjiang Ge+5 more|arXiv (Cornell University)|Nov 29, 2021

Domain Adaptation and Few-Shot Learning参考文献 35被引用数 9

ひとこと要約

本論文では、同一の1×1畳み込み演算をコア計算として共有することにより、自己注意機構と畳み込みを滑らかに統合する新しいニューラルネットワークモジュールACmixを提案する。このアプローチにより、効率的な連合特徴学習が可能となり、純粋な自己注意または畳み込みモデルと比較して計算コストの増加が最小限に抑えられる。本手法は、画像認識および物体検出ベンチマークにおいて最先端の性能を達成する。

ABSTRACT

Convolution and self-attention are two powerful techniques for representation learning, and they are usually considered as two peer approaches that are distinct from each other. In this paper, we show that there exists a strong underlying relation between them, in the sense that the bulk of computations of these two paradigms are in fact done with the same operation. Specifically, we first show that a traditional convolution with kernel size k x k can be decomposed into k^2 individual 1x1 convolutions, followed by shift and summation operations. Then, we interpret the projections of queries, keys, and values in self-attention module as multiple 1x1 convolutions, followed by the computation of attention weights and aggregation of the values. Therefore, the first stage of both two modules comprises the similar operation. More importantly, the first stage contributes a dominant computation complexity (square of the channel size) comparing to the second stage. This observation naturally leads to an elegant integration of these two seemingly distinct paradigms, i.e., a mixed model that enjoys the benefit of both self-Attention and Convolution (ACmix), while having minimum computational overhead compared to the pure convolution or self-attention counterpart. Extensive experiments show that our model achieves consistently improved results over competitive baselines on image recognition and downstream tasks. Code and pre-trained models will be released at https://github.com/LeapLabTHU/ACmix and https://gitee.com/mindspore/models.

研究の動機と目的

• 先行研究で見過ごされてきた、自己注意と畳み込みの間の根本的な計算的類似性を解明すること。
• 自己注意と畳み込みの長所を活かすハイブリッドニューラルモジュールを設計すること。この際、計算コストを倍増させないことを目的とする。
• 自己注意モジュールと畳み込みモジュールを組み合わせる際の通常の推論および学習コストを低減すること。
• 提案されたACmixモジュールが、複数のビジョンベンチマークにおいて、純粋な畳み込みモデルや自己注意モデルを上回ることを実証的に検証すること。

提案手法

• 標準的なk×k畳み込みを、k²個の個別的な1×1畳み込みに分解し、その後にシフトおよび和集合の操作を実行する。
• 自己注意のクエリ、キー、バリューの投影を、複数の1×1畳み込みとして再解釈し、その後に動的注意重みの計算とバリューの集約を実行する。
• ACmixにおいて、自己注意パスと畳み込みパスの両方に同じ1×1畳み込み投影を共有することで、重複する計算を排除する。
• 出力の結合に学習可能なスケーリングパラメータαとβを用い、特徴の動的ルーティングを可能にする。
• シフト操作をグループ畳み込みまたは学習可能なカーネルを用いて実装し、集約段階における効率性と柔軟性を維持する。
• 共有された1×1畳み込みと学習可能なα、βパラメータを同時に最適化するエンドツーエンドの最適化によりモデルを学習する。

実験結果

リサーチクエスチョン

• RQ1 基本的な演算に分解した場合、自己注意と畳み込みの間に根本的な計算的同値性が存在するか？
• RQ2 共有された1×1畳み込み演算が、自己注意モジュールと畳み込みモジュールの両方の基盤として機能できるか？
• RQ3 同じ特徴投影を再利用するハイブリッドモジュールは、別々の自己注意モジュールと畳み込みモジュールを用いる場合と比較して、より優れた性能と効率を達成できるか？
• RQ4 ネットワークの深さに応じて、自己注意と畳み込みの相対的寄与度はどのように変化するか？

主な発見

• ImageNet-1KにおいてACmixはトップ-1正解率81.9%を達成し、同じ設定下でSwin-T（81.5%）およびConv-Swin-T（81.4%）を上回った。
• アブレーションスタディの結果、パスの結合に学習可能なパラメータαとβを用いることで、固定比率や単一パスモデルよりも優れた性能が得られた。
• グループ畳み込みに基づくシフトモジュールを搭載したモデルは、単一のRTX 2080Ti GPUで419 FPSの推論速度を達成し、高い推論効率を示した。
• 学習されたαおよびβパラメータの可視化から、安定した特徴の選択傾向が観察された。初期層では畳み込みが優位に機能し、深層に進むにつれて自己注意の寄与度が高まる傾向が見られた。
• 両パスの組み合わせは、全評価ベンチマークおよび設定で、単一モジュールを用いたモデルを常に上回る性能を示した。
• 別々に自己注意と畳み込みモジュールを学習する場合と比較して、本手法は計算コストの増加を抑えつつ、性能を維持または向上させた。