[論文レビュー] Sparse MLP for Image Recognition: Is Self-Attention Really Necessary?

Chuanxin Tang, Yucheng Zhao+4 more|arXiv (Cornell University)|Sep 12, 2021

Advanced Neural Network Applications被引用数 13

ひとこと要約

この論文は、自己注意機構を排除した視覚モデルであるsMLPNetを提案する。sMLPNetは、軸方向のグローバルモデリングと重み共有を用いたスパースMLP（sMLP）モジュールによって自己注意機構を置き換える。パラメータ数とFLOPsを削減しながら過学習を回避することで、sMLPNetはImageNet-1Kで83.4%のトップ1精度を達成し、パラメータ数は65.9Mにとどまる。これは、Swin Transformerと同等の性能を25%小さいモデルサイズ、10%少ないFLOPsで達成した。この結果から、自己注意機構は最先端の画像認識に不可欠ではないことが示された。

ABSTRACT

Transformers have sprung up in the field of computer vision. In this work, we explore whether the core self-attention module in Transformer is the key to achieving excellent performance in image recognition. To this end, we build an attention-free network called sMLPNet based on the existing MLP-based vision models. Specifically, we replace the MLP module in the token-mixing step with a novel sparse MLP (sMLP) module. For 2D image tokens, sMLP applies 1D MLP along the axial directions and the parameters are shared among rows or columns. By sparse connection and weight sharing, sMLP module significantly reduces the number of model parameters and computational complexity, avoiding the common over-fitting problem that plagues the performance of MLP-like models. When only trained on the ImageNet-1K dataset, the proposed sMLPNet achieves 81.9% top-1 accuracy with only 24M parameters, which is much better than most CNNs and vision Transformers under the same model size constraint. When scaling up to 66M parameters, sMLPNet achieves 83.4% top-1 accuracy, which is on par with the state-of-the-art Swin Transformer. The success of sMLPNet suggests that the self-attention mechanism is not necessarily a silver bullet in computer vision. The code and models are publicly available at https://github.com/microsoft/SPACH

研究の動機と目的

• 視覚変換器における最先端性能を達成するために自己注意機構が不可欠であるという仮定に挑戦すること。
• 過学習と高いパラメータ数という主な制限要因を克服することで、MLPベースのモデルが自己注意機構なしで競争力のある性能を達成できるかどうかを調査すること。
• 軸方向計算と重み共有を活用して効率的なグローバル特徴モデリングを可能にする、新しいスパースMLP（sMLP）モジュールを設計すること。
• Swin Transformerなどの最先端の視覚変換器と同等またはそれ以上の性能を達成できる、自己注意機構を排除したネットワーク（sMLPNet）を構築すること。
• データ増強や追加事前学習を用いない標準的な学習プロトコル下で、ImageNet-1KにおけるsMLPNetの有効性を検証すること。

提案手法

• sMLPモジュールは、2次元画像トークンの行方向および列方向に1次元MLPを適用することで、軸方向のグローバル依存関係モデリングを実現する。
• すべての行（または列）にわたってパラメータを共有することで、モデルサイズと計算複雑性を著しく削減する。
• 各トークンがsMLP層内で自身の行および列の隣接トークンとのみ相互作用するようにスパarsityを導入することで、完全な2次関数的注意機構を回避する。
• sMLPモジュールは、ViTと同様のアーキテクチャに統合され、トークン混合ブロック内の自己注意機構が置き換えられる。
• 局所性バイアスを保持するために、トークン混合ブロックでは深度方向畳み込みが用いられ、チャネル混合は標準MLPと同一のままにされる。
• モデルスケーリングは幅と深さの調整により実現され、224×224解像度のImageNet-1Kで実験が実施された。

実験結果

リサーチクエスチョン

• RQ1 MLPベースの視覚モデルは、自己注意機構を使用せずに最先端の性能を達成できるか？
• RQ2 視覚変換器における自己注意機構は、代替的なグローバルモデリング機構に比べて根本的な優位性を提供するのか？
• RQ3 MLPにおけるスパース接続性と重み共有は、過学習を緩和し、計算コストを削減しながら性能を維持できるか？
• RQ4 局所性バイアスと軸方向グローバルモデリングを組み合わせることで、より小さなモデルサイズで競争力ある精度を達成できるか？
• RQ5 Swin Transformerの成功は、局所性に起因するのか、それとも自己注意機構に起因するのか？また、自己注意機構を排除したモデルはその性能を再現できるか？

主な発見

• sMLPNetは、24.1MパラメータでImageNet-1Kで81.9%のトップ1精度を達成し、同サイズ制約下でのほとんどのCNNおよび視覚変換器を上回った。
• 65.9MパラメータのsMLPNet-Bは、83.4%のトップ1精度を達成し、Swin-Bと同等の性能を示したが、パラメータ数は25%少なく、FLOPsは10%少なかった。
• 66Mパラメータでも過学習の兆候が見られず、MLP-Mixerやその他のMLPベースのモデルが深刻な過学習を示すのとは対照的であった。
• sMLPNet-T（19.2Mパラメータ）は、50億FLOPs未満のモデルの中で最高の81.3%のトップ1精度を達成した。
• DeiT-SおよびDeiT-Bと比較して、sMLPNet-Tははるかに少ないパラメータ数とFLOPsで高い精度を達成した。
• sMLPモジュールにより、シーケンス長に伴う計算複雑性が2次関数的から線形に低下し、高解像度入力の効率的処理が可能になった。