[論文レビュー] Delving into Masked Autoencoders for Multi-Label Thorax Disease Classification

Junfei Xiao, Yutong Bai+2 more|arXiv (Cornell University)|Oct 23, 2022

COVID-19 diagnosis using AI被引用数 6

ひとこと要約

本稿では、266,340枚のラベルなし胸部X線画像を用いて、マスクド自己符号化器（MAE）を用いた強力な事前学習および微調整のレシピを、ビジョントランスフォーマー（ViT）に提案する。この手法により、3つのベンチマークデータセットにおける多ラベル胸部疾患分類において、最先端の畳み込みニューラルネットワーク（CNN、例：DenseNet-121）と同等またはそれ以上の性能を達成し、それぞれNIH ChestX-ray14、CheXpert、COVIDxでmAUCが82.3%、89.2%、99.3%を記録した。

ABSTRACT

Vision Transformer (ViT) has become one of the most popular neural architectures due to its great scalability, computational efficiency, and compelling performance in many vision tasks. However, ViT has shown inferior performance to Convolutional Neural Network (CNN) on medical tasks due to its data-hungry nature and the lack of annotated medical data. In this paper, we pre-train ViTs on 266,340 chest X-rays using Masked Autoencoders (MAE) which reconstruct missing pixels from a small part of each image. For comparison, CNNs are also pre-trained on the same 266,340 X-rays using advanced self-supervised methods (e.g., MoCo v2). The results show that our pre-trained ViT performs comparably (sometimes better) to the state-of-the-art CNN (DenseNet-121) for multi-label thorax disease classification. This performance is attributed to the strong recipes extracted from our empirical studies for pre-training and fine-tuning ViT. The pre-training recipe signifies that medical reconstruction requires a much smaller proportion of an image (10% vs. 25%) and a more moderate random resized crop range (0.5~1.0 vs. 0.2~1.0) compared with natural imaging. Furthermore, we remark that in-domain transfer learning is preferred whenever possible. The fine-tuning recipe discloses that layer-wise LR decay, RandAug magnitude, and DropPath rate are significant factors to consider. We hope that this study can direct future research on the application of Transformers to a larger variety of medical imaging tasks.

研究の動機と目的

• ビジョントランスフォーマー（ViT）が、特に多ラベル胸部疾患分類において医療画像分野で果たせる可能性を評価すること。
• ViTのデータ依存性の高さを解消するため、マスクド自己符号化器（MAE）を用いて266,340枚のラベルなし胸部X線画像の大規模データセットで事前学習すること。
• 低コントラストな解剖学的構造や小さな病変特徴を示す医療画像の独自の特性に適合した、カスタマイズされた事前学習および微調整のレシピを開発すること。
• 同じデータと評価プロトコルを用いて、ViTと最先端のCNN（例：ImageNetで事前学習されたものや自己教師あり学習で事前学習されたもの）を公平に比較するベンチマークを提供すること。
• 再現可能性および医療ビジョン分野の今後の研究を支援するため、コードおよび事前学習済みモデル（ViT-Small、ViT-Base）を公開すること。

提案手法

• マスクド自己符号化器（MAE）を用いてビジョントランスフォーマーを事前学習し、画像パッチの90％をマスクし、可視パッチから欠損したピクセルを再構築する。
• 自然画像の事前学習とは異なり、標準的な0.2–1.0とは著しく狭い0.5–1.0のランダムリサイズクロップ範囲を採用した、変更を加えたデータ拡張戦略を用いる。
• 自然画像のMAEで用いられる75％とは異なり、医療画像の再構築に適した高いマスキング比（90％）を採用する。
• レイヤーごとの周辺学習率減少、マグニチュード15のRandAug拡張、およびDropPath率0.1を用いて、事前学習済みViTを微調整する。
• 同じデータと評価プロトコルを用いて、ViTの性能を、DenseNet-121などの最先端CNNと比較する。
• 解釈可能性および局在化性能を評価するため、ViTとCNNの注意マップをGrad-CAMを用いて分析・比較する。

実験結果

リサーチクエスチョン

• RQ1 マスクド自己符号化器（MAE）を用いて大規模な医療データセットで事前学習された、シンプルなビジョントランスフォーマー（ViT）は、多ラベル胸部疾患分類において、最先端のCNNと同等またはそれ以上の性能を発揮できるか？
• RQ2 医療画像、特に胸部X線画像に適用する際、MAEの最適な事前学習ハイパーパrameter（マスキング比、クロップ範囲、データスケール）は何か？
• RQ3 レイヤーごとの周辺学習率減少、RandAugのマグニチュード、およびDropPath率といった微調整ハイパーパrameterは、医療分類タスクにおけるViT性能にどのように影響を与えるか？
• RQ4 Grad-CAMとボクシングボックスIoUを用いて測定した場合、ViTの小さな病変領域の局在化能力は、CNNと比べてどうか？
• RQ5 医療データで事前学習するドメイン内転移学習は、ImageNetで事前学習するのと比べて、医療ビジョンタスクにおいてより効果的か？

主な発見

• 事前学習済みViTはNIH ChestX-ray14でmAUC 82.3％を達成し、最先端のDenseNet-121と同等の性能を示した。
• スタンフォードのCheXpertデータセットでは、ViTはmAUC 89.2％を達成し、文献に報告された最高のCNN結果を上回った。
• COVIDxデータセットでは、ViTはmAUC 99.3％を達成し、挑戦的で小規模かつ多様なデータセットにおいて優れた性能を示した。
• 小さな病変（例：結節）の局在化において、ViTはCNNよりも優れた性能を示したのに対し、CNNは比較的大きな異常（例：心臓肥大）の局在化に優れていた。
• 266,340枚の胸部X線画像でドメイン内事前学習を施したことで、ImageNet事前学習に比べてViTの性能が顕著に向上した。これは、ドメイン特化データの重要性を示している。
• レイヤーごとの周辺学習率減少、RandAugマグニチュード15、およびDropPath率0.1は、多ラベル分類タスクにおける効果的なViT微調整に不可欠なハイパーパrameterであると特定された。