【MLP-Mixer】MLPがCNN，Transformerを超える日
3つの要点
✔️ 多層パーセプトロン（MLP）だけのシンプルなアーキテクチャでSoTAに匹敵する画像分類性能を達成
✔️ 位置ごとの特徴と空間的な情報を繰り返しMixする
✔️ 計算コストを抑えつつ、同時に高い精度を達成
MLP-Mixer: An all-MLP Architecture for Vision
written by Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, Alexey Dosovitskiy
(Submitted on 4 May 2021 (v1), last revised 17 May 2021 (this version, v2))
Comments: Accepted by arXiv
Subjects: Computer Vision and Pattern Recognition (cs.CV); Artificial Intelligence (cs.AI); Machine Learning (cs.LG)
code：
はじめに
コンピュータビジョン領域では，長らくConvolution Neural Network (CNN)がスタンダードとして最も良く使われており，また最近ではVision Transformer (ViT)などのようなAttentionベースのネットワークも注目を集めています．しかし，本論文ではConvolutionとAttentionの両方とも必要がないことを示しています．このMLP-Mixerは2種類のレイヤーから構成されています．1つ目は画像パッチごとに適用される（位置ごとの特徴量をMixする）MLPで，2つ目は画像パッチを横断して適用される（空間情報をMixする）MLPです．このMLP-Mixerは十分な量のデータセットと正規化手法を用いることで，SoTAに匹敵する画像分類性能を達成しています．
アーキテクチャ
まずは，上図のMixerの全体像を説明します．画像を16×16などのパッチに分割し，入力とします．次にパッチごとにlinear embeddingsを行います．そして，それらに対して後述するMixer Layerを繰り返し適用します．最後に，一般的なCNNと同様にGlobal Average Poolingを適用して画像のクラス分類を行います．
次に，上図のMixer Layerの内容について説明します．線形射影された画像パッチ（token）を入力として，パッチ（空間的情報）×チャンネル（特徴量）の次元で扱います．Mixerはchannel-mixing MLPとtoken-mixing MLPの2種類のMLPを用います．Token-mixing MLPは異なる空間位置（token）の間で特徴量をMixする役割を果たし，独立にそれぞれのチャンネルに適用されてテーブルのそれぞれの列を入力として扱います．一方，channel-mixing MLPは，異なるチャンネル間で特徴量をMixする役割を果たし，独立にそれぞれのtokenに適用され，テーブルのそれぞれの行を入力として扱います．
極端なケースですが，channel-mixing MLPは1×1畳み込みを用いたCNN，token-mixing MLPは単一チャンネルのdepth-wise convolutionsとみなすこともできます．しかし，このMLP-MixerはそうしたCNNよりもはるかに単純なアーキテクチャとなっています．
さらにtoken-mixing MLPとchannel-mixing MLPを数式化してみます．なお，MLP-Blockの構造は以下で示されている通りで，全結合層＋GELU＋全結合層から成り，CNNやTransformerと比べると非常に単純であることが分かります．また，それぞれのMLPを適用する際には，ResNetのようにスキップ接続（skip-connection）でつなぎます．さらに，各MLPの前にはLayer Normalizationを行っています．
token-mixing MLP
・U: token-mixingのMLPブロックの出力
・X: 画像をパッチに分割して得られた特徴量
・W1: MLPブロックの1つ目の全結合層
・W2: MLPブロックの2つ目の全結合層
・σ: MLPブロックのGELU関数
channel-mixing MLP
・Y: channel-mixingのMLPブロックの出力
・U: token-mixingのMLPブロックの出力
・W3: MLPブロックの1つ目の全結合層
・W4: MLPブロックの2つ目の全結合層
・σ: MLPブロックのGELU関数
なお，ここで用いられているGELU活性化関数は，GPTやBERTなどの有名な自然言語処理モデルでも使われているものです．Dropoutのように，いくつかのアクティベーションに0をランダムに乗算することでモデルをロバストにしています．
Mixerのアーキテクチャは以上になります．行列の積や転置など，非常に基本的な処理だけで全ての演算が完結していることが分かります．
実験
実験では，従来の方法と同様に，大規模なデータセットで事前学習を行い，その後少ないデータセット（下流タスク）でfine-tuningを行います．実験の目的は以下の3つです．
1. 下流タスクにおける精度
2. 事前学習にかかる計算コスト
3. 推論時のスループット
ここで，あくまでの本論文の目的はSoTAを達成することではなく，MLPベースのモデルが今日のCNNやTransformerに匹敵し，超えるだけの可能性があることを示すことにあることに注意して下さい．
下流タスク
・ImageNet
・CIFAR-10/100
・Oxford-IIIT Pets
・Oxford Flowers-102
・Visual Task Adaptation Benchmark
事前学習データ
・ImageNet
・ImageNet-21k
・JFT-300M
Mixerモデル詳細
結果
ImageNet-21kを用いて事前学習を行った場合，Visual TransformerやBiTと比較して，同程度の精度を達成しながらも，同時に計算コストを低く抑えることができています．
さらに，上図に示されている通り，MLP-Mixerはデータセットのサイズを大きくすればするほど画像分類の精度も向上し，その程度は他のモデルよりも大きいことが分かっています．300Mまでデータセットを大きくすると，BiTよりも精度が高くなり，しかもより高速に行うことができました．
さらに，上図は精度と計算コストのトレードオフに関する性能を示しており，Mixerは精度単体ではSoTAに若干及ばないものの，精度と計算コストのトレードオフを考慮した場合，他のSoTAモデルに並び，Pareto Frontier上にあることが分かります．
まとめ
Mixerは2種類のMLP層からなる非常に単純なアーキテクチャですが，計算コストを抑えつつも，SoTAモデルに匹敵する画像分類性能を達成しています．本論文を元にして研究が発展すれば，分類精度，計算コストの両方においてCNNやTransformerをも超える日も近いかもしれません．
この記事に関するカテゴリー
