はじめに
最近、MLPベースの画像認識がVision業界を騒がせているのはご存知でしょうか？
2020年の終わりからViT[1]をきっかけに、Vision界にTransformerが本格的に流入してきました。しかし、Self-Attentionの計算コストが高い(2次)ということもあって、その問題を対処する研究も多く進められていました。そんな中、Google Brainのチームによる論文『MLP-Mixer: An all-MLP Architecture for Vision』[2]が発表されました。その内容は、ニューラルネットワークの基礎中の基礎であるMLPと転置などのテンソル形状の変更、正規化だけを使用したアーキテクチャで、ViTに対抗できるようなモデルが作れたいう衝撃的なものでした。
この記事では、MLP-Mixerがどのような構造についての解説を簡単にしたいと思います。続編以降で、MLP-Mixerの仲間たちのアイデアについても紹介していきたいと思います。間違いや質問などがございましたら遠慮なくコメントを頂ければ幸いです。
MLP-Mixer
アイデア
MLP-MixerはViT同様、16x16のようなパッチ分割をします(ViTについてご存じない方は例えばこちらの記事を御覧ください)。
分割したパッチに対して、MLPを使用した以下の操作を、繰り返しを行うアイデアです。
- 同じ空間の特徴間の関係
- 異なる空間位置の特徴の関係
CNNでは1x1Conv.が1を、3x3Conv.などが1,2の両方を担っており。TransformerではSelf-Attentionが1,2の両方を担っていました。MLP-Mixerの枠組みでは1にはchannel-mixing block、2にはtoken-mixing blockという名前がついております。2つのブロックに入力する前に、関係を捉える次元を切り替える必要があります。それは転置によって実現しています。
どちらのブロックもMLPによって構成されており、それ以外の構成要素としてはGELU, 残差接続, Layer Normalizationが使われています。
Mixer-B/16というモデルではこれらのモジュールが12回繰り返され、グローバルアベレージプーリングの後に、1000クラスの分類器に入力するだけ、という大変シンプルなモデルです。
実験結果の一部
実験では、ViTとImageNet, ImageNet-21k, JFT-300Mのデータセットを使用して事前学習の精度が比較されています。ViTよりも精度は少し落ちるものの、MLPだけでこの精度が叩き出せるとは驚くべきことですね。また、ViTに比べて計算コストが低いこともスループットやTPUv3 core-daysの結果からもわかります。
論文ではファインチューニングの実験もされていますが、ここでは割愛します。
筆者コメント
筆者自身もPyTorchで実装してみましたが、容易に実装することができました。CIFAR10で学習を試しましたが、高々10分程度の学習でも、90%近い精度を叩き出すことができて、正直驚いています。現状ではImageNetあたりの量や解像度ではさすがに学習はさすがに辛いです。リソースがないと辛いですし、実際筆者がImageNetで学習させてみても、label smoothingやら、CutMix, MixUpなどいろいろやって初めてまともに学習できます。筆者自身は生成モデルについての研究もしているので、生成モデルにもMLPベースの画像認識モデルの知見を取り込んで行けたらいいなと思っています。
おわりに
MLP-Mixerを始めとする、MLPベース画像認識モデルはCNNやTransformerに比べたら、まだ性能やあるいは学習速度、パラメータ数など課題が盛りだくさんですが、シンプルゆえに今後の発展が望まれます。CNNが過去の遺物になる時代もくるかもしれませんね。MLPベースのアイデアは他のチームも研究していたようで、この論文が公開された直後に、近いアイデアの論文が公開されております。続編ではその１つであるResMLP[4]について取り上げられたらと思います。
参考文献
[1] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, Alexey Dosovitskiy et al., 2020, https://arxiv.org/abs/2010.11929
[2] MLP-Mixer: An all-MLP Architecture for Vision, Ilya Tolstikhin et al., 2021, https://arxiv.org/abs/2101.07172
[3] 画像認識の大革命。AI界で話題爆発中の「Vision Transformer」を解説！, https://qiita.com/omiita/items/0049ade809c4817670d7
[4] ResMLP: Feedforward networks for image classification with data-efficient training, Hugo Touvron et al., 2021, https://arxiv.org/abs/2105.03404