MA-SAM: Modality-agnostic SAM Adaptation for 3D Medical Image Segmentation

Foundation models has recently been actively developed in computer vision, although to a lesser extent compared to their prevalence in natural language processing. Pioneering vision foundation models learn directly from vast image-text pairs sourced from the web in a self-supervised manner. Representative works CLIP [41] and ALIGN [24] leverage contrastive learning techniques to train both text and image encoders. However, these models primarily excel in tasks that involve mapping images to text, such as classification. Later on, Florence [54] incorporates universal visual-language representations, showing adaptability to more diverse computer vision tasks. One of the latest developments is SAM [29], a vision foundation model for general-purpose image segmentation. By pre-training on 1 billion masks, SAM demonstrates impressive zero-shot capability across numerous image segmentation tasks. Concurrently, SegGPT [51] and SEEM [62] have also emerged for general image segmentation, but are pre-trained on relatively smaller datasets compared to SAM.

With the remarkable performance exhibited by large models, the paradigm of pre-training large foundation models and subsequently fine-tuning for specific downstream tasks has gained increasing popularity. As the pre-trained large models continue to grow in scale, the research on PETL has emerged to achieve effective and efficient adaptation by optimizing only a small subset of model parameters while keeping substantial amount of parameters fixed. PETL techniques has been originally proposed in natural language processing, and can be categorized into three main groups [34], including additive methods, selective methods, and reparameterization-based methods. Additive methods, such as Adapters [17], aim to augment the existing pre-trained model by introducing additional parameters or layers, and then fine-tuning only these newly introduced components [15, 38]. Selective methods focus on updating a few selected influential layers or internal structure within the model [13, 56]. Reparameterization-based methods, such as LoRA [19] and FacT [26], leverage low-rank representations to minimize the number of trainable parameters, demonstrating robust and SOTA performance across various PETL tasks. Recently, PETL has also been actively studied in computer vision, enabling the effective adaptation of vision foundation models to a wide range of downstream tasks [60, 25, 39, 50].

Attracted by SAM’s outstanding zero-shot performance in natural images, a plethora of evaluation studies quickly emerged in various medical image segmentation tasks [21, 16, 47, 61, 10, 18, 7, 58]. However, due to the large domain gap between natural and medical images, directly applying SAM to medical applications typically resulted in unsatisfactory performance. For example, He et al. [16] assessed SAM’s segmentation accuracy on 12 medical imaging datasets and observed that SAM’s zero-shot performance lagged significantly behind models trained on domain-specific medical images, with performance gap as large as 70% in Dice in some tasks. Similar observations were reported in [21], even when using different types of prompts. These findings suggest the necessity of task-specific fine-tuning to adapt SAM for medical images for a better segmentation performance.

Subsequently, attention has shifted from evaluation to adaptation of SAM to medical images [57, 3, 53, 33, 12]. Driven by the improvements observed with the use of prompts, a majority of works leverage SAM’s prompt design during fine-tuning [7, 9, 8, 55]. For instance, SAM-Med2D [7] adopted more comprehensive prompts involving points, bounding boxes, and masks to tailor SAM for 2D medical images, and conducted comprehensive evaluations. MSA [52] employed point prompts and the Adapter technique to integrate medical domain knowledge into the SAM model. However, creating prompts for each 2D slice of 3D medical data is labor-intensive. In the case of SAM adaptation for fully automatic medical image segmentation [20, 40], SAMed [57] and Wang et al. [48] adopted LoRA for fine-tuning, showing superior performance than multiple 2D medical image segmentation methods. However, these methods do not take into account the critical 3D volumetric or temporal information, which is well-known to be valuable for enhancing medical image segmentation performance.

SAM is a promptable segmentation architecture consisting of three main components, i.e., the image encoder, the prompt encoder, and the mask decoder. The image encoder employs the Vision Transformer (ViT) [11] as the backbone, extracting essential features of the images with a set of transformer blocks. The prompt encoder takes in various types of prompts including points, boxes, or texts, and encodes these inputs into prompts embeddings to facilitate the segmentation task. The mask decoder is designed to be lightweight, which computes the cross-attention between embeddings of image and prompts, and utilizes transposed convolutional layers and multi-layer perception to generate segmentation masks. When applying to medical images, the model’s performance largely degrades since medical images present distinct texture and objects from natural images. This highlights the necessity for task-specific fine-tuning of SAM to address such challenges.

In order to effectively extract image features, SAM’s image encoder comprises a substantial portion of network parameters. Fine-tuning all these weights is computationally intensive. Previous research has shown that PETL techniques can achieve adaptation performance similar to full fine-tuning but with significantly fewer network parameters updated [19, 39]. In this regard, we adopt FacT [26], a SOTA PETL technique, that can obtain comparable or superior performance compared to other PETL methods while introducing a smaller number of trainable parameter.

Based on the common observations that the transformer-based models tend to be redundant in rank, FacT assumes the dense weight increment matrices $\Delta\bm{W}$ used for fine-tuning can be approximated by a set of low-rank factors with cross-layer weight sharing. Following the tensor decomposition in FacT, we decompose the weight increments $\Delta\bm{W}$ for each layer into three factors $\bm{U}\in\mathbb{R}^{d\times r}$, $\bm{V}\in\mathbb{R}^{d\times r}$, and $\bm{\Sigma}\in\mathbb{R}^{r\times r}$, where $d$ denotes the feature dimensions in ViT, $r$ stands for the rank of these factors with $r<<d$. It is worth noting that the two factors, $\bm{U}$ and $\bm{V}$, are shared across all layers, while the factor $\bm{\Sigma}$ is unique for each layer. The weight increments can then be calculated using the following equation:

where $s$ denotes a hyper-parameter for adjusting the learning rate of factors. We fix $s$ as 1 in our experiments and tune the overall learning rate with the optimizer to achieve a similar scaling effect. The FacT weight increments are applied to the query and value transformations within each transformer block, while all the other weights initialized from SAM remain frozen, as empirically there were no obvious improvements observed when applying FacT to other layers. With the FacT weight increments, the query and value transformations become:

where $\bm{W}_{q/v}$ denotes the query or value transformation after fine-tuning, $\bm{W}_{0}$ represents the SAM pre-trained weights.

SAM is initially pre-trained on 2D images, yet medical imaging typically involves more than two dimensions. For example, volumetric CT and MRI data contain crucial 3D spatial information for depicting anatomical structures or lesions, and surgical video data possesses valuable temporal relations between frames. Incorporating this volumetric or temporal knowledge inherent in medical imaging data, is pivotal for the successful transfer learning of SAM in medical applications. To address this key challenge, we propose to integrate a series of 3D adapters into the 2D transformer blocks within the SAM architecture. These adapters serve the purpose of extracting the essential volumetric or temporal insights needed for medical image analysis. By incorporating these adapters, we bridge the gap between the inherent complexities of medical imaging data and SAM’s pre-trained 2D backbone, enabling it to effectively handle multidimensional medical data.

Specifically, as shown in Fig. 1, each 3D adapter consists of a normalization layer, a linear down-projection layer, a 3D convolutional layer followed by an activation layer, and a linear up-projection layer. The core extraction of volumetric or temporal information primarily resides within the 3D convolutional layer. The purpose of the down-projection layer is to reduce the dimensionality of the original $d$-dimensional features into a more compact $c$-dimensional representation, so as to control the number of newly introduced parameters. Conversely, the up-projection layer restores the feature dimensions. With $\bm{M}$ denoting feature maps, the 3D adapter can be expressed as:

where Norm denotes the layer normalization, $\sigma$ denotes the activation function, $\bm{W}_{\mathrm{down}}\in\mathbb{R}^{d\times c}$ and $\bm{W}_{\mathrm{up}}\in\mathbb{R}^{c\times d}$ denote the linear down- and up-projection layer respectively, and Conv3D denotes the 3D convolutional layer with a kernel size of $3\times 1\times 1$ to specifically extract the third dimensional information.

To make the 3D adapters compatible with the 2D SAM backbone, for the network inputs, we extract a set of adjacent slices $\bm{x}=\{x_{i-\frac{N-1}{2}},...,x_{i},...,x_{i+\frac{N-1}{2}}\}_{i=1}^{B},\bm{x}\in\mathbb{R}^{B\times N\times H\times W}$. Here, $B$ denotes the batch size, $N$ denotes the number of adjacent slices, and $H\times W$ denotes the slice dimensions. Before the inputs are passed into the SAM backbone, a reshape operation is applied to transform $\bm{x}\in\mathbb{R}^{B\times N\times H\times W}$ into $\bm{x}\in\mathbb{R}^{BN\times H\times W}$ by merging the adjacent slices into the batch dimension. Then for the feature maps, prior to feeding into the 3D convolutional layer of a 3D adapter, they are reshaped from $[BN,H/16,W/16,c]$ to $[B,c,N,H/16,W/16]$. Here $H/16$ and $W/16$ denote the spatial dimensions of feature maps, which are down-sampled by 16 times because of the patch embedding process in transformer. After the 3D convolutional operation, the shape of feature maps are changed back again. In this way, the volumetric or temporal information can be effectively extracted within a 2D network backbone. For each transformer block, we incorporate two 3D adapters before and after the attention layers, as empirically superior performance can be obtained with such a design.

The mask decoder within original SAM comprises only two transformer layers, two transposed convolutional layers, and a single multilayer perception layer. Considering its lightweight architecture, it is feasible to apply full fine-tuning on the complete mask decoder for effective adaptation on medical images. During the patch embedding process of the transformer backbone within SAM’s image encoder, each $16\times 16$ patch is embedded as a feature vector, leading to $16\times 16$ times down-sampling of the inputs. The SAM mask decoder utilizes two consecutive transposed convolutional layers to up-sample the feature maps by 4 times, yet the final predictions generated by SAM remain 4 times lower in resolution than the original shapes. Nevertheless, since many anatomical structures or lesions in medical images are quite small, achieving a higher resolution is often necessary to ensure improved discrimination in the context of medical imaging [43].

To address this issue, we explore two approaches to tailor the mask decoder for enhanced suitability in medical image segmentation. For the first approach, termed as “progressive upsampling”, we introduce modest adjustments to the SAM decoder by integrating two additional transposed convolutional operations. With each layer up-samples the feature maps by a factor of 2, the four transposed convolutional layers progressively restore feature maps to their original input resolution. The second approach, termed as “multi-scale fusion”, entails creating a design resembling a “U-shaped” network [43]. This involves connecting the multi-scale feature maps of the image encoder with corresponding stages of the mask decoder using skip connections, a concept akin to that of U-Net. To achieve this, we uniformly divide the image encoder into four stages, establishing connections between the feature maps of each stage and those of the decoder through a series of up-sampling and convolutional operations. In our experiments, we have observed that the gradual up-sampling mechanism yields superior outcomes compared to multi-layer feature aggregation, showing the efficacy and simplicity of the progressive up-sampling approach.

Task1: The Beyond the Cranial Vault (BTCV) challenge dataset [30] contains 30 CT volumes with manual annotations for 13 abdominal organs. Each CT scan contains 85 to 198 slices with the slice thickness varying from 2.5 $mm$ to 5.0 $mm$. The axial size is $512\times 512$ for all scans, but with in-plane resolution ranging from 0.54 $\times$ 0.54 $mm^{2}$ to 0.98 $\times$ 0.98 $mm^{2}$. We use the same data split as [45], which contains 24 cases for training and 6 cases for testing.

Task2: We perform prostate segmentation on 6 MRI data sources [37], i.e., Site A to F, that were collected from NIC-ISBI13 [4], I2CVB [32], and PROMISE12 [35] datasets. The case number for each site is 30, 30, 19, 13, 12, 12 respectively, which were randomly divided into 80% and 20% for training and testing. These MRI scans from different sites were acquired with varying imaging protocols and present heterogeneous data distributions, thus were commonly used in previous domain generalization studies [36].

Task3: The 2018 MICCAI Robotic Scene Segmentation Challenge (EndoVis18) dataset [1] comprises 19 sequences, captured using the da Vinci X or Xi system. Each sequence contains either 149, 249, or 250 frames at a resolution of $1280\times 1024$. The dataset encompasses the surgical scene, with 12 classes annotated for various anatomical structures and robotic instruments. The dataset is officially split into 15 sequences for training and 4 sequences for testing.

Task4: The Pancreas Tumor Segmentation task within 2018 MICCAI Medical Segmentation Decathlon Challenge (MSD-Pancreas) dataset [2] contains 281 CT scans with annotations for pancreas and tumor. Each scan comprises 37 to 751 slices with an axial size of 512$\times$512. We follow [14] to utilize only tumor labels in our experiments and employ the same data split as in their work.

In addition, we use the Multi-Modality Abdominal Multi-Organ Segmentation Challenge (AMOS 22) dataset [23] for the evaluation of model generalization. This dataset contains abdominal CT and MRI data that were acquired from different patients. Each scan was annotated with 15 organs, but we focus on the 12 organs that overlap with the BTCV dataset. 300 CT scans and 60 MRI scans in the training and validation sets of AMOS 22 are used for our generalization evaluation.

For data pre-processing, the intensity values of each CT scan in BTCV and MSD-Pancreas datasets were truncated within the range of [-200, 250] Hounsfield Units (HU) and [-50, 200] HU respectively. The intensity of each MRI scan was truncated at the 99th percentile. Each CT or MRI scan was normalized to zero mean and unit variance. For surgical video data, each frame was normalized to [0, 1] range. We resized all images to $512\times 512$ for the axial plane of CT and MRI data, as well as for each frame of surgical video sequences. For model evaluation, we employ the common metrics, i.e., the Dice score and Hausdorff Distance (HD) to assess pixel-wise segmentation accuracy and the segmentation boundary quality respectively. We also report the mean intersection-over-union (mIoU) for the EndoVis18 dataset and the normalized surface distance (NSD) for the MSD-Pancreas dataset to align with previous studies.

The fine-tuning process was supervised using a hybrid segmentation loss, which combines the cross-entropy loss and Dice loss as: $\mathcal{L}_{\text{seg}}=\alpha\mathcal{L}_{\text{ce}}+\beta\mathcal{L}_{\text{Dice}}$. The weighting factors $\alpha$ and $\beta$ were set as 0.2 and 0.8 following [57], except for the surgical video data for which only the Dice loss was utilized. Every five consecutive slices were taken as the network inputs. For data augmentation, we applied a range of transformations including random rotation, flip, erasing, shearing, scaling, translation, posterization, contrast adjustment, brightness modification, and sharpness enhancement. Our model was trained using Adam optimizer with a batch size of 24. As in [57], we adopted a warmup training strategy to increase the learning rate linearly to the specific value and then exponentially decrease it towards the end of training to stabilize the training. We employed ViT_H as the backbone of the image encoder and conducted a total of 400 epochs of training, ensuring that the model converged effectively. Our framework was implemented in PyTorch 2.0 using 8 NVIDIA A100 GPUs.

For CT and MRI datasets, we extensively compare our method with various SOTA 3D medical image segmentation methods including CNN-based approaches nnU-Net [22], which is a U-Net [43] based self-configuring framework, showing robust performance on various medical image segmentation competitions, and 3D UX-Net [31], which is a very recent large kernel volumetric ConvNet for 3D medical image segmentation, as well as transformer-based methods SwinUNETR [46], which is a 3D transformer-based model with a hierarchical encoder, and nnFormer [59], which is a model combining local and global volume-based self-attention mechanism. We also compare our method with the most recent SAM adaptation methods SAMed_h [57], which is an automatic 2D medical image segmentation model for organ segmentation, and 3DSAM-adapter [14], which is a promptable 3D medical image segmentation model for tumor segmentation. For surgical video data, we compare our method with SOTA surgical scene segmentation methods, NCT [44], UNC [42], and OTH [6], which are the top-three approaches reported in the challenge, Noisy-LSTM [49] which uses ConvLSTM to learn temporal cues, STswinCL [27] which is a transformer-based model capturing intra- and inter-video relations, and nnU-Net. For all comparison experiments, the dataset splits remain consistent across all the methods.

Table 1 to Table 4.1 present comparative results for the four different tasks: abdominal multi-organ segmentation in CT data, prostate MRI segmentation across 6 sites, scene segmentation in surgical video, and tumor segmentation in CT data, respectively. When prompts are not specified, all methods generate results automatically without using any prompt. With our dedicatedly designed fine-tuning strategy for SAM, our method consistently and significantly outperforms other comparison approaches across all the four tasks. In terms of fully automatic segmentation for the first three tasks, our method improves the Dice score by 0.9%, 2.6%, 5% compared to the second-best performing approach, respectively. Notably, nnU-Net proves to be a strong competitor, showing robust segmentation performance across CT and MRI datasets. However, in surgical scene segmentation, nnU-Net obtains lower results compared to methods specifically tailored for processing surgical videos. Our method demonstrates strong performance across both volumetric and video medical data, indicating the potential of unifying the network architecture in these two domains of medical imaging, where previous methods were developed separately. When comparing with the pure 2D SAM fine-tuning method SAMed_h, which employs the same network backbone as ours, our method also achieves significantly better results, demonstrating the benefits of incorporating volumetric or temporal information for 3D medical image segmentation. The visual comparison results are presented in Fig. 2 to Fig. 4.

Pancreas tumor segmentation presents a substantial challenge due to the irregular contours and unclear margins of pancreas tumors in CT scans. As can be seen in Table 4.1 and Fig. 5, all automatic segmentation models struggle to correctly delineate pancreas tumor regions, obtaining merely a 41.6% Dice score for the best-performing model. We consider in such a demanding segmentation task, the use of prompts become valuable. By adding prompts in the form of one tight 3D bounding box per volume into the model, our method remarkably boosts the Dice score from 41.6% to 80.35%, demonstrating the effectiveness of leveraging prompts for tumor segmentation. However, if allowing 5% relaxation on the tightness of provided bounding box, the performance drops to 74.7%, showing the effect of prompts quality on segmentation performance. Our method also significantly outperforms the recent holistic 3D SAM adaptation method 3DSAM-adapter, with 10% Dice improvement when using automatic segmentation. This can be attributed to our method’s effective incorporation of third-dimensional information into model fine-tuning, as well as its substantial utilization of pre-trained weights from SAM to retain general discriminative knowledge. We notice that our method’s automatic segmentation performance falls slightly behind nnU-Net on tumor segmentation. This observation might indicate that SAM fine-tuning might be less effective for objects with ill-defined margins and small sizes, as these characteristics differ from the natural images on which SAM was originally trained.

One of the most appealing advantage of foundation models lies in their impressive generalization capability. To investigate the generalization of our models adapted from SAM, we first compare the zero-shot and few-shot capability of nnU-Net and our method by applying models trained on the BTCV CT scans to the AMOS22 CT and MRI scans. In Fig. 6, “nnU-Net 0 shot” and MA-SAM 0 shot” denote that the models trained on BTCV data are directly used to perform inference on AMOS22 images, and “nnU-Net 5 shot” and MA-SAM 5 shot” denote the models are further fine-tuned with 5 additional training cases from the AMOS22 dataset. From the results, we can see that our method exhibits better zero-shot and few-shot segmentation performance on AMOS22 CT and MRI images, demonstrating higher generalization capability. Especially for MRI images, nnU-Net encounters complete failure in the zero-shot context, obtaining only 10.9% Dice score, while our model still retains the performance of 60.4% Dice score. In the five-shot context, our method also shows 9% improvements than nnU-Net, further underscoring the advantages of generalization.

We also compare the model generalization on prostate MRI segmentation. In Table 5, the results of nnU-Net and our models are obtained by directly applying the models fine-tuned on Site A to make predictions for Site B to F. We include two recent SOTA test-time domain generalization methods into comparison, i.e., TTST [28] and TASD [36], which employ additional domain generalization techniques when performing predictions on each specific site. The results demonstrate that our method not only outperforms nnU-Net by a large margin when generalizing to different sites, but also achieves superior performance than SOTA domain generalization approaches. All these results on AMOS22 dataset and different prostate MRI datasets underscore the impressive generalization capability of our method, which is an importance characteristic for critical medical applications.

We conduct extensive ablation experiments on the BTCV dataset to investigate several key aspects regarding our SAM fine-tuning strategy: 1) effectiveness of each important component in our method, 2) effect of mask decoder design, 3) influence of network backbone, 4) choice of location for 3D adapters, 5) choice of rank for parameter-efficient fine-tuning.

1) Effectiveness of each component: We first validate the contribution of key components within our method, i.e., SAM’s pre-trained weights, the parameter-efficient fine-tuning strategy with FacT, and the incorporation of 3D information with 3D adapters. In Table 9, the “Full FT” model denotes whether full fine-tuning is used for the image encoder, since the mask decoder is fully fine-tuned for all the models. By comparing the results between the first and third rows in Table 9, we observe a substantial 13.1% improvement in Dice when the model is initialized with SAM’s pre-trained weights. This underscores the benefits of utilizing SAM’s original weights that were pre-trained on a large-scale and diverse dataset. The second row shows that if the entire image encoder remains frozen without being fine-tuned, and only the mask decoder is updated, the performance is unsatisfactory. This suggests that due to the significant difference in image texture between natural and medical images, the image encoder trained solely on natural images struggles to extract essential features from medical images. Moreover, Table 9 demonstrates that FacT is capable of delivering performance on par with full fine-tuning, by adjusting a small portion of weight increments. The models equipped with 3D adapters achieve superior performance, validating the importance of incorporating the third-dimensional information for medical image segmentation.

2) Effect of mask decoder design: We compare the performance of different mask decoder designs, including the original SAM mask decoder, the progressive up-sampling strategy, and the multi-scale fusion strategy. Table 4.4 shows that the straightforward progressive up-sampling strategy yields superior results, validating its simplicity and effectiveness. These results demonstrate the importance of recovering prediction resolution for medical images, which often contain small objects. However, no significant improvements were observed with the multi-scale fusion strategy. This might because of the extensive modifications it introduces to the original SAM decoder, resulting in less effective utilization of the pre-trained SAM weights.

3) Influence of network backbone: We conduct experiments with different network backbones, i.e., ViT_B, ViT_L, and ViT_H, to assess their impact on the performance of our method. As can be observed in Table 4.4, there is a noticeable improvement in Dice performance as the model size increases from ViT_B to ViT_H, signifying the advantages of using a larger model size to enhance overall model performance.

4) Choice of location for 3D adapters: We perform ablation experiments to investigate the placement of 3D adapters within our model. Specifically, we compare the performance when incorporating a 3D adapter in one of three locations: before the multi-head self-attention block (MHSA), after MHSA, or in both of these positions. As demonstrated in Table 4.4, the configuration with two 3D adapters positioned both before and after MHSA yields superior performance for our final model.

5) Choice of rank: We investigate how the model’s performance changes with varying decomposition rank $r$, by considering the rank value from the set {4, 8, 16, 32, 64}. As expected, Table 10 shows that with an increase in rank, there is a corresponding improvement in average Dice performance, but the performance tends to saturate when $r\geq 32$. We thus set $r=32$ in our experiments to seek a balance between performance gains and the number of parameters introduced.