[論文レビュー] S$^2$-MLP: Spatial-Shift MLP Architecture for Vision

Tan Yu, Li Xu+3 more|arXiv (Cornell University)|Jun 14, 2021

Advanced Neural Network Applications参考文献 39被引用数 29

ひとこと要約

S2-MLPは、クロスパッチ通信のためのパラメータフリーの空間シフト演算を使用する純粋なMLPアーキテクチャで、ViTとMLP-MixerよりもFLOPsとパラメータが少なく競争力のあるImageNet-1Kパフォーマンスを達成する。

ABSTRACT

Recently, visual Transformer (ViT) and its following works abandon the convolution and exploit the self-attention operation, attaining a comparable or even higher accuracy than CNNs. More recently, MLP-Mixer abandons both the convolution and the self-attention operation, proposing an architecture containing only MLP layers. To achieve cross-patch communications, it devises an additional token-mixing MLP besides the channel-mixing MLP. It achieves promising results when training on an extremely large-scale dataset. But it cannot achieve as outstanding performance as its CNN and ViT counterparts when training on medium-scale datasets such as ImageNet1K and ImageNet21K. The performance drop of MLP-Mixer motivates us to rethink the token-mixing MLP. We discover that the token-mixing MLP is a variant of the depthwise convolution with a global reception field and spatial-specific configuration. But the global reception field and the spatial-specific property make token-mixing MLP prone to over-fitting. In this paper, we propose a novel pure MLP architecture, spatial-shift MLP (S$^2$-MLP). Different from MLP-Mixer, our S$^2$-MLP only contains channel-mixing MLP. We utilize a spatial-shift operation for communications between patches. It has a local reception field and is spatial-agnostic. It is parameter-free and efficient for computation. The proposed S$^2$-MLP attains higher recognition accuracy than MLP-Mixer when training on ImageNet-1K dataset. Meanwhile, S$^2$-MLP accomplishes as excellent performance as ViT on ImageNet-1K dataset with considerably simpler architecture and fewer FLOPs and parameters.

研究の動機と目的

• 中規模データ上の vision backbones において畳み込みと自己注意の代替案の必要性を動機づける。
• パラメータフリーの空間シフトブロックを用いたパッチ通信を持つ純粋なMLPアーキテクチャ(S2-MLP)を提案する。
• 空間シフトベースの通信が、ViTやMLP-Mixerより少ないパラメータとFLOPsで競争力のある精度をもたらすことを示す。
• ImageNet-1K上でS2-MLPを評価し、深さ、幅、シフト方向、入力スケールを理解するためのアブレーションを行う。

提案手法

• パッチ単位の全結合埋め込みを導入する。
• 各ブロックには4つの全結合層に加え2つのGELU活性化および2つのLayerNormを含む、N個のスタックされたS2-MLPブロックを使用する。
• トークンミキシングを、チャネルをグルーピングし各グループを一方向にシフトさせて局所的なパッチ間通信を可能にする空間シフトモジュールに置換する。
• 空間シフトを、隣接パッチ間での固定深さ方向のシフトに等しいパラメータフリー演算として定義する。
• PFL、S2-MLPブロック、最終分類器層の複雑さの分析を提供する。

実験結果

リサーチクエスチョン

• RQ1 純粋なMLPアーキテクチャがトークンミキシングなしで中規模データ上のImageNet-1K精度で競争力を持てるか。
• RQ2 パラメータフリーの空間シフト機構がクロスパッチ通信を十分に提供し、MLP-Mixerと同等またはViTの性能に近づくか。
• RQ3 深さ(N)、隠れ層サイズ(c)、拡張比(r)、シフト方向、および入力スケールが精度と効率にどのように影響するか。
• RQ4 S2-MLPアーキテクチャは入力スケールに不変か、トークンミキシングMLPモデルとどう比較されるか。

主な発見

Model	Resolution	Top-1	Top-5	Params (M)	FLOPs (B)
S2-MLP-wide	224×224	80.0	94.8	71	14.0
S2-MLP-deep	224×224	80.7	95.4	51	10.5
Mixer-B/16	224×224	76.4	-	59	11.6
FF	224×224	74.9	-	59	11.6
ResMLP-36	224×224	79.7	-	45	8.9
ViT-B/16	384×384	77.9	-	55.5	-

• S2-MLP-wideはImageNet-1KでTop-1 80.0%とTop-5 94.8%を達成し、71Mパラメータと14B FLOPsで、同程度の規模のMLP-Mixer (Top-1 76.4%)を上回る。
• S2-MLP-deepはImageNet-1KでTop-1 80.7%とTop-5 95.4%を達成し、51Mパラメータと10.5B FLOPsで、同様の条件下でResMLP-36を超える。
• S2-MLPの性能は、報告された設定でViTと比較して競争力があり、より単純なアーキテクチャと低いFLOPs/パラメータを維持している。
• アブレーション結果は、深さを1から12ブロックに増やすと精度が向上（ImageNet100でTop-1が56.7%から87.1%）、ただし12–16ブロックを超えると小規模データセットで過学習のため性能が飽和またはわずかに低下する。
• 隠れサイズcを拡張するとある点まで精度が向上する（c=768でTop-1が87.1%）、ただしcを大きくするとパラメータとFLOPsが増加する。
• 4方向へのシフト（デフォルト）は強いクロスパッチ通信を提供する。方向を増やすとある点まで性能は向上するが、シフトをなくすと精度は大幅に低下する。