PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ziyan Huang, Haoyu Wang, Zhongying Deng, Jin Ye, Yanzhou Su, Hui Sun, Junjun He, Yun Gu, Lixu Gu, Shaoting Zhang, Yu Qiao

arXiv:2304.06716v1cs.CV

TL;DR

Medical image segmentation lacks large, transferable models that can handle varied tasks and modalities. STU-Net scales an nnU-Net-derived architecture, pre-trains it on TotalSegmentator, and evaluates transfer across downstream datasets, finding robust transfer and improved performance with larger models.

  • Problem

    Existing medical segmentation models are generally small, require task-specific tuning, and lack evidence for one model handling diverse modalities and segmentation targets.

  • Method

    STU-Net refines nnU-Net for scalability and transferability, scales depth and width together, and uses supervised pre-training on TotalSegmentator.

  • Results

    Scaling model size yields significant performance improvements across downstream tasks, while STU-Net-H shows robust transfer in both direct inference and fine-tuning.

  • Takeaways & Limitations

    Large-scale pre-trained models show practical potential for transferable medical image segmentation across multiple downstream datasets.

Abstract

from arXiv · show

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.

1. Introduction

STU-Net addresses limited transferability and scale in medical image segmentation by building scalable models on nnU-Net and pre-training them on TotalSegmentator. The resulting models span 14 million to 1.4 billion parameters and transfer across diverse downstream tasks.

  • Medical segmentation supports downstream tasks including registration, quantification, and image-guided surgery.
  • Existing models often require careful task-specific tuning and rarely handle varied modalities and segmentation targets with one model.
  • STU-Net scales from 14 million to 1.4 billion parameters, with the 1.4B model identified as the largest medical image segmentation model to date.
  • Scaling depth and width together was found optimal, producing STU-Net-L with 450 million and STU-Net-H with 1.4 billion parameters.
  • Replacing transpose convolution and fixing weight-related architectural hyper-parameters makes STU-Net weights transferable between pre-training and fine-tuning tasks.
  • Pre-training on TotalSegmentator yields strong transfer through direct inference and fine-tuning across diverse downstream datasets.

2. Related Works

Prior medical segmentation models are generally small and task-specific, while STU-Net scales them by an order of magnitude and evaluates transfer on large-scale datasets.

  • Medical segmentation models include CNN- and transformer-based architectures, but commonly contain only several million parameters and lack scalability and transferability.
  • Model scaling commonly varies depth and width, while prior work studies scaling relationships among computation, data size, and model size.
  • STU-Net scales models an order of magnitude beyond previous work and evaluates their transfer capacities on large-scale datasets.

3. Methods

STU-Net refines nnU-Net with residual and transfer-friendly components, then scales depth and width symmetrically and pre-trains across a family of model sizes.

  • Architectural refinements: STU-Net builds on nnU-Net’s task-adaptive encoder-decoder while refining its architecture for scalability and transferability.
  • Architectural refinements: Residual connections address gradient diffusion when increasing the number of convolutional blocks, enabling deeper stages.
  • Architectural refinements: Fixing six resolution stages and isotropic 3 × 3 × 3 kernels keeps weight-related architecture consistent across tasks.
  • Architectural refinements: Integrating downsampling into the first residual block uses parallel branches to reduce resolution while matching output shapes.
  • Architectural refinements: Weight-free interpolation followed by a 1 × 1 × 1 convolution replaces transpose convolution to avoid task-dependent weight-shape mismatch.
  • Scaling strategy: Depth and width are scaled symmetrically across encoder and decoder stages, with S, B, L, and H denoting Small, Base, Large, and Huge models.
  • Pre-training and transfer: Pre-training uses a 105-channel output for TotalSegmentator and extends training to 4000 epochs with mirror augmentation.
  • Pre-training and transfer: Direct inference reuses pre-trained models for downstream CT labels within the upstream classes, while novel labels or modalities require a reinitialized segmentation head.

4. Experiments

The experiments train STU-Net models on TotalSegmentator and compare their segmentation performance across model scales and established baselines. Larger STU-Net variants achieve progressively stronger results on the validation dataset.

  • Experimental setup: 1204 images with 104 anatomical structures form the TotalSegmentator training dataset, split into 1081 training, 57 validation, and 65 testing cases.Images are resampled to 1.5 × 1.5 × 1.5 mm isotropic resolution.
  • Experimental setup: 1000 epochs are used to compare STU-Net with other methods within the nnU-Net framework.The comparison uses reported optimizer, learning-rate, and decay settings for competing methods.
  • Quantitative results: 0.36% and 4.48% mean DSC gains allow STU-Net-B to surpass nnU-Net and SwinUNETR-B, respectively, across all TotalSegmentator classes.Mean DSC (%) is reported for all classes and five subclass groups.
  • Quantitative results: 1.59% and 2.94% improvements in mean DSC result from scaling the STU-Net base model to large and huge sizes, respectively.STU-Net-H achieves the highest mean DSC across all classes and the five subclass groups.

4.2. Transferability of Trained Models

Transferability is evaluated through direct inference on 14 downstream datasets and fine-tuning on three datasets spanning new structures, modalities, and domains. Larger pre-trained STU-Net models generally perform better, including on non-CT modalities and lesion segmentation.

  • Direct inference: 14 downstream datasets containing 2494 cases are used for direct inference on categories included in TotalSegmentator.Annotation inconsistencies are addressed by merging lesion labels with corresponding organs and reconciling left-right labels.
  • Direct inference: Larger pre-trained models usually achieve higher mean DSC across all 14 direct-inference datasets and generally on each specific dataset.The results support stronger transfer capacities for large-scale STU-Net models pre-trained on TotalSegmentator.
  • Fine-tuning: 80.69% mean DSC is achieved by STU-Net-H-ft across the three fine-tuning datasets, outperforming models trained from scratch.The datasets include AutoPET22 with lesions and PET, AMOS22 with MR, and FLARE22 with multiple domains.
  • Fine-tuning: Limited downstream training samples make the huge model slightly worse than the large model from scratch, whereas pre-training enables the huge model to exceed the large model after fine-tuning.The huge model usually surpasses the large model by a clear margin when TotalSegmentator pre-training supplies sufficient samples.
  • Fine-tuning: Pre-trained models improve performance after fine-tuning on non-CT modalities, including AMOS-MR and AutoPET-PET, despite CT-focused pre-training.The paper suggests that pre-training learns features and structures that generalize across modalities.

4.3. Ablation Study

The ablations show that STU-Net’s architectural refinements improve performance and transferability, while compound scaling and larger pre-training setups strengthen segmentation results. Larger models also remain advantageous with limited training data.

  • 4.3.1 Efficacy of Architectural Refinements: STU-Net-B’s refinements outperform parameter-matched nnU-Net*, demonstrating the effectiveness of its architectural changes.nnU-Net* increases the maximum feature number from 320 to 512 to match STU-Net-B’s parameters.
  • 4.3.1 Efficacy of Architectural Refinements: The default upsampling design improves performance and running speed while preserving weight transferability for downstream fine-tuning.Replacing it with transpose convolution decreases performance by 0.16% and makes weights non-transferable; trilinear interpolation also reduces performance and slows execution.
  • 4.3.2 Scaling Strategy: Compound scaling of depth and width is more effective and efficient for STU-Net than scaling either dimension alone.Across comparable scaling settings, STU-Net consistently outperforms nnU-Net* with similar parameters and FLOPs.
  • 4.3.3 Pre-training Settings: 0.84% and 3.67%: large-scale pre-training improves fine-tuned performance on FLARE22 and AutoPET, respectively.The comparison uses STU-Net-L trained with and without TotalSegmentator pre-training.
  • 4.3.3 Pre-training Settings: Mirror augmentation in both pre-training and fine-tuning achieves the best downstream result, while 4k pre-training epochs produce the best result among tested durations.The 4k-epoch setting is described as sufficient for convergence.
  • 4.3.4 Model Scaling with Dataset Size: Larger STU-Net models outperform smaller models even with fewer training cases, indicating greater data efficiency, while performance continues improving as case counts increase.STU-Net-H surpasses STU-Net-S with 5% of cases and STU-Net-B with 20% of cases; the case-count trend has not saturated.

4.4. Universal Model Versus Expert Models

The universal model and expert models trade advantages across anatomical subcategories, but the largest universal model achieves the best overall TotalSegmentator performance. Increasing model size generally benefits both model types.

  • Model-size comparison: As model size increases, performance generally improves for both universal and expert STU-Net models.The comparison covers five expert models targeting anatomical subcategories and one universal model trained on all 104 classes.
  • Subcategory comparison: Expert models excel on organs, vertebrae, and cardiac categories, whereas universal models perform better on muscles and ribs.The relative advantage varies by anatomical subcategory rather than favoring one model type everywhere.
  • Overall performance: 90.06%: STU-Net-H’s universal model achieves the highest overall mean DSC, compared with 89.07% for the best expert model.The comparison is reported across all TotalSegmentator classes.
  • Scaling effects by category: With increasing model size, the universal–expert performance gap narrows for organs and vertebrae, reverses for cardiac, and widens for muscles and ribs.The direction and magnitude of the gap therefore depend on the anatomical subcategory.

5. Outlook

The authors position STU-Net as an initial step toward medical foundation models and MedAGI. They connect this direction to large-scale datasets and models that generalize across medical imaging tasks.

  • MedAGI motivation: The authors frame STU-Net as an attempt toward Medical Artificial General Intelligence, or MedAGI.They argue that MedAGI should rely on foundation models and large-scale datasets, as in computer vision.
  • Foundation-model gap: Medical image processing has emerging large-scale datasets but lacks foundation models with broad generalization capabilities.The authors contrast this gap with large-scale foundation models in computer vision.
  • Future directions: The authors propose extending STU-Net toward models that segment everything in medical images and then generalize beyond segmentation to tasks such as detection and classification.These directions are presented as possible future developments rather than demonstrated capabilities of the current work.

6. Conclusions

The paper introduces scalable, transferable STU-Net models based on nnU-Net and evaluates them through large-scale TotalSegmentator training. Larger models show stronger transfer performance, supporting their practical potential for medical image segmentation.

  • Contributions: STU-Net provides scalable medical image segmentation models based on nnU-Net, with the largest model containing 1.4 billion parameters.The 1.4B model is described as the largest medical image segmentation model to date.
  • Transfer performance: Training on TotalSegmentator shows that increasing model size yields significant performance improvements when transferring to varied downstream tasks.The conclusion summarizes transfer across multiple downstream settings without specifying a single representative metric.
  • Practical value: STU-Net-H demonstrates robust transferability in both direct inference and fine-tuning across multiple downstream datasets.The authors identify large-scale pre-trained models as practically valuable for medical image segmentation.
  • Conclusion: The authors conclude that scalable and transferable STU-Net models could advance medical image segmentation research and innovation.This is stated as a potential direction rather than a guaranteed outcome.

A. Consistency of overlapped annotations

The evaluation harmonizes overlapping anatomical targets across datasets by merging related labels and, where needed, inference results. These adjustments address annotation-protocol variation and hierarchical organ–tumor relationships.

  • Kidney and Lung targets were reconstructed by combining related TotalSegmentator annotations to maintain consistency across protocols.Kidney combines left and right kidney targets; Lung combines five TotalSegmentator targets, although the passage truncates the list.
  • Organ ground truths were combined with their inner tumors for MSD Liver, MSD Pancreas, and KiTS2021.The resulting composite targets were Liver, Pancreas, and Kidney, respectively.

B. Dataset Details

The evaluation spans 15 publicly available datasets, using 14 for direct inference and 3 for fine-tuning across CT, MR, PET, and varied anatomical targets. Dataset sizes range from small specialized cohorts to large multi-center collections.

  • 15 publicly available datasets were evaluated: 14 for direct inference and 3 for fine-tuning.Direct inference used all annotated cases and overlapping targets with TotalSegmentator where applicable.
  • Dataset scale varies from 30-case cohorts such as BTCV and BTCV-Cervix to 1,000 abdominal CT cases in AbdomenCT-1K.The datasets also include KiTS2021 with 300 annotated CT scans and CT-ORG with 140 CT images.
  • Specialized CT datasets include Verse2020 spine scans, WORD multi-organ scans, and SegThor thoracic-organ scans.Verse2020 evaluation uses its 61-case official train set; WORD evaluates 120 annotated cases, while SegThor provides 40 cases.
  • Fine-tuning evaluations use official or specified splits, including FLARE22, AMOS, and AutoPET training and test partitions.FLARE22 uses 50 labeled cases for training and 20 validation cases; AMOS and AutoPET use their stated training and evaluation subsets.
  • The benchmark includes CT, MR, and paired CT/PET data covering organs, tumors, lesions, and other anatomical structures.AMOS includes CT and MR, while AutoPET provides paired 3D CT and PET studies with tumor-lesion annotations.
Loading 2304.06716v1…