Masked Autoencoders (MAE)

Masked autoencoder（MAE）は、画像を対象とする自己教師あり学習手法です。モデルは画像パッチの大部分（通常は75%）をランダムに隠し、残りの可視部分から欠損ピクセルを再構成するように学習します。Meta AIのKaiming Heらは、BERTのマスク単語予測をビジョントランスフォーマーに応用し、2021年の論文Masked Autoencoders Are Scalable Vision Learnersでこの手法を発表しました。エンコーダーは可視パッチだけを処理するため、MAEによる事前学習は、すべてのパッチを処理する場合より3倍以上高速でした。この方法で事前学習したViT-Hugeモデルは、ImageNetで87.8%の精度を達成し、当時、ImageNet-1Kのデータだけを使った手法の中で最高の結果となりました。事前学習済みエンコーダーは、分類、検出、セグメンテーションのバックボーンとして利用できます。

Masked Autoencoderの仕組み#

MAEでは、画像を重なり合わないパッチのグリッドに分割します。224x224の入力では16x16ピクセルのパッチが196個になり、その大部分をランダムに隠します。エンコーダーとデコーダーは、意図的に異なる構成にされています。標準的なViTであるエンコーダーは可視パッチだけを処理するため、計算量とメモリ使用量を約4分の1に削減できます。デフォルトでは、軽量なデコーダーはTransformerブロックを8つ備え、トークンあたりの計算量はエンコーダーの10%未満です。デコーダーは、エンコードされた可視パッチと、隠された各位置に共通の学習可能なマスクトークンを受け取り、各隠れパッチのピクセル値を予測します。損失関数には隠されたパッチのみの平均二乗誤差を使い、パッチごとに正規化したピクセルをターゲットにすると、学習される特徴がさらに向上しました。事前学習後、デコーダーは破棄します。

高いマスク率が、このタスクを有用なものにしています。隣接する画像パッチには高い冗長性があるため、隠すパッチが少ないと、モデルは周囲から補間して各領域を埋められます。75%を隠すと、断片から物体全体やシーンを推測する必要が生じ、形状、部位、文脈の学習が促されます。一方、BERTがマスクする単語は15%だけです。各単語がすでに豊富な意味を持っているためです。

MAEの中核となるマスキング処理を、ここではPyTorchで4枚の画像のバッチを使って示します。

import torch

torch.manual_seed(0)
images = torch.rand(4, 3, 224, 224)  # a batch of 4 RGB images

# Patchify: each 224x224 image becomes 196 patches of 16x16 pixels, flattened to 768 values
patches = images.unfold(2, 16, 16).unfold(3, 16, 16).permute(0, 2, 3, 1, 4, 5).reshape(4, 196, 768)

# Random masking as in the MAE reference code: rank patches by random noise and keep the first 25%
keep = int(196 * (1 - 0.75))
ids_keep = torch.argsort(torch.rand(4, 196), dim=1)[:, :keep]
visible = torch.gather(patches, 1, ids_keep.unsqueeze(-1).expand(-1, -1, 768))

# Mask of hidden patches (1 = hidden): the decoder predicts these, and the loss is averaged over them only
mask = torch.ones(4, 196).scatter(1, ids_keep, 0)
print(f"Encoder input: {visible.shape[1]} of {patches.shape[1]} patches per image")
print(f"Reconstruction targets: {int(mask[0].sum())} hidden patches per image")

Encoder input: 49 of 196 patches per imageとReconstruction targets: 147 hidden patches per imageが出力されます。したがって、エンコーダーのTransformer層が処理するシーケンスは、画像全体の4分の1の長さになります。アテンションは各トークンを他のすべてのトークンと比較するため、アテンションスコア行列のサイズは、196個のパッチすべてを使う場合の16分の1になります。

結果とバリアント#

MAEの特徴は、ファインチューニング後に最も強力になります。線形プロービングでは、MAEで事前学習したViT-LエンコーダーのImageNet精度は73.5%で、対照学習手法を下回りました。しかし、ファインチューニングしてCOCOの物体検出やセグメンテーションに転移すると、MAEによる事前学習は教師あり事前学習を上回りました。この手法は静止画像以外にも広く応用されています。VideoMAEは動画クリップの時空間パッチの90%以上をマスクし、SatMAEは衛星画像で事前学習を行います。また、ConvNeXt V2は、スパース畳み込みを使ってマスク付きオートエンコーディングを畳み込みネットワークに取り入れています。

マスク付き事前学習は、広く利用されている基盤モデルにも採用されています。Ultralyticsパッケージで動作するMetaのSegment Anything Model (SAM)は、MAEで事前学習したViTを画像エンコーダーに採用しています。そのViT-Bバージョンであるsam_b.ptは、Ultralyticsのauto_annotateでデフォルトのマスクモデルとして使われ、YOLO26検出器のボックスをセグメンテーションラベルに変換します。

Masked autoencoderと関連用語の比較#

• オートエンコーダー：入力全体を小さなコードに圧縮し、再構成します。MAEは不完全な入力を受け取り、見ていない部分を予測する必要があります。また、エンコーダーの出力には可視パッチごとに1つのベクトルが保持されます。
• マスク画像モデリング：MAEが属する、より広い手法群です。BEiTはマスクされたパッチの離散的な視覚トークンを予測し、SimMIMはエンコーダーにマスクトークンを入力します。一方、MAEはピクセルを予測し、エンコーダーにはマスクトークンを入力しません。
• 対照学習：同じ画像の拡張ビューを照合して学習し、デコーダーを必要としません。凍結した状態でより良く機能する特徴が得られますが、MAEの特徴はファインチューニングによってさらに向上します。
• Joint Embedding Predictive Architecture (JEPA)：画像領域を隠す点は同じですが、ピクセルではなく埋め込みを予測するため、モデルの容量をテクスチャの細部に使わずに済みます。