[論文レビュー] Self-Supervised Pre-Training of Swin Transformers for 3D Medical Image Analysis

Yucheng Tang, Dong Yang+6 more|arXiv (Cornell University)|Nov 29, 2021

Medical Imaging and Analysis被引用数 46

ひとこと要約

本論文は Swin UNETR を導入する。自己教師付きの proxy タスクで事前学習した 3D トランスフォーマー型エンコーダを用い、finetune 後に BTCV および MSD ベンチマークで最先端のセグメンテーションを達成する。

ABSTRACT

Vision Transformers (ViT)s have shown great performance in self-supervised learning of global and local representations that can be transferred to downstream applications. Inspired by these results, we introduce a novel self-supervised learning framework with tailored proxy tasks for medical image analysis. Specifically, we propose: (i) a new 3D transformer-based model, dubbed Swin UNEt TRansformers (Swin UNETR), with a hierarchical encoder for self-supervised pre-training; (ii) tailored proxy tasks for learning the underlying pattern of human anatomy. We demonstrate successful pre-training of the proposed model on 5,050 publicly available computed tomography (CT) images from various body organs. The effectiveness of our approach is validated by fine-tuning the pre-trained models on the Beyond the Cranial Vault (BTCV) Segmentation Challenge with 13 abdominal organs and segmentation tasks from the Medical Segmentation Decathlon (MSD) dataset. Our model is currently the state-of-the-art (i.e. ranked 1st) on the public test leaderboards of both MSD and BTCV datasets. Code: https://monai.io/research/swin-unetr

研究の動機と目的

• 医用画像解析に適した 3D トランスフォーマー型エンコーダを開発する（Swin UNETR）。
• 解剖学的文脈に合わせた自己教師付き proxy タスクを設計・統合する（インペインティング、回転、対照学習）。
• 大規模なラベルなし CT コーパスでの効果的な事前学習を実証し、セグメンテーションタスクへの転移を行う。
• 公開 BTCV および MSD ベンチマークで性能を検証し、従来の最先端と比較する。

提案手法

• 階層型 3D Swin Transformer エンコーダと skip 連結を持つ CNN ベースのデコーダを備えた Swin UNETR を提案する。
• エンコーダを三つの自己教師付き proxy タスクで事前学習する：マスクされたボリューム・インペインティング、3D 回転予測、対照学習。
• 重み付き複数目的損失 L_tot = λ1 L_inpaint + λ2 L_contrast + λ3 L_rot を等しい重みで用いる（λ1 = λ2 = λ3 = 1）。
• 頭頸部、胸部、腹部/pelvis 領域にまたがる ROI 対応表現を学習するため、事前学習中にランダムなサブボリュームのクロップおよび拡張を行う。
• BTCV の多臓器分割および MSD タスクに対して、CNNデコーダと skip connections を備えた 4 段階 Swin Transformer エンコーダを用いて事前学習済みエンコーダをファインチューニングする。

実験結果

リサーチクエスチョン

• RQ1 自己教師付きタスクで事前学習した 3D Swin Transformer エンコーダは、CT データに対して頑健で ROI 配慮のある表現を学習できるのか？
• RQ2 マスクされたボリューム・インペインティング、回転予測、対照学習が下流の 3D 医用画像分割を協調して改善するのか？
• RQ3 大規模なラベルなし CT データセットでの事前学習が、BTCV および MSD ベンチマークの性能、データ効率、収束にどのように影響するのか？

主な発見

• Swin UNETR は自己教師付き事前学習とともに BTCV の多臓器分割で最先端の Dice スコアを達成する。
• MSD では Swin UNETR がいくつかのタスクで最高の性能を、10 タスク中の総合 Dice で最良を達成する。
• アブレーションにより、すべての proxy タスクを組み合わせると Dice が最良となり（彼らの研究では BTCV で 84.72%）、インペインティングが単一タスクとして強力な利益を提供している。
• 事前学習はアノテーション作業を削減し、ラベル付きデータが少なくても高い性能を達成する（例：BTCV ラベルの 10% で Dice が約 10% 改善）。
• 事前学習データ量を増やし、すべての proxy タスクを使用することで収束を加速し、下流の精度を向上させる。