MLP-Mixer: An all-MLP Architecture for Vision

nogawanogawa commented 2 years ago

論文URL

著者

Ilya Tolstikhin, Neil Houlsby, Alexander Kolesnikov, Lucas Beyer, Xiaohua Zhai, Thomas Unterthiner, Jessica Yung, Andreas Steiner, Daniel Keysers, Jakob Uszkoreit, Mario Lucic, Alexey Dosovitskiy

会議

NeurIPS 2021

背景

CVの領域では、長らく畳み込みを用いたCNNが主要な手法として進歩してきた。また、最近ではVision TransformerのようなTransfomerを用いたモデルも登場してきている。

目的

畳み込みもTransformerも使用しない新しいネットワーク構造によって高い精度を実現できることを示す

アプローチ

MLP-Mixer
- MLP (Multi-Layer Perceptron)だけを使用した新しい構造のモデル
- 従来の畳み込みやTransformerの構造に比べてシンプルな構造
- 学習・推論に必要なリソースの割に、従来手法と同程度の精度を達成した

nogawanogawa commented 2 years ago

背景

CVの領域では、長らく畳み込みを用いたCNNが主要な手法として進歩してきた。また、最近ではVision TransformerのようなTransfomerを用いたモデルも登場してきている。

nogawanogawa commented 2 years ago

目的

畳み込みもTransformerも使用しない新しいネットワーク構造によって高い精度を実現できることを示す

アプローチ

MLP-Mixer
- MLP (Multi-Layer Perceptron)だけを使用した新しい構造のモデル
- 従来の畳み込みやTransformerの構造に比べてシンプルな構造
- 学習・推論に必要なリソースの割に、従来手法と同程度の精度を達成した

nogawanogawa commented 2 years ago

MLP-Mixer

モダンなDeep vision archtectureの構成

モダンなCNNやTransformerによるDeep vision archtectureを抽象化すると、

per-location (channel-mixing) operations : 与えられた空間位置における特徴のmix
cross-location (token-mixing) operations : 異なる空間位置間での特徴のmix

の2つのレイヤに分けられ、構造によってはこれらが同時に実行される。

CNNでは2.は畳み込み（N>1）とpoolingによって実行され、同時に1×1の畳み込みが1. の役割を担っている。 Vision Transformerはself-attentionが1. と2.の役割を同時に担い、MLPブロックが1.の役割で実行されていると考えることができる。

MLP-Mixerの構成

MLP-Mixerにおいても、上記の1. , 2. の役割が含まれ、それぞれがMLPのみで実現される。 MLP-Mixerのアーキテクチャを下記に示す。

動作としては、まず元画像を小さなパッチに分割され、それらがMLP-Mixerの入力となる。 MLP-Mixerは同一サイズの複数の層で構成され、各層は2つのMLPブロックで構成されます。

上図の左側のMLPブロックはtoken-mixingを行うMLPで、チャンネルとパッチで構成される行列が転置された状態で入力されている。その後再び転置され、2つ目のchannel-mixingを行うMLPに入力されるような動作になっている。それぞれのMLP層は2層のFC層とGELUで構成されている。

nogawanogawa commented 2 years ago

評価

中〜大規模なデータセットで事前学習してから、小〜中規模のデータセットを用いた分類タスクを用いて評価を行う。

観点は、

小〜中規模のデータセットを用いた下流タスクの精度
中〜大規模なデータセットで事前学習の計算コスト
実用化するにあたっての推論時のスループット

としている。

結果

Avg5は下流タスク（ImageNet、CIFAR-10、CIFAR-100、Pets、Flowers）の平均を表す。 MLP-Mixerは（正則化を行えば）、他のモデルに若干劣るものの、全体的に高い性能を示していることがわかる。

正則化がない場合にはoverfitし、これはViTでも同様の事象が発生することがわかっている。

ImageNet-21KまたはJFT、web image pairを用いた事前学習の結果を下記（左図）に示す。（左上に行けば行くほど事前学習の計算コストが小さく高い性能が出ることを表す。）

HaloNet、ViT、NFNet、MPLと同程度の高い学習効率を達成している。

また、右図は学習データのサイズに対する精度を示しており、MLP-Mixerは、ResNetやViTより同じ学習データ量に対して、より早く精度を向上させることがわかる。また、大規模なMLP-MixerとViTのモデル間の精度の差は小さくなっている。

計算リソース、スループットに対する精度に関して、同規模のViTモデルやAdamで事前学習したBiTモデルと比較した結果を下記に示す。

各モデルとも学習損失が非常に似た値を達成していることがわかる。図3は、Mixerがモデル規模の下限では他のモデルよりやや劣るものの、上限では優れていることを意味している。（計算リソースを十分に投下できれば従来手法より高い性能を達成できるかもしれないことを意味する）

また、上の表で、Mixer-H/14は、ImageNetにおいてViT-H/14とわずか0.3%の差であり、2.2倍高速に実行されており、より効果的に学習ができていることを表す。

nogawanogawa / paper_memo