ついに画像認識のBERTが実現？！MAEについて
3つの要点
✔️ Vision Transformer (ViT)を用いてCVのBERTを目指す
✔️ Encoderに75%のパッチをマスクしてDecoderにTransformerを用いるMAEを提案
✔️ ImageNet-1kのラベルなしデータで学習した表現が、初めて87.8%の精度を達成
Masked Autoencoders Are Scalable Vision Learners
written by Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, Ross Girshick
(Submitted on 11 Nov 2021 (v1), last revised 19 Dec 2021 (this version, v3))
Comments: Published on arxiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV)
code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
はじめに
画像認識(CV)分野においては深層学習が大きな成功(ViT, Vision Transformersの未来に託す研究)を挙げてきました。その中でラベル付き画像が重要な役割を担いました。一方で自然言語処理(NLP)分野では自己あり教師学習が凄まじい成果を上げています。例えば、1000億単位のパラメータを持つ自己回帰モデルGPT-3(GPT-3の真価を引き出す方法 : Promptプログラミング)やMasked AutoEncoderモデルBERTなどが有名です。
CVとNLPにおけるこの差は一体どんなところから生まれたのでしょうか？本記事では言語と画像の違いについて分析を行った上で、その差を埋めるMasked AutoEncoder (MAE)を提案した論文について紹介します。
CVではMasked AutoEncoderの一つであるDenoising Autoencoderは前から研究されています。ただCVにおけるBERTみたいな流行りを埋め出せていません。その原因について以下のように考えられます。
1. アーキテクチャの違いです。CVは最近まで畳み込みネットワークCNNに支配されており、CNNは画像の各領域から情報を抽出するためマスク機構が取り入れるのが難しいです。ただ、この問題はVision Transformer (ViT)によって解決されました。
2. 情報の密度の違いです。自然言語は人間が生み出したものであり、高い抽象度や単位あたりの情報密度が濃いです。よってNLPでは単語にマスクをかけて予測させるのは難易度の高いタスクになるが、CVの情報密度が疎です。したがって、本論文では画像に非常に高い比率のマスクをかけることを提案します。難易度の高い予測タスクがより広域な情報に注目するように促す効果が期待できます。
3. AutoEncoderのDecoderの役割が異なります。NLPではDecoderは表現空間から単語を予測するにはシンプルなMLPで行うが、CVではPixelレベルの再現を行わなければいけません。本論文のDecoderはTransformerを用います。
以上の３点から、著者らが画像表現学習のためにMasked AutoEncoder (MAE)を提案します。Vision Transfer (ViT)がTransformerをそのままCVに適応したというなら、MAEはBERTをそのままCVに適応したということが言えるため、BERTがもたらしたNLP革命をCVに期待しています。 最後にMAEの結果 (Figure 2)をお見せしておきましょう。
Masked AutoEncoder (MAE)
MAEは入力の一部から残りの部分を再構築する構造で、Encoderで入力を潜在空間に射影して、Decoderで潜在空間から入力空間に射影するAutoEncoder (Figure 1)です。それぞれについて説明をします。
まず画像を重複のないようにパッチに分割します。そして、高い割合(75%)のパッチをランダムでマスクをかけます。MAE encoder (ViT)にマスクされていないパッチのみを入力します。具体的には各パッチに対して線形射影とPositional Embeddingを行います。
MAE decoderでは、(1)エンコードされたパッチと(2)マスクのトークンを入力とします。マスクトークンは共通な学習されたものであり、予測すべきパッチを表しています。またDecoder全てのトークンにPositional Embeddingを足します。Decoderは事前学習時にのみ用いるため任意なモデルで良いが、MAEではEncoderの1/10計算量の軽量なTransformerを用います。さらに、Decoderは256次元のベクトルを出力して16 x 16に変換してから正解パッチとMSE誤差で学習します。
Experiments
MAEのBackboneはViT-Large/16を用います。実験では先行研究で示されたViT-L/16の精度76.5%をより強く制限を加えて82.5%まで改良しました。MAEはImageNet-1Kで自己教師ありの事前学習を行った後、fine-tuning (ft)とはLinear Probing (lin)で学習した表現の評価を行います。
ImageNetでのアブレーション実験
MAEのアブレーション実験結果をTable 1に示します。
(a)はDecoderの深さがtfの結果に与える影響が小さく、blocks=8の時lin (最後の一層だけfine-tuning)が最も良い結果を示します。またftはlinより全体的に精度が高いことがわかります。Decoderの幅が512が良いことが(b)で示されています。
面白いことに(c)よりEncoderの入力にMaskedされたパッチを使わない方が計算量が少ないだけでなく、精度も良いことが分かります。またFigure 5はマスクの割合を調べた図であり、ftの場合は40%以上が良く、linの場合は75%程度のマスクが最も精度が良いことが言えます。
(d)は再構築の目的を調べてもので、dVAEのように離散的なトークンに射影することも精度が良いことがわかります。ただpixel空間で正則化を行うことで同等な精度が得られるため、シンプルなpixel空間を採用します。
(e)はデータ拡張を調べたもので、シンプルにrandom sizeの切り抜きだけで精度が良いことがわかります。
(f)からランダムにマスクをサンプリングした結果が良いことが分かります。Figure 6に異なる種類のマスクを可視化したものになります。
Table 2は計算時間を調べたもので、Decoderの深さが1の時に同等な精度を出しながら約4倍のスピードアップが確認できます。最後にFigure 7は事前学習のepoch数が多ければ多いほど精度がまだ上がることを示しています。
ImageNetでの比較実験
Table 3に比較手法の精度を示しており、MAEが最も良いことが分かります。
Figure 8はViTの事前学習データの違いを調べたものです。最も精度が良いのはJFT300Mで学習した表現であり、それと近い精度を出しているのはImage-1Kで学習したMAEです。これは少ないデータ数かつラベルを使わないで学習したMAEの有効性を示しています。
ImageNetでの部分的なFine-tuning実験
Figure 9ではFine-tuningする層の数を調べました。4層までのfine-tuningは大きく精度に貢献していることが確認でき、これはトップ4層の表現はタスクに大きく相関していることを示しています。
その他のデータセットにおける転移学習実験
MAEといった表現学習の目的は下流タスクに適応することであるため、本研究ではImageNet以外の分類タスク(Table 6)、COCO物体検出(Table 4)とADE20Kセグメンテーションタスク(Table 5)にて実験を行いました。
いずれの表からもMAEがSOTAを達していることが分かります。
まとめ
本記事はMasked AutoEncoder (MAE)について紹介しました。ViTを利用してBERTのようなモデルを画像認識で実現しました。特徴は2つで、Encoderに高いレート（75%）のパッチをマスクすることでタスク難易度を上げました。またDecoderはPixelレベルの予測を行うためTransformerを利用しました。さらにViT以降のテクニックを取り入れた結果、ImageNet-1Kのラベルなしデータセットで事前学習した表現が、87.8%と教師あり学習のViTを超えた結果が得られました。
本研究がラベル使わずに良い表現を学習できることを示し、特に生成モデルが画像表現学習における可能性を示したことが大きな貢献であると考えており、MAEを紹介させてもらいました。
この記事に関するカテゴリー
