MixMAE:階層的Vision Transformerの効率的プレトレーニングのための混合およびマスク自動エンコーダ

本稿では、さまざまな階層的Vision Transformerに適用可能な、シンプルながら効率的な事前学習手法である「Mixed and Masked AutoEncoder(MixMAE)」を提案する。既存の階層的Vision Transformer向けのマスク画像モデリング(MIM)手法は、入力トークンのランダムなサブセットを特別な[MASK]記号に置き換え、破損した画像から元の画像トークンを再構成することを目的としている。しかし、本研究では、[MASK]記号の使用が、特に大きなマスク率(例:SimMIMにおける60%)により、学習の速度を著しく低下させるとともに、事前学習と微調整の間に不整合を引き起こすことを発見した。一方、MAEはエンコーダー段階で[MASK]トークンを一切導入しないが、階層的Vision Transformerには適用できない。この問題を解決し、階層モデルの事前学習を高速化するために、本研究では、ある画像のマスクされたトークンを別の画像の可視トークンで置き換える、いわゆる「混合画像(mixed image)」を生成する手法を採用する。その後、混合入力から2つの元の画像を同時に再構成する「二重再構成(dual reconstruction)」を実施することで、大幅な効率向上を達成した。MixMAEはさまざまな階層的Transformerに適用可能であるが、本研究では大窓サイズを有するSwin Transformerを用い、巨大なモデルサイズ(最大6億パラメータ)までスケールアップした。実証的な結果から、MixMAEは高品質な視覚表現を効率的に学習できることを示した。特に、Swin-B/W14を用いたMixMAEは、600エポックの事前学習によりImageNet-1Kで85.1%のトップ1精度を達成した。さらに、他の6つのデータセットにおける転移性能評価から、MixMAEは従来の代表的なMIM手法と比較して、FLOPs対性能のトレードオフが優れていることが明らかになった。コードはhttps://github.com/Sense-X/MixMIMにて公開されている。