[論文レビュー] STU-Net: Scalable and Transferable Medical Image Segmentation Models Empowered by Large-Scale Supervised Pre-training

Ziyan Huang, Haoyu Wang+9 more|arXiv (Cornell University)|Apr 13, 2023

COVID-19 diagnosis using AI被引用数 49

ひとこと要約

STU-Netは最大14億パラメータまでのスケーラブルなU-Net系を導入し、TotalSegmentatorで事前学習を行い、14の下流データセットおよびファインチューニングシナリオへの高い転移性能を示す。

ABSTRACT

Large-scale models pre-trained on large-scale datasets have profoundly advanced the development of deep learning. However, the state-of-the-art models for medical image segmentation are still small-scale, with their parameters only in the tens of millions. Further scaling them up to higher orders of magnitude is rarely explored. An overarching goal of exploring large-scale models is to train them on large-scale medical segmentation datasets for better transfer capacities. In this work, we design a series of Scalable and Transferable U-Net (STU-Net) models, with parameter sizes ranging from 14 million to 1.4 billion. Notably, the 1.4B STU-Net is the largest medical image segmentation model to date. Our STU-Net is based on nnU-Net framework due to its popularity and impressive performance. We first refine the default convolutional blocks in nnU-Net to make them scalable. Then, we empirically evaluate different scaling combinations of network depth and width, discovering that it is optimal to scale model depth and width together. We train our scalable STU-Net models on a large-scale TotalSegmentator dataset and find that increasing model size brings a stronger performance gain. This observation reveals that a large model is promising in medical image segmentation. Furthermore, we evaluate the transferability of our model on 14 downstream datasets for direct inference and 3 datasets for further fine-tuning, covering various modalities and segmentation targets. We observe good performance of our pre-trained model in both direct inference and fine-tuning. The code and pre-trained models are available at https://github.com/Ziyan-Huang/STU-Net.

研究の動機と目的

• 複数のモダリティとターゲットを扱える、スケーラブルで転移性の高い医用画像分割モデルの動機付け。
• nnU-Netを改良して、より高いスケーラビリティと転移性を持つSTU-Net系統を開発する。
• 下流タスクへの転移を高めるため、 大規模な医用分割データセットで事前学習する。
• 多様なデータセットとモダリティに対する直接推論とファインチューニングによる転移性を評価する。

提案手法

• nnU-Netブロックを残差接続で改良し、より深いアーキテクチャを可能にする。
• 転置型アップサンプリングを、重みなし補間と1x1x1畳み込みを組み合わせた方式に置換して転移性を高める。
• アーキテクチャのハイパーパラメータ（例: 段数、等方性カーネル）を固定して、タスク間の転移性を維持する。
• 深さと幅を複合的に同時にスケールさせ、パラメータ数を増やした STU-Net-S、STU-Net-B、STU-Net-L、STU-Net-H を作成する。
• TotalSegmentator CTデータセット（104臓器、1204体積）で4000エポック、ミラー拡張を用いて事前学習する。
• 下流データセットでファインチューニングまたは直接推論を実行し、必要に応じてチャネル適応を行う。

実験結果

リサーチクエスチョン

• RQ1 STU-Netは、大規模な医用分割データで深さと幅を共にスケールさせることで、スケーラブルな性能向上を達成できるか？
• RQ2 タスク固有のアップサンプリングを排除（重みなし補間を介して）転移性がモダリティやタスク全体で改善されるか？
• RQ3 TotalSegmentatorによる大規模監督付き事前学習が、さまざまな下流データセットでの転移性能にどう影響するか？
• RQ4 複数のCT/MR/PETデータセットにおける直接推論とファインチューニングの転移効果の違いは何か？

主な発見

• STU-Net-H（深さ3倍、幅3倍）は14億パラメータに達し、TotalSegmentatorクラス全体で最も高い平均DSCを示す。
• STU-Net-BはTotalSegmentatorの平均DSCでnnU-NetとSwinUNETR-Bを上回り、STU-Net-LおよびSTU-Net-Hへスケールするとさらに大きな利得を得る。
• 事前学習済みSTU-Netモデルは、直接推論で14の下流CTデータセットへ効果的に転送され、モデルが大きいほど平均DSCが高い傾向。
• 3つの下流データセット（AutoPETを含む）でSTU-Net-H-ftをファインチューニングすると最良の平均DSCを得られ、ベースラインを上回る。
• アーキテクチャの改良（残差ブロック、重みなしアップサンプリング）と複合スケーリングは、同程度の計算量下でnnU-Net系より常に優れる。