[論文レビュー] X-volution: On the unification of convolution and self-attention

Xuanhong Chen, Hang Wang+1 more|arXiv (Cornell University)|Jun 4, 2021

Computability, Logic, AI Algorithms参考文献 34被引用数 23

ひとこと要約

本稿では、シフトされた特徴マップの学習可能な畳み込み変換によってグローバル自己注意を近似することで、畳み込みと自己注意を統合する新しいニューラルネットワークモジュールであるX-volutionを提案する。この手法により、マルチブランチアーキテクチャを通じて局所的および非局所的特徴学習を統合可能となり、推論時には単一の標準畳み込みに再パラメータ化可能である。ImageNetでは+1.2%のトップ-1精度向上、COCOオブジェクト検出では+1.7 APの向上を達成し、最先端の性能を実現した。

ABSTRACT

Convolution and self-attention are acting as two fundamental building blocks in deep neural networks, where the former extracts local image features in a linear way while the latter non-locally encodes high-order contextual relationships. Though essentially complementary to each other, i.e., first-/high-order, stat-of-the-art architectures, i.e., CNNs or transformers lack a principled way to simultaneously apply both operations in a single computational module, due to their heterogeneous computing pattern and excessive burden of global dot-product for visual tasks. In this work, we theoretically derive a global self-attention approximation scheme, which approximates a self-attention via the convolution operation on transformed features. Based on the approximated scheme, we establish a multi-branch elementary module composed of both convolution and self-attention operation, capable of unifying both local and non-local feature interaction. Importantly, once trained, this multi-branch module could be conditionally converted into a single standard convolution operation via structural re-parameterization, rendering a pure convolution styled operator named X-volution, ready to be plugged into any modern networks as an atomic operation. Extensive experiments demonstrate that the proposed X-volution, achieves highly competitive visual understanding improvements (+1.2% top-1 accuracy on ImageNet classification, +1.7 box AP and +1.5 mask AP on COCO detection and segmentation).

研究の動機と目的

• 畳み込みと自己注意の計算パターンが異なっているため、それらを単一のニューラルネットワークモジュールに原理的かつ統合的に統合する方法が不足しているという問題に対処すること。
• ビジョンタスクにおける効率的な計算を可能にする、低複雑性で微分可能なグローバル自己注意の近似を構築すること。
• 最適化の安定性を畳み込みが保つ一方で、自己注意によって長距離依存関係を捉えることができるマルチブランチモジュールを設計すること。
• 訓練済みのマルチブランチモジュールを、推論時に単一の効率的な畳み込み層に変換可能な構造的再パラメータ化を可能にすること。

提案手法

• グローバル自己注意のO(n)近似として、ピクセルシフト自己注意（PSSA）を提案。ドット積み注意を8方向にシフトされた特徴マップ上の要素ごとの演算に置き換える。
• 標準の3×3および拡張係数5×5の畳み込みブランチに加え、非局所的特徴相互作用に特化したPSSAブランチを備えたマルチブランチX-volutionモジュールを構築する。
• 元の特徴マップとシフトされた特徴マップの間で学習可能なシフト操作と要素ごとの乗算を適用し、動的注意マップを生成する。
• 訓練の安定化と一般化性能の向上を図るため、最終出力の前段階にバッチ正規化を適用する。
• 訓練後、マルチブランチアーキテクチャを単一の標準畳み込みに統合する構造的再パラメータ化を活用し、推論の効率性を実現する。

実験結果

リサーチクエスチョン

• RQ1 変換された特徴に対して畳み込み操作のみを用いて、グローバル自己注意を効果的かつ効率的に近似できるか？
• RQ2 畳み込みと自己注意を統合したユニフィードモジュールは、安定して訓練可能であり、独立したCNNやTransformerよりも優れた性能を達成できるか？
• RQ3 構造的再パラメータ化を伴うマルチブランチ設計により、畳み込みの最適化利点を維持しながら、自己注意の表現力も得られるか？
• RQ4 得られたX-volutionモジュールは、アーキテクチャの変更なしに、既存のネットワークにプラグインとして容易に統合可能か？

主な発見

• X-volutionは、ResNet-50ベースラインと比較してImageNet-1K分類タスクで+1.2%のトップ-1精度向上を達成した。
• COCOオブジェクト検出タスクでは、ボックスAPが+1.7ポイント、マスクAPが+1.5ポイント向上した。
• 長時間の訓練スケジュール（例：3xおよび6x）において、X-volutionの性能向上が顕著に増大し、3x設定ではCOCOで+2.0 APの向上を記録した。
• PSSA近似は、短時間の訓練スケジュールでは標準畳み込みと同等の性能を示したが、訓練が進むにつれてそれを上回り、より優れた長距離モデリング能力を示した。
• 再パラメータ化されたX-volutionモジュールは、標準畳み込みとしてデプロイ可能であり、高い性能を維持しながら効率的な推論を実現した。