Transformerよりもシンプル?「MLP-Mixer」爆誕(4日目) ~Experiments 2編~

19 min read読了の目安(約17200字

ニツオです。TwitterでAIやMLについて関連する話題を紹介してます。お気軽にフォローください。

さて、2021年5月にMLP-Mixerというモデルが爆誕しました。本日はその解説シリーズ4日目です。

  • 1日目: Abstract / Introduction
  • 2日目: Mixer Architecture
  • 3日目: Experiments 1
  • 4日目: Experiments 2
  • 5日目: Related Work
  • 6日目: Conclusion
  • 7日目: Appendix
  • 8日目: Source Code

「MLP-Mixer: An all-MLP Architecture for Vision」の原文はこちらです。2021年5月4日にGoogle ResearchとGoogle Brainの混合チームから発表され、関係者のTwitterでもかなり話題になっています。

シリーズ関連記事は一番下にリンク貼ってます。
早速みていきましょう。

3.1 Main results

Table 2 presents comparison of the largest Mixer models to state-of-the-art models from the literature.

表2は、Mixerモデル(LargeとHuge)と、参考文献に掲載されている最先端モデルとの性能比較です。

“ImNet” and “ReaL” columns refer to the original ImageNet validation labels [13] and cleaned-up ReaL labels [5].

"ImNet" と "ReaL" の列は、オリジナルのImageNet検証ラベル[13]とクリーンアップされたReaLラベル[5]を参照しています。

ReaLとは、Reassessed Labelの略で、オリジナルのImageNetの人間がつけたラベルをつけなおしてみた学習セットのこと。

参考論文はこちら。本論文の著者でもあるLucas BeyerらGoogle BrainとグループのDeep Mindが2020年6月に公開した論文でReaLは定義されてます。
Are we done with ImageNet?

“Avg.5” stands for the average performance across all five downstream tasks (ImageNet, CIFAR-10, CIFAR-100, Pets, Flowers).

3列目の "Avg.5"は、5つの下流Fine-Tuning学習タスク(ImageNet、CIFAR-10、CIFAR-100、Pets、Flowers)すべてにおける、平均性能を表しています。

Figure 2 (left) visualizes the accuracy-compute frontier.

図2(左)は、縦軸‐横軸が、精度-計算量でみた、フロンティア(最前線のモデルたちを示す)を可視化したものです。

When pre-trained on ImageNet-21k with additional regularization, Mixer achieves an overall strong performance (84.15% top-1 on ImageNet), although slightly inferior to other models ^2.

ImageNet-21kに正則化を加えて事前学習した場合、Mixerは他のモデル ^2 に比べて若干劣るものの、全体的に高い性能を達成した(ImageNetで84.15% top-1)。

他のモデル、には注釈がありましたので、下記に記載しました。

Annotation 2

In Table 2 we consider the highest accuracy models in each class, that use the largest resolutions (448 and above).

表2では、最大の解像度(448以上)を使用した、各クラスで最も精度の高いモデルを記載しています。

However, fine-tuning at smaller resolution leads to substantial improvements in the test-time throughput, with often only a small accuracy penalty.

しかし、より小さな解像度で微調整を行うと、テスト時間のスループットが大幅に向上しますが、多くの場合、精度面ではわずかなペナルティしかありません。

つまり、これまでに出てきましたが、通常はDownstreamタスクはUpstreamタスクに比較して、解像度が大きい、よりリッチな画像を用いて学習することで微調整するのが慣習だが、意外にも、解像度の小さい画像でそれをやっても、精度がちょっと落ちるだけで、その代わり計算処理能力が大幅に上がる、といってます。

For instance, when pre-training on ImageNet-21k the Mixer-L/16 model fine-tuned at 224 resolution achieves 82.84% ImageNet top-1 accuracy at throughput 420 img/sec/core;

例えば、ImageNet-21kで事前トレーニングを行った場合、解像度224で微調整されたMixer-L/16モデルは、420img/sec/coreのスループットで82.84%のImageNet top-1精度を達成しました。

つまり、表ではFine-Tuningは解像度448で微調整して、ImNetの精度84.15%でしたが、解像度224に落としたもので微調整したとしても、82.84%に少し下がるものの微差であり、一方で処理能力(img/sec/core)は105が420と4倍に上がるので、ほら意外とわるくないでしょ、と言えそうです。

the ViT-L/16 model fine-tuned at 384 resolution achieves 85.15% at 80 img/sec/core [14];

ちなみにViTでもやってみた文献がある。
表では同じく解像度448→384で微調整したViT-L/16モデルでは、32→80img/sec/coreで3倍弱に性能UPしつつも、精度は85.30%→85.15%と微小に下がっただけを達成しています[14]。

and HaloNet fine-tuned at 384 resolution achieves 85.5% at 258 img/sec/core [49].

HaloNetでも同様で、解像度448→384にしたら、精度は85.8→85.5%とほとんど変わらなかったが、スループットは120→258と倍増した。

Regularization in this scenario is necessary and Mixer overfits without it, which is consistent with similar observations for ViT [14].

参考文献によると、この方法の場合、Regularization(行列の正則化処理。正則行列にする処理で正規化とは異なる処理)が必要で、MixerはRegularizationなしだと、実際オーバーフィット(学習データに異常に適用しすぎて、汎用データに弱くなる過学習)してしまう。

We also report the results when training Mixer from scratch on ImageNet in Table 3 and the same conclusion holds. Supplementary B details our regularization settings.

また、ImageNetで、Mixerをゼロから学習させた場合の結果を表3に示しますが、同様の結論が得られました。補足Bでは、正則化の設定を詳しく説明しています.

同様といわれると、表3は縦に長い表のわりに乱暴な説明だなあと思うのですが、実際に、表の2グループ目のImageNet-21kのグループを見てみると、確かに with extra regularizationと書かれていて、正則化プロセスが追加されてる感があり、それぞれのモデルの精度も、MixerはViTには微妙に及ばず・・・ながらもほとんど変わらないレベルをたたき出しつつ、そのスループットはViTよりも格段に上、という感じです。つまり、事前学習ありのパラメータを使わずに、念のため一から学習させた場合も、Mixerは確かにViTなどと比較してCompetitiveだなとは思います。

正則化と正規化を同じだと思いがちなのですが、まったく異なるものであり、正規化は対象行列の各要素を0~1の間に収まるように変更することで、正則化は正則化項を呼ばれる項を追加することで対象行列が正則行列になるように調整することであり、上述した通り、正則化を行うと、過学習が抑制されることが知られている。

こちらの記事が参考になりました。

https://qiita.com/7y2n/items/96abd25319786c9445d2
https://qiita.com/ryouka0122/items/a7fbad253680bb7f815e

また、この正則化、具体的には補足Bにあるということですが、その中のWeight Decay(重み減衰)とDropoutがあたります。

正則化のテクニックはこちらの記事が参考になりました。

https://www.tensorflow.org/tutorials/keras/overfit_and_underfit?hl=ja
https://reference.wolfram.com/language/tutorial/NeuralNetworksRegularization.html.ja?source=footer

When the size of the upstream dataset increases, Mixer’s performance improves significantly.
In particular, Mixer-H/14 achieves 87.94% top-1 accuracy on ImageNet, which is 0.5% better than BiT ResNet152x4 and only 0.5% lower than ViT-H/14.

上流の事前学習のデータセットのサイズが大きくなると、Mixerのパフォーマンスは大きく向上します。

特に、表2の中段にあるMixer-H/14というモデルで見てみると、ImageNetでは、87.94%のtop-1精度を達成しており、これはBiT ResNet152x4よりも0.5%良く、ViT-H/14よりも0.5%低いだけです。

Mixer-H/14のH/14は既出ですが、今回のモデルにおけるMixerレイヤーのレイヤー数に応じて呼び方を定義していて、Hはそのレイヤー数が下記の表1のように32のモデルです。

Remarkably, Mixer-H/14 runs 2.5 times faster than ViT-H/14 and almost twice as fast as BiT.

驚くべきことに、Mixer-H/14はViT-H/14の2.5倍、BiTの約2倍の速さで動作しています。下図のマーカー部分ですね。

Overall, Figure 2 (left) supports our main claim that in terms of the accuracy-compute trade-off Mixer is competitive with more conventional neural network architectures.

全体的に見て、図2(左)は、精度と計算のトレードオフの観点から、ミキサーが従来のニューラルネットワークアーキテクチャと競合するという我々の主張を支持しています。

横軸の計算時間を増やせば、他のBiTやMPLなどのモデルのように精度は90%に向けて上がっていくが、その精度を多少犠牲にすれば、計算時間を圧倒的に省力化できるからですね。

The figure also demonstrates a clear correlation between the total pre-training cost and the downstream accuracy, even across architecture classes.

また、この図では、アーキテクチャのタイプが異なっていても、事前学習の総計算時間的なコストとDownstreamタスクの精度には明確な正の相関関係があることを示しています。

右肩あがりだからですね。ただし、横軸はTPUv3を使って1コアあたり何日計算にかかるか時間ですが、対数表示になっているので、実際は、かなり計算時間的なコスト(そしてそれは実際の金銭的なコストにもつながる)を支払って、ようやくちょっと性能があがるんだ、という感覚値になります。そうとらえると、少ない計算時間でそこまで負けない性能を達成できるということは、かなり魅力である、Competitiveといっていいでしょう。

BiT-ResNet152x4 in the table are pre-trained using SGD with momentum and a long schedule.

表のBiT-ResNet152x4は、momentum項ありのSGD(最適化アルゴリズムの1種で、Stochastic Gradient Descentの略。日本語だと確率的勾配降下法)を使って長時間のスケジュールで事前学習したものです。

Since Adam tends to converge faster, we complete the picture in Figure 2 (left) with the BiT-R200x3 model from Dosovitskiy et al. [14] pre-trained on JFT-300M using Adam.

Adamの方が早くconverge(収束)する傾向があるので、JFT-300Mで事前学習したDosovitskiyら[14]のBiT-R200x3モデルに、最適化アルゴリズムはAdamを使って、図2(左)の図を完成させました。

This ResNet has a slightly lower accuracy, but considerably lower pre-training compute.

このResNet(Bitを指す)は、精度はやや低いが、事前の学習計算量はかなり少ない。

ResNetはResidual Network(残差接続)を持つアーキテクチャーのことで、Bit(Big Transfer)もそのアーキテクチャーを持つため、言い換えられています。

Finally, the results of smaller ViT-L/16 and Mixer-L/16 models are also reported in this figure.

ViT-L/16とMixer-L/16の小型モデルの結果もこの図で報告されています。

3.2 The role of the model scale

The results outlined in the previous section focus on (large) models at the upper end of the compute spectrum.

前節の結果は、計算能力敵に上位に位置する(大規模な)モデルに焦点を当てたものでした。

We now turn our attention to smaller Mixer models.

ここで、これまでより小型なMixerについて着目してみる。

We may scale the model in two independent ways:

そもそも、モデルをスケール(大きくする=性能をあげる)しようとした時、2つの異なるアプローチがある。

(1) Increasing the model size (number of layers, hidden dimension, MLP widths) when pre-training.

1つは、事前学習時のモデルのサイズを大きくすること。サイズとは、レイヤーの数、隠れ層の次元数、MLPブロックの幅(おそらくMLP自体の入力次元のこと)を指す。

(2) Increasing the input image resolution when fine-tuning.

もう1つは、Fine-Tuningに使用する入力画像の解像度をあげること。解像度とは1インチあたりのドット数もしくはピクセル数のこと。

While the former affects both pre-training compute and test-time throughput, the latter only affects the throughput.

前者の方法は、事前学習の計算時間と検証時の処理能力両方に悪影響を与えてしまうのに対して、後者は処理能力だけにしか影響を与えない。

Unless stated otherwise, we fine-tune at resolution 224.

特に断りのない限り、解像度224で微調整しています。

We compare various configurations of Mixer (see Table 1) to ViT models of similar scales and BiT models pre-trained with Adam.

Mixerの様々な構成(表1参照)を、同規模のViTモデルやAdamで事前に学習したBiTモデルと比較します。

The results are summarized in Table 3 and Figure 3.

その結果が表3と図3です。


When trained from scratch on ImageNet, Mixer-B/16 achieves a reasonable top-1 accuracy of 76.44%.
This is 3% behind the ViT-B/16 model.

ImageNetで、ゼロから学習した場合、Mixer-B/16は76.44%という妥当なトップ1精度を達成した。表3の1行目です。
これは、ViT-B/16モデルと比べて3%の差があります。

The training curves (not reported) reveal that both models achieve very similar values of the training loss.
In other words, Mixer-B/16 overfits more than ViT-B/16.

トレーニングカーブ、つまり学習カーブ(本論文に掲載はない)を見ると、どちらのモデルもトレーニングロスの値は非常に似通っています。
つまり、ViT-B/16よりもMixer-B/16の方が過学習しているのです。

For the Mixer-L/16 and ViT-L/16 models this difference is even more pronounced.

Mixer-L/16とViT-L/16では、この違いがさらに顕著になります。

As the pre-training dataset grows, Mixer’s performance steadily improves.

事前学習用データセットの増加に伴い、Mixerの性能は着実に向上しています。

Remarkably, Mixer-H/14 pre-trained on JFT-300M and fine-tuned at 224 resolution is only 0.3% behind ViT-H/14 on ImageNet whilst running 2.2 times faster.

驚くべきことに、JFT-300Mで事前に学習され、224の解像度で微調整されたMixer-H/14は、ImageNetでViT-H/14に0.3%劣っているだけで、逆に2.2倍の速度で動作しています。


Figure 3 clearly demonstrates that although Mixer is slightly below the frontier on the lower end of model scales, it sits confidently on the frontier at the high end.

図3は、Mixerがモデルサイズの小さい方ではフロンティア(いわゆるトップの記録)をわずかに下回っているものの、モデルサイズの大きい方では自信を持ってフロンティアに位置していることを明確に示しています。

3.3 The role of the pre-training dataset size

The results presented thus far demonstrate that pre-training on larger datasets significantly improves Mixer’s performance.

ここまで書いてきた内容の通り、より大きなデータセットで事前学習を行うことで、Mixerのパフォーマンスが大幅に向上することが示されています。

Here, we study this effect in more detail.

ここでは、この効果をより詳細に調べてみました。

To study Mixer’s ability to make use of the growing number of training examples we pre-train Mixer-B/32, Mixer-L/32, and Mixer-L/16 models on random subsets of JFT-300M containing 3%, 10%, 30% and 100% of all the training examples for 233, 70, 23, and 7 epochs.

事前学習データが増えるとMixerの性能が向上するのかどうか調べるために、JFT-300Mの全学習データの3%、10%、30%、100%を使ったランダムなサブセットに対して、Mixer-B/32、Mixer-L/32、Mixer-L/16モデルで、233、70、23、7エポックで事前学習した。

通常、今回もそうですが、学習データが N 個ある場合に、それは n 個のデータであるバッチにわけます。この時、バッチ数は N/n 個となります。このバッチごとに学習をして、重みのパラメータの最適化をします。それが N/n 回終わったら1通り1回回したことになるので、それで1エポックとなります。2エポック目以降は、バッチの中身をランダムにシャッフルして、再度繰り返していきます。

Thus, every model is pre-trained for the same number of total steps.

このように、どのモデルも同じステップ数で事前学習を行います。

We use the linear 5-shot top-1 accuracy on ImageNet as a proxy for transfer quality.

Top-1の精度としては、5回実施した中で一番良かった値を使用しています(as a proxy for transfer qualityは意味がわからず割愛・意訳)

For every pre-training run we perform early stopping based on the best upstream validation performance.

すべての事前トレーニングの実行において、上流の検証での最高のパフォーマンスに基づいて、早期の停止(過学習を防ぐため)を行います。

早期の停止とは、何に比べて早期?という感じですが、学習用データに対してエポック数が進むにつれて、検証用データに対する性能もよくなっていくものですが、ある程度進むと、学習用データに適合しすぎて繁華性能が落ち始め、検証用データに対する性能が劣化し始めるタイミングがあります。これを過学習と呼んでいて、過学習が起き始める前にトレーニングを停止することを早期停止と表現しています。

Results are reported in Figure 2 (right), where we also include ViT-B/32, ViT-L/32, ViT-L/16, and BiT-R152x2 models.

結果を図2(右)に示します。ここでは、ViT-B/32、ViT-L/32、ViT-L/16、BiT-R152x2のモデルも含めています。

When pre-trained on the smallest subset of JFT-300M, all Mixer models strongly overfit.

JFT-300Mの最小サブセットで事前学習した場合、すべてのMixerモデルが強くオーバーフィットした。

性能が下がってしまう前に学習を止めているので、グラフ上は結果的に事前学習のデータセットのサイズが大きくなると、性能があがる、という結果のグラフになっています。

横軸は事前学習のデータセットサイズで、データセットはJFT300Mというデータセット。左からそのデータセットから3%、10%、30%、100%抽出したサブセットで行った場合を横軸で表現してます。また、それぞれ、学習エポック数は左から、233回、70回、23回、7回と減らした上でこの曲線なので、データセットを大きくすると、必要なエポック数は少なくなったようです。(つまり、例えば3%しか使わなかった場合はエポック数=233回が一番よい結果だったが、データセットを増やすと、同じ回数、例えば233回もいらなかったという結果なので、つまり過学習している、といえる)

BiT models also overfit, but to a lesser extent, possibly due to the strong inductive biases associated with the convolutions.

BiTモデルもオーバーフィットを起こしますが、その程度は低く、おそらく畳み込みに伴う強い誘導バイアスが原因です。これはグラフからはわからない情報だと思いますが、恐らく必要なエポック数の減衰がゆるやかだったのではと思います。

また、詳細書かれてませんが、強い誘導バイアス、つまりバイアスの値が大きかったことにより、重み減衰の効果が得られ、結果、過学習を防ぐ効果が相対的にあった、ということだと思います。

As the dataset increases, the performance of both Mixer-L/32 and Mixer-L/16 grows faster than BiT; Mixer-L/16 keeps improving, while the BiT model plateaus.

データセットの増加に伴い、Mixer-L/32とMixer-L/16の性能はBiTよりも速く成長する。特に、BiTモデルが停滞しても、Mixer-L/16は向上し続けている。これはグラフの傾きが明らかに黒の点線のBiTより、Mixer-Lを示すオレンジと緑の折れ線の方が大きいからですね。

The same conclusions hold for ViT, consistent with Dosovitskiy et al. [14].

同じ結論がViTにも当てはまり、Dosovitskiyら[14]と一致しています。

However, the relative improvement of larger Mixer models are even more pronounced.

しかし、Mixerの大型化による相対的な向上(率の上昇)はさらに顕著である。最も大型なのが緑の折れ線のMixer-L/16です。

The performance gap between Mixer-L/16 and ViT-L/16 shrinks with data scale.

Mixer-L/16とViT-L/16の性能差は、データ規模に応じて縮小します。緑の実線と点線が、右にいくにつれて、確かに均衡してきます。

It appears that Mixer models benefit from the growing pre-training dataset size even more than ViT.

ViTよりもMixerの方が、事前学習用データセットのサイズが大きくなることにメリットを感じているようです。

One could speculate and explain it again with the difference in inductive biases.

帰納的バイアスの違いで再度説明することもできます。

Perhaps, self-attention layers in ViT lead to certain properties of the learned functions that are less compatible with the true underlying distribution than those discovered with Mixer architecture.

おそらく、ViTのSelf-Attentionレイヤーは、Mixerアーキテクチャで発見されたものよりも、学習された関数の特性が真の基本的な分布とは一致しないものになるからなのでしょう。

3.4 Visualization

It is commonly observed that the first layers of CNNs tend to learn Gabor-like detectors that act on pixels in local regions of the image.

一般的に、CNNの最初の層では、画像の局所領域のピクセルに作用するGaborフィルターのような検出器を学習する傾向があると言われています。

Gaborフィルターはこちらが参考になりましたが、特定の向きをよく検出するフィルターのようです。

http://zellij.hatenablog.com/entry/20131003/p1
https://jp.mathworks.com/help/images/texture-segmentation-using-gabor-filters.html

In contrast, Mixer allows for global information exchange in the token-mixing MLPs, which begs the question whether it processes information in a similar fashion.

一方、Mixerでは、トークン混合MLPで(行列のある列のどの行の要素に対しても)グローバルな情報交換が可能であるのだが、CNNと同様の方法で情報を処理しているのではないかという疑問がわいてくる。

Figure 4 shows the weights in the first few token-mixing MLPs of Mixer trained on JFT-300M.

図4は、JFT-300Mで学習したMixerの最初の数個のトークン混合MLPの重みを示したものである。

Recall that the token-mixing MLPs allow communication between different spatial locations (see Figure 1).

トークン混合MLPは、(ある列の中の)異なる空間位置(行)間のコミュニケーションを可能にすることを思い出してください(図1参照)。

Some of the learned features operate on the entire image, while others operate on smaller regions.

学習された特徴は、画像全体を対象とするものと、より小さな領域を対象とするものがあります。

The first token-mixing MLP contains many local interactions, while the second and third layers contain more mixing across larger regions.

第1層のトークン混合MLPには、(1つのマス内において)局所的に反応する重みが多く可視化されていますが、第2層と第3層には、より広い領域での混合(反応部位)が多く含まれています。

Higher layers appear to have no clearly identifiable structure.

高層部(深層部)は明確に識別できる構造(斜めの線を見分けやすいとか)を持たないように見える。

Similar to CNNs, we observe that many of the low-level feature detectors appear in pairs with opposite phases [38].

CNNと同様に、低レベル(絵の上の方)の特徴検出器の多くは、反対の位相を持つペアで現れることが観察される[38]。

In Supplementary D, we show that the structure of learned units depends on the hyperparameters and input augmentations.

補足Dでは、学習したユニットの構造は、ハイパーパラメータと入力補強に依存することを示しています。

In the linear projection of the first patch embedding layer we observe a mixture of high and low frequency filters; we provide a visualization in Supplementary Figure 6.

第1パッチ埋込層の線形射影では、高域と低域のフィルタが混在していることがわかります。補足図6にその様子を示します。

おわり

「MLP-Mixer」を解説するシリーズ4日目は以上です。
感想や要望・指摘等は、本記事へのコメントか、TwitterのリプライやDMでもお待ちしております。

シリーズ関連記事はこちら

https://zenn.dev/attentionplease/articles/532a3de6308f57
https://zenn.dev/attentionplease/articles/7a11a56d767280
https://zenn.dev/attentionplease/articles/df6170f8581b71
https://zenn.dev/attentionplease/articles/7a3e74ad1bc9bf
https://zenn.dev/attentionplease/articles/a0d88939f9ceed
https://zenn.dev/attentionplease/articles/719580daf5a2d1