[論文レビュー] Self Pre-training with Masked Autoencoders for Medical Image Classification and Segmentation

Lei Zhou, Huidong Liu+4 more|arXiv (Cornell University)|Mar 10, 2022

Radiomics and Machine Learning in Medical Imaging被引用数 8

ひとこと要約

本論文は、外部データセット（例：ImageNet）に依存せずに、ターゲットデータセットのトレーニングデータのみを用いて、医療画像解析のための自己教師付きマスク付き自己符号化器（MAE）事前学習手法を提案する。3D画像パッチの12.5–75%をマスクし、ビジョントランスフォーマーを用いて再構築することで、文脈特徴の学習を向上させ、チークX線分類、腹部CTセグメンテーション、MRI脳腫瘍セグメンテーションの3つのタスクで最先端の性能を達成し、ランダム初期化やImageNet事前学習に比べて顕著な向上を示した。

ABSTRACT

Masked Autoencoder (MAE) has recently been shown to be effective in pre-training Vision Transformers (ViT) for natural image analysis. By reconstructing full images from partially masked inputs, a ViT encoder aggregates contextual information to infer masked image regions. We believe that this context aggregation ability is particularly essential to the medical image domain where each anatomical structure is functionally and mechanically connected to other structures and regions. Because there is no ImageNet-scale medical image dataset for pre-training, we investigate a self pre-training paradigm with MAE for medical image analysis tasks. Our method pre-trains a ViT on the training set of the target data instead of another dataset. Thus, self pre-training can benefit more scenarios where pre-training data is hard to acquire. Our experimental results show that MAE self pre-training markedly improves diverse medical image tasks including chest X-ray disease classification, abdominal CT multi-organ segmentation, and MRI brain tumor segmentation. Code is available at https://github.com/cvlab-stonybrook/SelfMedMAE

研究の動機と目的

• 医療画像分野において大規模かつドメイン特化された事前学習データの不足に応えるために、ターゲットデータセットのトレーニングデータのみを用いた自己教師付き事前学習パラダイムの提案。
• 解剖学的構造が機能的・機械的に相互に依存する医療画像において、マスク付き自己符号化が文脈特徴の学習をどのように向上させるかを調査すること。
• 低データ環境下における、分類およびセグメンテーションを含む多様な3D医療画像処理タスクに対するMAE事前学習の有効性を評価すること。
• 医療画像ベンチマークにおける下流タスクの性能を評価し、MAE自己事前学習とImageNet事前学習、ランダム初期化を比較すること。
• 医療画像タスクにおける最適なハイパーパrameter（例：マスク比率、事前学習エポック数）を特定すること。

提案手法

• MAE再構築中に干渉を避けるために、正弦・余弦関数を用いた位置エンコーディングを採用したビジョントランスフォーマー（ViT）バックボーンを用いる。
• マスク付き自己符号化器アーキテクチャを採用し、3D画像パッチの12.5%～75%をランダムにマスクし、可視パッチと学習可能なマスクトークンのみを用いて再構築する。
• エンコーダーは可視パッチのみを処理し、パッチごとの表現を出力する。軽量なデコーダーは、これらの表現とマスクトークンから元の画像を再構築する。
• マスクされたパッチに対して平均二乗誤差損失を用い、周囲のパッチからの文脈的情報を活用して欠損領域を推定するようネットワークを促進する。
• 事前学習後、エンコーダーの重みを下流モデルの初期化に転送する：分類タスクには線形分類器、セグメンテーションタスクにはUNETRに類似したデコーダーを用いる。
• ファインチューニングにはレイヤーごとの学習率スケーリングとコサインスケジューリングを適用し、各データセットに特化したハイパーパrameterを設定する。

Fig. 1 : Segmentation Pipeline with MAE Self Pre-training. Left: A ViT encoder is first pre-trained with MAE. A random subset of patches is input to the encoder and a transformer decoder reconstruct the full image. Right: The pre-trained ViT weights are transferred to initialize the segmentation enc

実験結果

リサーチクエスチョン

• RQ1 ターゲットデータセットのデータのみを用いた自己教師付きマスク付き自己符号化による事前学習が、医療画像分類およびセグメンテーションタスクの性能向上に寄与するか？
• RQ2 小規模な医療画像データセットに適用した場合、MAE事前学習がImageNet事前学習を上回る性能を発揮するか？
• RQ3 特に3DCTおよびMRIデータを対象とした医療画像セグメンテーションタスクにおいて、MAE事前学習の最適なマスク比率は何か？
• RQ4 自己事前学習は、腫瘍微小環境など解剖学的に接続された領域の文脈的依存性を捉える能力にどのように影響を与えるか？
• RQ5 位置エンコーディングの選択（正弦・余弦 vs. 学習可能なもの）が、医療画像MAE事前学習における再構築品質および下流タスクの性能に与える影響は何か？

主な発見

• BTCVデータセットにおける腹部多臓器セグメンテーションでは、MAE自己事前学習により平均DSCがベースラインの78.8%から83.5%に向上し、ImageNet事前学習を顕著に上回った。
• BraTS MRI脳腫瘍セグメンテーションタスクでは、MAE事前学習によりDSCが77.4%から78.91%に上昇し、HD95が7.78mmから7.22mmに低下した。
• CXR14胸部X線分類タスクでは、MAE自己事前学習がImageNet事前学習を0.8%上回り、自己教師付き事前学習を用いた最先端のCNNベースモデルをも凌駆した。
• 最高のセグメンテーション性能は12.5%のマスク比率で達成されたが、より高い比率（例：75%）では効果が薄く、タスクに応じた最適化が不可欠であることが示された。
• 事前学習エポック数を延ばすことで性能向上が見られたが、100～150エポックを超えると特に小規模データセットでは過学習が発生した。
• 正弦・余弦関数を用いた位置エンコーディングを事前学習段階で採用した場合、学習可能なエンコーディングに比べて再構築品質および下流タスクの性能が向上した。

Fig. 2 : MAE reconstruction. First row: Original image. Second row: Masked image where masked regions are colored with gray/black. Third row: Reconstructed images from unmasked patches. Every two columns show the results of one dataset, i.e., CXR, BTCV and BraTS from left to right.