A Segmentation Foundation Model for Diverse-type Tumors

Prior to initiating model training, it is imperative to address the inadequacy of available data. To tackle this challenge, we adopt a dataset integration method akin to the recent work MultiTalent [7]. This method seamlessly combines multiple datasets, effectively mitigating the issue of distinct label IDs for the same target across different datasets. The method employ the binary Binary Cross-entropy loss (BCE) and a modified Dice loss for each class over all images in a batch. And to compensate for the varying number of training images in each dataset, we choose a sampling probability per case that is inversely proportional to $\sqrt{n}$, where $n$ is the number of training cases in the corresponding source dataset. In our approach, we integrate ten datasets to construct dataset pool, comprising seven tumor datasets (Liver, Lung, HepaticVessel, Colon, Pancreas from MSD [17], Breast cancer [27], KiTS2021 [21]) and three multi-organ datasets (AMOS22 [16], BTCV [18], Totalsegmentor-part1 [15]). As detailed in Section 1, we create a dataset pool to address the scarcity of tumor data and enhance the tumor segmentation capabilities of TSFM. Directly fusing datasets will cause a challenge of inconsistent labels, where in the labels for the same type of target in different datasets vary. To overcome this, we establish a label remapping table to correlate the labels of the original datasets with the fused dataset pool. This ensures that each target has a unique label during training. Noteworthy is the innovative fusion of multi-organ datasets and tumor datasets, distinguishing our approach from MultiTalent.

Our network architecture preserves the U-shaped structure and skip connections, integrating a Resblock-backbone and Transformer-bottleneck. This design facilitates the extraction of finer-grained features compared to pure convolutional networks while addressing concerns regarding excessive memory consumption. Fig 1(a) illustrates the network structure, while detailed information about each block is provided in Fig 1(b).

Transfer learning in deep learning typically involves leveraging pre-trained models to initialize new models, thereby transferring learned features to related tasks. Our proposed TSFM demonstrates strong transferability due to the following design considerations: (1) The TSFM architecture combines CNN with transformer, enabling it to maintain learned weights from the CNN component even when dealing with downstream images of varying sizes. This ensures that most of the learned weights on CNN are preserved, with adjustments made only to the transformer-bottleneck part. (2) The concise network structure of TSFM reduces the risk of overfitting, enhances transferability, and yields improved transfer performance. (3) The dataset pool used for training TSFM encompasses diverse variations, multiple tumor types, and a sufficient number of datasets. This rich dataset pool effectively supports TSFM in extracting foundational features from medical images, making it suitable for various other medical segmentation tasks.

The environment of the experiments is all based on Python 3.9.18, CentOS 7, Pytorch 2.1.0+cu121 and nnU-Net v1.

We utilizes preprocessing, postprocessing, and similar operations to those in the nnU-Net v1 framework to process the data. The optimizer for the model is selected as SGD optimizer with Nestrov momentum of 0.99 and a weight decay of $1e^{-3}$. The batch size is fixed to 2 and each epoch contains 250 iterations, with a total of 1000 epochs, an initial learning rate of 0.001, and a patch size of 112x160x192. In the transfer experiment, all parameters will be set according to the default settings of nnU-Net.

The comparison in this section involves TSFM, nnU-Net, Swin UNTER, and SAM-Med3D [12, 13, 30], with separate Dice scores provided for each type of data. nnU-Net and Swin UNTER are trained using their default parameters from the original official code implementation on 10 datasets separately for 1000 epochs. The experimental results of SAM-Med3D are obtained using its pre-trained model with original parameters. Detailed data is presented in Table 1, which lists all tumor datasets. Due to space constraints, the results for multi-organ datasets will be provided in the supplementary materials. As depicted in Table 1, our model demonstrates strong performance on tumor datasets, surpassing nnU-Net in most cases, with significant improvements observed. These results underscore the effectiveness of our model, achieving an average improvement of 3$\%$ over nnU-Net in tumor segmentation and comparable performance in multi-organ segmentation. For specific results, please refer to Fig. 2.

To evaluate the transferability of TSFM, we deploy it in other downstream tasks, including BraTS2020, KiTS2019 [22] and Adbomen-1K [14], which are not included in the original dataset pool. We design experiments with 4 different training epochs (10, 20, 50, 100). Regarding the loading model, we use the best model stored in pre-training for downstream tasks. We conduct training both with and without loading the pre-trained model weights across these epochs. Additionally, we compare our approach with STU-Net [5], known for its strong transferability in medical image segmentation tasks, using its huge version for the same transfer process. Finally, we compare the results with nnU-Net trained for 1000 epochs.

Among the three datasets, BraTS2020 and KiTS2019 are tumor datasets comprising 368 and 210 cases, respectively. In the transfer learning experiments, our focus is on comparing the tumor segmentation performance. Abdomen-1K, on the other hand, is a multi-organ dataset consisting of 50 cases, and we compare the average segmentation scores across all organs. The detailed results are shown in Table 2. From the table, it can be seen that the pre-trained TSFM has excellent transfer learning ability. TSFM shows good performance after fine-tuning at 100 epochs, surpassing the nnU-Net results trained at 1000 epochs and also surpassing the STU-Net fine-tuning results at 100 epochs. After only training for 50 epochs, the performance is also close to the result of nnU-Net. These results indicate that TSFM can migrate to downstream datasets with only 10$\%$ of the trained epoch, and achieve good performance. With only 5$\%$ of the trained epoch, we can achieve a performance close to nnU-Net. The remarkable performance of our model in transfer learning can be attributed to the meticulous design of the model structure, the abundance of parameters, and the extensive dataset, underscoring the robust transferability of TSFM. Additionally, the enhancement of the network structure through the incorporation of a transformer component proves instrumental in optimizing tumor segmentation.

Since the primary contributions of our work lie in the model architecture and dataset pool construction, ablation experiments are conducted on these two aspects.

Model ablation To demonstrate the effectiveness of our model, we compare it with nnU-Net trained on the dataset pool (the same dataset used for TSFM’s training) for 1000 epochs. Comparing the results of nnU-Net and TSFM, it becomes evident that the TSFM network structure outperforms nnU-Net. Detailed performance metrics are presented in Table 1, illustrating the superior performance of our method when compared to nnU-Net trained on the dataset pool. Our model’s advantage lies in its ample number of parameters, enabling it to adeptly capture diverse characteristics within larger datasets, thereby yielding superior results. In contrast, nnU-Net, constrained by a relatively limited parameter count, struggles to comprehensively learn each feature within extensive datasets. This discrepancy significantly contributes to the superior performance of our model in ablation experiments.

Dataset pool ablation To validate the effectiveness of the dataset pool, we train and transfer our model on each dataset within the pool. We compare the results of TSFM trained on the dataset pool with those trained on individual datasets to demonstrate the promoting effect of the dataset pool on model performance. We select two representative datasets for our experiments: a tumor dataset (Pancreas) and a multi-organ dataset (AMOS22). Training our model on each dataset for 1000 epochs and then performing downstream transfer learning on the same dataset, we compare the results with those obtained in Section 3.3. As depicted in Table 3, our model exhibits significantly better transferability when trained with the dataset pool compared to training without it. The ablation experiment on the dataset underscores the importance of a sufficient dataset in the pre-training process of the cornerstone model. When the dataset is inadequate, the model’s performance may be compromised, highlighting the crucial role of a comprehensive dataset in achieving optimal results.