[論文レビュー] Swin MAE: Masked Autoencoders for Small Datasets

Zi'an Xu, Yin Dai+6 more|arXiv (Cornell University)|Dec 28, 2022

AI in cancer detection被引用数 11

ひとこと要約

Swin MAE は、小規模な医療画像データセットにおける自己教師あり表現学習のため、Swin Transformer をバックボーンとするマスクドオートエンコーダーを提案する。事前学習済み重みなしでも、ImageNet で微調整された自己教師あり Swin Transformer モデルと同等またはわずかに優れた転移学習性能を達成しており、マスク画像再構成を用いて数千年の画像からの強力な特徴学習が可能であることが示された。

ABSTRACT

The development of deep learning models in medical image analysis is majorly limited by the lack of large-sized and well-annotated datasets. Unsupervised learning does not require labels and is more suitable for solving medical image analysis problems. However, most of the current unsupervised learning methods need to be applied to large datasets. To make unsupervised learning applicable to small datasets, we proposed Swin MAE, which is a masked autoencoder with Swin Transformer as its backbone. Even on a dataset of only a few thousand medical images and without using any pre-trained models, Swin MAE is still able to learn useful semantic features purely from images. It can equal or even slightly outperform the supervised model obtained by Swin Transformer trained on ImageNet in terms of the transfer learning results of downstream tasks. The code is publicly available at https://github.com/Zian-Xu/Swin-MAE.

研究の動機と目的

• 深層学習における限られた、正確にアノテーションが付与された医療画像データセットの課題に対処する。
• ビジョントランスフォーマーが大規模な事前学習に依存するのを克服するため、Swin Transformer を用いてインダクティブバイアスを導入する。
• データオーグメンテーションがしばしば医療画像では非現実的であるため、データオーグメンテーションなしで、小規模なデータセットにおける効果的な自己教師あり表現学習を可能にする。
• 小規模な医療画像での自己教師あり事前学習のみを用いて、下流タスクで競争力のある転移学習性能を達成する。
• マスクドオートエンコーディングを用いた Swin MAE が、事前学習済み重みが一切ない状態でも、ImageNet で微調整された教師ありモデルと同等またはそれを上回る性能を示すことを実証する。

提案手法

• 事前学習中に 75% の画像パッチをランダムにマスクするマスクドオートエンコーダー（MAE）フレームワークを採用する。
• 局所的なインダクティブバイアスを導入するため、シフトされたウィンドウを用いた Swin Transformer をエンコーダーバックボーンとして使用し、小規模データセットにおける学習安定性を向上させる。
• 医療画像の特性に適応させるために、MAE の 16×16 パッチとは異なり、重複のない 4×4 パッチを用いたパッチベースの画像トークン化を実装する。
• 空間的整合性を保つために、4×4 パッチを単位としてマスクするウィンドウベースのマスキング戦略を採用する。
• 元の画像と再構成画像の間の平均二乗誤差（MSE）損失を用いて、マスク領域の予測に焦点を当てる。
• 軽量なデコーダーヘッドを用いて、マスクされた潜在表現から完全な画像を再構成する。MAE と Swin MAE の間でアーキテクチャを共有することで、公平な比較を実現する。

実験結果

リサーチクエスチョン

• RQ1 Swin Transformer をバックボーンとするマスクドオートエンコーダーは、事前学習済み重みが一切ない状態でも、小規模な医療画像データセットから有用なセマンティック特徴を学習できるか？
• RQ2 同じ小規模データセットで、Swin MAE は標準的な MAE と比べて収束速度と学習安定性に優れているか？
• RQ3 Swin MAE を用いた自己教師あり事前学習は、ImageNet で微調整された教師ありモデルと比較して、下流タスクの転移学習性能をどの程度向上できるか？
• RQ4 Swin Transformer のインダクティブバイアスは、標準的なビジョントランスフォーマーと比較して、小規模・低リソースの医療画像データセットにおける自己教師あり学習を改善するか？
• RQ5 データオーグメンテーションに依存せずに、マスクドオートエンコーディングを単一チャンネルで完全な医療スキャンに効果的に適用できるか？

主な発見

• Swin MAE は、数千年の医療画像での学習にもかかわらず、ImageNet で事前学習された教師あり Swin Transformer モデルと同等またはわずかに優れた転移学習性能を達成した。
• 同じ小規模データセットと同一のハイパーパrameterを用いて学習した場合、標準的な MAE と比較して、Swin MAE は収束が早く、学習損失が低く、滑らかな学習曲線を示した。
• Swin MAE は、Swin Transformer のインダクティブバイアスのおかげで、事前学習なしでも MAE よりも優れた性能を発揮した。これは、学習の安定性と特徴学習の向上をもたらした。
• 再構成結果から、Swin MAE は 75% のマスキング率であっても、境界の輪郭や色の整合性を含めた欠損領域の効果的な回復が可能であることが示された。
• 事前学習済み重みが存在しないにもかかわらず、性能に悪影響を及げなかった。Swin MAE は、小規模データでの自己教師あり再構成から、強力な表現を学習した。
• 上流の学習損失曲線から、Swin MAE は小規模データセットにおいて、MAE よりも優れた最適化安定性と高速な収束を達成しており、優れた自己教師あり学習能力を示していることがわかった。