Skip to main content

MAE

MAE (Masked Autoencoder) は、ViT を backbone にした masked image modeling の代表的手法です。シンプルな設計と高い mask ratio によって、scalable な vision pretraining を確立しました。

核心的な設計

二つの工夫が重要です。

  1. High mask ratio (~75%)
    Language では 15% 程度ですが、image は redundancy が高いため、大胆に隠さないと簡単すぎて学習が進みません。

  2. Asymmetric encoder-decoder
    Encoder は visible patches のみを処理し、mask token は decoder で初めて加わります。これによって encoder の計算量が大幅に減り、scalable な pretraining が可能になります。

なぜ動くのか

Image の隣接 patch は高い相関を持つため、25% の visible patch でも残りの 75% を概ね補えます。Encoder は、その復元に必要な abstract な structure を学ぶことを強制されます。

Pretrain と fine-tune

ImageNet 上で MAE pretrain した ViT-L / ViT-H は、ImageNet 分類、COCO 検出、ADE20K segmentation で強い結果を出します。

数式で見る MAE

画像を patch の集合 {xi}i=1N\{x_i\}_{i=1}^{N} とし、mask された patch の index 集合を M\mathcal{M}、見えている patch の集合を Mˉ\bar{\mathcal{M}} とします。MAE は、encoder EθE_\theta に visible patches だけを入力し、decoder DϕD_\phi が mask token を補ったうえで masked patches を復元します。

x^i=Dϕ ⁣(Eθ(xMˉ),mask tokeni)\hat{x}_i = D_\phi\!\left(E_\theta(x_{\bar{\mathcal{M}}}),\mathrm{mask\ token}_i\right)

Pixel を target にする場合、MAE の loss は masked patch 上の mean squared error です。

LMAE=1MiMx^ixi22\mathcal{L}_{\mathrm{MAE}} = \frac{1}{|\mathcal{M}|} \sum_{i\in\mathcal{M}} \left\|\hat{x}_i-x_i\right\|_2^2

ここで、xix_i は元の patch、x^i\hat{x}_i は decoder が予測した patch です。この式の気持ちは、「見えている約 25% の patch から、隠された約 75% の patch を復元するために、画像全体の構造を encoder に持たせる」というものです。

MAE の設計で特に重要なのは、mask token を encoder に入れず、decoder で初めて追加する点です。ViT の self-attention は token 数に対しておおよそ二乗で重くなるため、encoder が処理する token 数を visible patches に絞ると、事前学習の計算量を大きく下げられます。つまり、高い mask ratio は単に task を難しくするだけではなく、大規模な encoder を現実的に学習するための計算上の工夫でもあります。

MAE の影響

  • Video MAE、3D MAE、point cloud MAE、graph MAE など、多様な領域に拡張
  • JEPA 系の基盤 (mask + predict の枠組み)
  • DINO 系との組み合わせ (DINOv2)
  • Diffusion model の backbone 初期化

SimMIM との違い

観点MAESimMIM
Encoder 入力Visible onlyAll (mask token も入れる)
DecoderLight な ViT decoder1-layer prediction head
Mask ratio高い (~75%)適度
計算効率高いやや劣る

関連ページ

主なソース