[論文メモ] Masked Autoencoders Are Scalable Vision Learners

arxiv.org

FAIR

ViTは入力画像をパッチにして入力するがBERTのように一部をマスクし、それを復元する自己教師ありの事前学習。


f:id:Ninhydrin:20211118091528p:plain

手法

提案する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が大きいとメモリ消費が大きくなる)。

Reconstruction target

decoderは入力トークンに対応する箇所のパッチをピクセルレベルで復元する。
decoderの出力をreshapeして対応するパッチと同じ形にしてMSE lossで学習。lossを計算するのはBERTと同じくマスクされた部分のみ。
また正規化したパッチの復元タスクも実験。
パッチ内のピクセルの平均・標準偏差を使って正規化したパッチをターゲットとして復元するというもの。特徴量の質が向上したとのこと。

実験

ImageNet-1Kを使って事前学習。
できた特徴量をend-to-endのfine-tuningとlinear probingの2つで評価。
ベースラインモデルはViT-Large(ViT-L/16)。提案論文、スクラッチでの学習、MAEでのfine-tuningのaccを以下に示す。
f:id:Ninhydrin:20211118094543p:plain

クラッチ学習は寄り強い正則化を使ったそう。
fine-tuningは50epoch、スクラッチは200epoch。

Masking ratio

マスク量をどれくらいにすると良いかの比較実験。75%が良さそう。
また図4のように、入力と違うパッチを復元することが多い。
オブジェクトやシーンのゲシュタルトを理解してると推測(単純に線の補完とかではない)。
f:id:Ninhydrin:20211118094839p:plain
f:id:Ninhydrin:20211118095410p:plain

Decoder design

decoderの深さ、広さについて。結果は表1の(a)、(b)。
linear probingはある程度の大きさが必要だが、fine-tuningはあまり影響を受けない。何なら1blockでも十分な性能。

f:id:Ninhydrin:20211118095727p:plain

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は簡単すぎて再構成はきれいだが、出来上がった特徴が不十分。
ランダムが簡単かつ精度がよい。

f:id:Ninhydrin:20211118101400p:plain

Comparisons with self-supervised methods

MoCoやDINO等との比較結果が表3。全体としてMAEが良い。
特筆すべきはViT-Hで448のサイズでfine-tuningするとImageNet1Kだけで87.8%のacc。
f:id:Ninhydrin:20211119085211p:plain

Object detection and segmentation

詳しい実験設定は論文を参照。
f:id:Ninhydrin:20211119090209p:plain

所感

いくつかのタスクに対して既存手法と比較しており、実験のボリュームがすごい(一部省略した)。
BERTとかと同じでマスクした箇所を予測するという自己教師ありタスクで性能は良さそう。
また、encoderとdecoderで入力を変えるというのは省メモリ・高速化に繋がりありがたい。
decoderは小さくても問題ないという実験結果も自分で実験する指針になり助かる。
decoderに入れるときにはマスク部分をトークンで埋めて入力するが、各画像で異なるマスクを使った場合、マスクトークンで埋めるのをきれいに書くのは難しそう。