はじめに

畳み込みやAttention機構を用いず、多層パーセプトロンのみを用いて最先端のモデルに見劣りしない性能をたたき出した画像認識モデルMLP-Mixerというモデルが提案されました。

今回はこのMLP-Mixerについて解説したいと思います。

論文データ

arxiv :「MLP-Mixer: An all-MLP Architecture for Vision」

[Submitted on 23 Mar 2021 (v1), last revised 30 Mar 2021 (this version, v2)]

論文の概要

画像パッチに分割しそれらのパッチを、空間方向及びチャンネル方向に関して多層パーセプトロンで変換(mix)させるモデルMLP-Mixerを提案した
このMLP-Mixerは畳み込みやAttentionを使用していないがSOTAと比較しても見劣りしない精度を誇り、さらにスループットに関しては特に優秀な性能をたたき出した

MLP-Mixerの解説

画像 :"MLP-Mixer: An all-MLP Architecture for Vision"

モデルの全体像は上の画像の通りです。

そして、MLP-Mixerは以下の3つのステップで画像認識を行います。

画像をP×Pのパッチに分割し、それぞれベクトルにEmbeddingする。
1.で取得したベクトル達を複数回Mixer Layer(後述)で処理する。
MixerLayerで処理されたベクトルをクラス分類処理する。

このMLP-Mixerの特徴は1.と2.のステップで、1.の画像のパッチ分割はVisionTransformerでも行われていた手法です。そして、MLP-Mixer最大の注目点はMLP-Mixerでの処理になります。そのため次の項でこのMixer Layerについて説明したいと思います。

Mixer-Layerの説明

図１:Mixer Layerの全体像

図2:MLP-BLockの構造

Mixer Layerは図１で表されるように以下の２つの部分

token-mixing
channel-mixing

によって構成されます。

token-mixing

図3：token-mixingの概要図

token-mixingの処理は以下の数式で表現されます。
$$
U_{∗,i} = X_{∗,i} + W_2 σ(W_1 \mathrm{LayerNorm}(X)_{∗,i}) \qquad \mathrm{for}: i = 1\cdots C
$$
ここで、上式に出てくる登場人物の紹介をします。

$X_{n,m}$：画像を分割して得られたパッチテンソルの集合をn×ｍの2次元テンソルに射影したもの。
ここで
- n：パッチの空間方向の情報を表す軸
- m：パッチのチャンネル方向の情報を表す軸
$U_{n,m}$：token-mixingのMLPブロックからの出力
$W_{1}$：MLPブロック内の１つ目の全結合層（図2の赤枠部分）
$\sigma$：MLPブロック内のGELU関数(図2の緑枠部分)
$W_{2}$：MLPブロック内の2つ目の全結合層（図2の青枠部分）

この時token-mixingの式の意味するところは、

画像パッチのEmbeddingテンソル($X$)の列方向ベクトル(パッチの空間方向)に対して

全結合ネット層$W_1$による変換を実施する。
活性化関数$\sigma$(GELU)を適用する。
全結合ネット層$W_2$による変換を実施する。

という処理を適用するということを表現しています。

この一連の処理によって、各パッチの空間成分の情報が混合(mix)されます。

そして、この空間成分の情報の混合(mix)がtoken-mixingブロックでの目的になります。

なお、図3の中で度々表れる転置処理は

1回目の転置処理：$X$の列方向に対して処理MLP-Blockの処理を実施するために実施
2回目の転置処理：１回に行った転置を元に戻すための処理

のために導入されたものになります。

channel-mixing

図4:channel-mixingの概要図

channel-mixingの処理は以下の数式で表現されます。
$$
Y_{j,∗} = U_{j,∗} + W_4 σ(W_3 \mathrm{LayerNorm}(U)_{j,∗}) \qquad \mathrm{for}: j = 1\cdots S
$$
ここで、上式に出てくる登場人物の紹介をします。

$U_{n,m}$：token-mixingのMLPブロックからの出力
$Y_{n,m}$：channel-mixingのMLPブロックからの出力
$W_{3}$：MLPブロック内の１つ目の全結合層（図2の赤枠部分）
$\sigma$：MLPブロック内のGELU関数(図2の緑枠部分)
$W_{4}$：MLPブロック内の2つ目の全結合層（図2の青枠部分）

この時channel-mixingの式の意味するところは、

token-mixingからの出力($U$)の行方向ベクトル(パッチのチャンネル方向)に対して

全結合ネット層$W_3$による変換を実施する。
活性化関数$\sigma$(GELU)を適用する。
全結合ネット層$W_4$による変換を実施する。

という処理を適用するということを表現しています。

そして、このチャンネル成分の情報の混合(mix)がchannel-mixingブロックでの目的になります。

まとめると

Mixer-Layerでは画像分割して得られた各パッチに関して、空間方向及びチャンネル方向の情報をそれぞれ混合(mix)させることで画像の特徴を学習させるモデルであると言えると思います。

実験結果

今回の記事では、論文中のMLP-Mixerのサイズ毎の複数のモデルとSOTAとの比較結果について解説したいと思います。

モデル

MLP-Mixerは層数に応じてS、B、L、Hの４種類、さらにそれらの中でさらに隠れ層のサイズS、CやMLPブロックのサイズ$D_C$、$D_S$が異なる計７種類のモデルを用いて実験を行ったと述べられています。

参考:VisionTransformerとのパラメーター数比較

なお、同程度の規模（同程度の層数）のVisionTransformerのモデルとパラメーター数を比較するとMixerの方が$\frac{2}{3}$程度となっています。

SOTAとの比較

論文ではImageNet-21k(中規模データ)や JFT-300M(大規模)データで事前学習したモデルをILSVRC2012 “ImageNet” などで転移学習させた結果に関してSOTAモデルと比較しています。

ここで、表の列ReaLとAvg５は

ReaL：ImageNetをより正確にアノテーションされた評価用ラベルでラベル付けしたもの
Avg５：ImageNet, CIFAR-10, CIFAR-100, Pets, Flowersに対して転移学習した時の平均精度

を表します。

比較結果の表より以下のことが分かると言えます

ImageNet-21kで事前学習した場合もJFT-300Mで事前学習した場合もSOTAと比較して精度に関しては見劣りしない(1～2％劣る程度)性能をたたき出した。
ImageNet-21k、JFT-300Mどちらの場合も他のSOTAと比較して上位のスループットを持つという特徴がMLP-Mixerにはある
大規模データ（JFT-300M）の場合、学習に要する時間(TPUv3 core-days)が最も少なくなる

まとめ

このモデルの結果を通して著者らは、従来の畳み込みやAttentionでは考慮できていなかった新しい帰納バイアスが存在する可能性を指摘しています。この帰納バイアスの考察が進むことで画像認識の更なるブレークスルーが生まれる可能性を私個人としては大いに期待しています。

ともあれ、去年のVisionTrandformerで面白くなってきた画像認識分野がこれから益々目が離せなくなりそうです。