[論文レビュー] MixMAE: Mixed and Masked Autoencoder for Efficient Pretraining of Hierarchical Vision Transformers

Jihao Liu, Xin Huang+3 more|arXiv (Cornell University)|May 26, 2022

Advanced Neural Network Applications被引用数 7

ひとこと要約

MixMAE は、マスクされたパッチを2つの画像から得たパッチに置き換えることで、[MASK] テンキーを用いずに二重再構成を可能にする、階層的 Vision Transformer のための新しいマスク自己符号化フレームワークを提案する。これにより、600 フェッチでの pretraining で ImageNet-1K で 85.1% のトップ-1 正答率を達成し、FLOPs/パフォーマンスのトレードオフにおいて、先行する MIM メソッドを上回る。

ABSTRACT

In this paper, we propose Mixed and Masked AutoEncoder (MixMAE), a simple but efficient pretraining method that is applicable to various hierarchical Vision Transformers. Existing masked image modeling (MIM) methods for hierarchical Vision Transformers replace a random subset of input tokens with a special [MASK] symbol and aim at reconstructing original image tokens from the corrupted image. However, we find that using the [MASK] symbol greatly slows down the training and causes pretraining-finetuning inconsistency, due to the large masking ratio (e.g., 60% in SimMIM). On the other hand, MAE does not introduce [MASK] tokens at its encoder at all but is not applicable for hierarchical Vision Transformers. To solve the issue and accelerate the pretraining of hierarchical models, we replace the masked tokens of one image with visible tokens of another image, i.e., creating a mixed image. We then conduct dual reconstruction to reconstruct the two original images from the mixed input, which significantly improves efficiency. While MixMAE can be applied to various hierarchical Transformers, this paper explores using Swin Transformer with a large window size and scales up to huge model size (to reach 600M parameters). Empirical results demonstrate that MixMAE can learn high-quality visual representations efficiently. Notably, MixMAE with Swin-B/W14 achieves 85.1% top-1 accuracy on ImageNet-1K by pretraining for 600 epochs. Besides, its transfer performances on the other 6 datasets show that MixMAE has better FLOPs / performance tradeoff than previous popular MIM methods. Code is available at https://github.com/Sense-X/MixMIM.

研究の動機と目的

• 階層的 Vision Transformer のマスク画像モデリング（MIM）における pretraining の非効率性と pretrain-finetune の不一致を解消すること。
• encoder における [MASK] テンキーの使用を排除することで、学習の遅さと fine-tuning 時との分布の不一致を解消すること。
• Swin Transformer のような階層的 ViT と互換性があり、高い表現品質を維持しながら、学習効率を向上させる手法を開発すること。
• SimMIM や MAE などの既存の MIM メソッドと比較して、より優れた FLOPs/パフォーマンスのトレードオフを達成すること。

提案手法

• MixMAE は、1つの画像のマスクされたパッチを別の画像からの可視パッチに置き換えることで混合入力を生成し、encoder における [MASK] テンキーの使用を回避する。
• 階層的 Vision Transformer の encoder は、混合入力を処理し、元の2つの画像の両方の隠れ表現を生成する。
• デコードの前に、隠れ表現を元に戻し、[MASK] テンキーで埋める。これにより、二重再構成が可能になる。
• 軽量なデコーダーは、混合表現から両方の元の画像を再構成し、特徴品質の向上を図るために二重再構成損失を用いる。
• 複数の画像を混合することで、変動するマスキング比に対応可能であり、4枚の画像を混合した 75% のマスキング比が最適なパフォーマンスを発揮する。
• アプローチはデコーダーでマスク自己注意を用い、encoder では [MASK] テンキーを避けるため、pretrain-finetune の一貫性が保証される。

Figure 1 : Overview of MixMAE. For pretraining, two images are mixed with a random mixing mask to create a mixed image. MixMAE takes the mixed image as input and reconstructs the two original images. Right before decoding, the token embeddings are unmixed and filled with mask tokens for dual reconst

実験結果

リサーチクエスチョン

• RQ1 階層的 Vision Transformer に対して、encoder における [MASK] テンキーの使用を回避するマスク自己符号化アプローチを設計可能か？ これにより、学習効率と一貫性が向上するか？
• RQ2 異なる画像からのパッチにマスクされたパッチを置き換えることで、[MASK] テンキーまたは学習可能なベクトルを使用するのと比較して、表現学習にどのような影響を与えるか？
• RQ3 高いパフォーマンスを最小限の計算コストで達成するための、最適なマスキング比と混合戦略は何か？
• RQ4 混合入力からの二重再構成は、単一画像再構成と比較して、より優れた下流タスクへの転移性能をもたらすか？
• RQ5 MixMAE は、既存の MIM メソッドと比較して、より少ない pretraining フェッチ数で強力なパフォーマンスを達成できるか？

主な発見

• MixMAE は、600 フェッチでの pretraining で ImageNet-1K で 85.1% のトップ-1 正答率を達成し、高い効率性と優れたパフォーマンスを示した。
• 900 フェッチでの pretraining において、ADE20K で 51.0 mIoU を達成し、密度予測タスクへの優れた一般化性能を示した。
• 4枚の画像を混合して得た 75% のマスキング比は、同じ比率で [MASK] テンキーを使用した場合を上回る性能を示し、混合戦略の有効性を確認した。
• 二重再構成は性能向上に顕著に寄与し、特に ADE20K では単一再構成と比較して mIoU を 2.6 ポints 向上させた。
• 画像分類、物体検出、セマンティックセグメンテーションを含む7つの下流ベンチマークにおいて、MixMAE は SimMIM や MAE と比較して、FLOPs/パフォーマンスのトレードオフにおいて優れた性能を示した。
• encoder における [MASK] テンキーの排除により、pretraining の効率性が顕著に向上し、計算の無駄が削減され、pretrain-finetune の不一致が回避された。

Figure 2 : Tradeoffs of FLOPs vs. (left) top-1 accuracy on ImageNet-1K, (middle) AP box on COCO, (right) and mIoU on ADE20K. All results are from various self-supervised pretraining methods followed by supervised finetuning. All entries on COCO [ 31 ] use Mask RCNN [ 22 ] framework. All entries on A