恥ずかしいほどシンプルなVision Transformer

Transformer 2022年01月04日

3つの要点
✔️ ViTの肝は、MetaFormerというメタ構造
✔️ パラメータのないPooling層を用いるPoolFormerを提案
✔️ PoolFormerが少ないパラメータ数で比較手法より高い精度を実現

MetaFormer is Actually What You Need for Vision
written by Weihao Yu, Mi Luo, Pan Zhou, Chenyang Si, Yichen Zhou, Xinchao Wang, Jiashi Feng, Shuicheng Yan
(Submitted on 22 Nov 2021 (v1), last revised 29 Nov 2021 (this version, v2))
Comments: Published on arxiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV); Artificial Intelligence (cs.AI); Machine Learning (cs.LG)

code：

本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。

はじめに

様々なタスクにおいてSOTAを更新しまくったVision Transformer (ついに、ViTがビデオ認識の分野にやってきた)が2021年のホットトピックの一つでした。その後なぜViTがうまくいっているのか？について多くの分析が行われていました。

参考図右のアーキテクチャを見てください。一般的にはTransformer構造にあるAttention (注意機構)が入力トークンの情報を統合する重要な役割を担う部分(Token information mixing)だと考えられました。よってAttentionの計算複雑度O(N^2)がボトルネックであるため、大変多くの研究が計算を軽くすることに取り組みました。

しかし、2021年3月にGoogleの研究者がAttentionではなく、単にMLPを用いてもSOTA程度の精度が出せることを明かしました(AttentionがないのにTransformerは活躍するのか)。その後MLPを用いたTransformerに関する研究も盛り上がってきて、ViTの何が本質的に重要なのか？という謎がさらに深まりました。

本記事では、『ViTの成功の鍵は、本当にToken information mixingのアーキテクチャなのか？』この問いに答えを求めた研究について紹介します。

衝撃な結果ですが、『Token information mixningがどうでも良い！情報さえ共有できれば、Pooling層だけでも良い！』ということが示されました。この気恥ずかしいほどシンプル(著者はembarrassingly simpleと表現しています)なPooling層を用いたPoolFormerの有効性を多くの実験を通じて示されています。

以降第2節で提案手法を紹介してから、第3節で重要な実験について紹介します。

提案手法

MetaFormer

論文では図(a)のMetaFormerという新たな概念を提案しました。

MetaFormerはアーキテクチャのコンセプトであり、具体的なアーキテクチャではないです。つまり、(a)のMetaFormerに示すメタ構造を満たすもの、例えば今まで研究されてきたAttention構造を用いたTransformer、MLPを用いたMLP-likeモデルと本研究で提案したPooling層を用いるPoolFormerが全部MetaFormerと呼びます。

MetaFormerはすごくシンプルです。まずは入力IをPatch Embeddingを行ってXを得ます (式1)。

次にXを正則化してToken Mixerに入力します。Token MixerではXの各Patch間の情報を融合します。ここでは残差接続を利用して、Yを得ます(式2)。

さらに、Yを正則化してMLPおよび活性化関数を通した後、残差接続で入力と出力を繋げます(式3)。

PoolFormer

MetaFormerはPython言語における抽象クラスという概念に類似しています。PoolFormerは(式2)におけるTokenMixerをPooling層で実装したものであり、MetaFormerをインスタンス化した一つに過ぎません。

Poolingは単にK*Kの値を平均したもの出力するだけです。ただし、(式4)の最後にはTを引き算しているのは、後に残差接続(式2の+X)を行うためその調整を行うという意図です。式4からも分かるように、当初の計算複雑度O(N^2)のAttentionを学習パラメータなしのPoolingに置き換えました。

本論文で提案するアーキテクチャの紹介は以上となるので、呆れるほどシンプルであることも実感していただけたでしょうか？最後にPytorchのサンプルコードと、入力から出力までの型の変化を可視化した図2を参考までに載せておきます。

実験：MetaFormerの有効性を検証

Image classification

実験ではImageNet-1K (1.3Mの訓練画像と50Kの検証画像で1Kクラス分類)において検証を行いました。MixUp, CutMix, CutOut, RandAugmentと4種類のデータ拡張を用いました。Optimizer AdamW, decay rate=0.05, batch size=4096, 学習率 lr = 0.001 * batch size / 1024で300epochsを訓練しました。またwarmup epoches = 5のcosine scheduleを用いて学習率を減衰さえました。さらにlabel smoothing = 0.1 に設定しています。label smoothingについては(label smoothingに隠された真実！)を参考してみてください。

長々と実験設定を紹介しましたが、実験結果はTable 2に示した通り、スッキリしています。Token Mixer別にモデルをグループ分けしています。評価指標はパラメータ数Params (M), 積和演算数MACs (G), Top-1精度(%)です。またRSB-ResNetは'ResNet Strikes Back'という訓練方法を改善したResNetを300epochs訓練したものになります。

Table 2最後のブロックに『少ないパラメータ数で高い精度』を出せるという結果を示しています。例えばパラメータ数が 21Mと31Mの小さなPoolFormerでも、80.3%と81.4%のTop-1精度を達しており、多くのAttentionベースとMLPベースのMetaformerを超えていることが分かります。

さらに図3の可視化から、PoolFormerが少ない計算量MACsとパラメータ数を持ちながら、比較手法より高い精度を出していることが一目瞭然です。

Object detection and instance segmentation

次にPoolFormerを110Kの訓練データと5Kの評価データを持つCOCOva12017データセットで検証しました。COCOの物体検出(Object Detection)とインスタンス分割(Instance Segmentation)のタスクに対して、それぞれRetinaNetとMask R-CNNをヘッドに持ったPoolFormer(Backbone)で実験をしました。

Table 3とTable 4はそれぞれの結果を示しています。各評価指標においていずれベースラインResNetより『少ないパラメータ数で高い精度』を出しています。