[論文レビュー] Masked autoencoders are effective solution to transformer data-hungry

Jiawei Mao, Honggu Zhou+2 more|arXiv (Cornell University)|Dec 12, 2022

Advanced Neural Network Applications被引用数 8

ひとこと要約

本稿では、小規模データセット向けに特化した Small Dataset Masked Autoencoders (SDMAE) を提案する。この手法は、Vision Transformers (ViTs) の性能を向上させるために、MAE デコーダーを単純化し、位置予測タスクとクラストークンに配慮した対照学習タスクを導入する。SDMAE は、ViT のアーキテクチャを変更せずに、過学習を低減し、局所的不変性と表現学習を向上させることで、Tiny-ImageNet や CIFAR-100、医療画像ベンチマークなど、小規模データセットで最先端の性能を達成する。

ABSTRACT

Vision Transformers (ViTs) outperforms convolutional neural networks (CNNs) in several vision tasks with its global modeling capabilities. However, ViT lacks the inductive bias inherent to convolution making it require a large amount of data for training. This results in ViT not performing as well as CNNs on small datasets like medicine and science. We experimentally found that masked autoencoders (MAE) can make the transformer focus more on the image itself, thus alleviating the data-hungry issue of ViT to some extent. Yet the current MAE model is too complex resulting in over-fitting problems on small datasets. This leads to a gap between MAEs trained on small datasets and advanced CNNs models still. Therefore, we investigated how to reduce the decoder complexity in MAE and found a more suitable architectural configuration for it with small datasets. Besides, we additionally designed a location prediction task and a contrastive learning task to introduce localization and invariance characteristics for MAE. Our contrastive learning task not only enables the model to learn high-level visual information but also allows the training of MAE's class token. This is something that most MAE improvement efforts do not consider. Extensive experiments have shown that our method shows state-of-the-art performance on standard small datasets as well as medical datasets with few samples compared to the current popular masked image modeling (MIM) and vision transformers for small datasets.The code and models are available at https://github.com/Talented-Q/SDMAE.

研究の動機と目的

• 医療画像や科学的画像のような小規模データセットにおける Vision Transformers (ViTs) のデータ依存性の高い性質に対処すること。
• アーキテクチャの変更なしに、Masked Autoencoders (MAE) が ViT のデータ依存性を効果的に低減できるかどうかを調査すること。
• 小規模データセットにおける過学習を防ぐために、MAE の最適で軽量なデコーダー構成を同定すること。
• 自己教師付きタスクを通じて、局所化や不変性といったインダクティブバイアスを MAE に組み込むこと。
• MAE プレトレーニング中に、対照学習タスクを用いて ViT のクラストークンの効果的なプレトレーニングを可能にすること。

提案手法

• 標準的な ViT モデルの小規模データセット学習に特化した、修正された MAE フレームワークである Small Dataset Masked Autoencoders (SDMAE) を提案する。
• 小規模データセットにおける過学習を回避する最小限のデコーダー構成（深さ=1、d_model=128）を同定するために、広範なアブレーションスタディを実施する。
• 人為的アノテーションを一切不要とする、パラメータが極めて少ない軽量な位置予測ヘッドを導入する。
• 特徴の不変性を向上させるとともに、ViT のクラストークンの専用プレトレーニングを可能にする対照学習タスクを設計する。
• マスク画像モデリングと位置予測、対照学習を統合した統一されたプレトレーニング目的関数を組み合わせる。
• すべての実験において、公平な比較を保つために、一貫したプレトレーニングおよびファインチューニングプロトコルを用いる。

実験結果

リサーチクエスチョン

• RQ1 ViT のアーキテクチャを変更せずに、MAE を小規模データセットに効果的に適応できるか？
• RQ2 小規模データセットで学習する際、過学習を防ぐために MAE の最適なデコーダー構成は何か？
• RQ3 軽量な位置予測タスクは、小規模データ設定下で ViT の局所的インダクティブバイアスを回復するのにどの程度有効か？
• RQ4 対照学習タスクは、表現学習の向上と MAE ベースの ViT におけるクラストークンのプレトレーニングを可能にするか？
• RQ5 マスク画像モデリングに加えて自己教師付きタスクを組み合わせることで、小規模および医療画像データセットで最先端の性能が達成できるか？

主な発見

• SDMAE は Tiny-ImageNet でトップ-1精度 72.24% を達成し、標準的な MAE や既存の ViT・CNN ベースラインを上回る性能を示した。
• アブレーションスタディの結果、位置予測と対照学習の両方を組み合わせた場合が最高の性能を示し、位置予測を除外すると精度が 1.3% 減少した。
• 最適なデコーダー構成（深さ=1、d_model=128）は、過学習を顕著に低減し、小規模データセットにおける一般化性能を向上させた。
• SDMAE は医療画像タスク（APTOS 2019 など）でも最先端の性能を達成し、高度な CNN や他の MIM 手法を上回った。
• 医療データに対するアテンション可視化の結果、SDMAE は臨床的に関連する領域に注目する能力を学習していることが確認され、実用的価値を示した。
• 対照学習タスクにより、通常の MAE 手法で無視されがちなクラストークンの効果的なプレトレーニングが可能になった。