STU-Net: Scalable and Transferable Medical Image Segmentation Models Empowered by Large-Scale Supervised Pre-training

要約

タイトル： STU-Net：大規模な教師付き事前学習によって強化されたスケーラブルで転移可能な医療画像セグメンテーションモデル

要約：
– 大規模データセットで事前学習を行った大規模モデルは、ディープラーニングの開発を深く進めることを実証している。
– しかしながら、医療画像セグメンテーションの最先端のモデルは依然として小規模であり、パラメータは数千万に留まっている。
– これらをより高次数にスケーリングしていくことはめったに探求されていない。
– 大規模モデルを探求する総合目標の1つは、大規模な医療画像セグメンテーションデータセットでトレーニングし、より良い転移能力を持たせることである。
– 本研究では、14百万から14億のパラメータサイズを持つ一連のスケーラブルで転移可能なU-Net（STU-Net）モデルを設計した。特に、14億のSTU-Netはこれまでに医療画像セグメンテーションで最大のモデルである。
– STU-NetはnnU-Netフレームワークに基づいており、その人気と優れたパフォーマンスのために採用された。
– 最初に、nnU-Netのデフォルトの畳み込みブロックをスケーラブルに調整した。その後、ネットワークの深さと幅の異なるスケーリングの組み合わせを実証的に評価し、モデルの深さと幅を一緒にスケーリングすることが最適であることがわかった。
– STU-Netモデルを大規模なTotalSegmentatorデータセットでトレーニングし、モデルサイズを増やすことで性能が向上することを発見した。大型モデルが医療画像セグメンテーションに有望であることが示唆された。
– さらに、異なるモダリティとセグメンテーションターゲットをカバーする、直接推論用の14のダウンストリームデータセットと、さらなるファインチューニング用の3つのデータセットで、モデルの転移性能を評価した。事前にトレーニングされたモデルの直接推論とファインチューニングの両方で良好なパフォーマンスを観察した。
– コードと事前学習モデルは、https://github.com/Ziyan-Huang/STU-Netで入手可能である。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, OpenAI