[論文メモ] Masked Autoencoders Are Scalable Vision Learners
FAIR
ViTは入力画像をパッチにして入力するがBERTのように一部をマスクし、それを復元する自己教師ありの事前学習。
手法
提案するmasked autoencoder (MAE)について。
encoderとdecoderの入力は非対称。
encoderはマスクされていないパッチのみを入力、decoderはencoderの出力とマスク部分を入力とする。
Masking
単純なランダムサンプリング。マスクする割合はある程度大きいほうがよい(75%)。
gridとかでマスキングするとセマンティックな再構成ではなく、単純なテクスチャや線の補完になってしまうので良くない。
MAE encoder
ViTと同じく、パッチを線形変換してPositional Encoding(PE)を足して入力とする。
入力するパッチをマスクされていないパッチだけにすることで巨大なViTモデルの学習でもメモリ消費等を抑えられる。
MAE decoder
入力はencoderの出力のトークンと、マスクトークン(図1を参照)。
encoderはマスクしてないパッチのみを入力とするが、decoderはエンコードしたパッチとマスク部分にあたるトークンを入力とする。
マスク部分のトークンはすべて共通かつ学習可能で、decoderはそのトークン部分を復元することになる。
また、マスク部分も補完したトークン列に対してはencoderと同様にPEを足す。
decoderはマスク部分復元タスク(事前学習)にのみ使うのでencoderと独立していれば何を使ってもいい。
実際彼らは非常に小さいモデルを利用した(すべてのトークンが入力になるのでdecoderが大きいとメモリ消費が大きくなる)。
実験
ImageNet-1Kを使って事前学習。
できた特徴量をend-to-endのfine-tuningとlinear probingの2つで評価。
ベースラインモデルはViT-Large(ViT-L/16)。提案論文、スクラッチでの学習、MAEでのfine-tuningのaccを以下に示す。
スクラッチ学習は寄り強い正則化を使ったそう。
fine-tuningは50epoch、スクラッチは200epoch。
Masking ratio
マスク量をどれくらいにすると良いかの比較実験。75%が良さそう。
また図4のように、入力と違うパッチを復元することが多い。
オブジェクトやシーンのゲシュタルトを理解してると推測(単純に線の補完とかではない)。
Decoder design
decoderの深さ、広さについて。結果は表1の(a)、(b)。
linear probingはある程度の大きさが必要だが、fine-tuningはあまり影響を受けない。何なら1blockでも十分な性能。
Mask token
encoderにマスク部分を使うか否か。結果は表1(c)で使った結果のほうが全体的に劣化。使わない方が断然速い。
Reconstruction target
再構成のターゲットパッチをいくつか試した結果が表1(d)。正規化したパッチの復元が良い。dVAEはVAEでトークン化してそれをターゲットとするやり方だが、VAEの学習やトークン化のオーバーヘッドが大きい。
Data augmentation
ランダムサイズ or 固定サイズに対してのcrop(水平flipはあり)が良さそう。BYOLやSimCLRと違って何もなくてもかなり良い。
マスクは毎回異なるし、マスキングだけで十分なデータ増強になっている。
Mask sampling strategy
マスクのやり方を比較したの結果が表1(f)で図6にそれぞれのマスクの例を示す。
blockは難易度が高く(lossも大きかった),75%以上マスクすると精度が下がった。
gridは簡単すぎて再構成はきれいだが、出来上がった特徴が不十分。
ランダムが簡単かつ精度がよい。
Comparisons with self-supervised methods
MoCoやDINO等との比較結果が表3。全体としてMAEが良い。
特筆すべきはViT-Hで448のサイズでfine-tuningするとImageNet1Kだけで87.8%のacc。
Object detection and segmentation
詳しい実験設定は論文を参照。