When converting SAM to auto-mode (without prompts), we can simply use the default embeddings generated by using “None” as inputs to the prompt encoder, i.e.

 sparse_embeddings, dense_embeddings = sam.prompt_encoder(points=None, boxes=None, masks=None), 

where the “sparse_embeddings” and “dense_embeddings” serve as inputs to SAM’s mask decoder. Empirical evidence shows the effectiveness of this simple approach [96]. More advanced methods, including adding a detection network to automatedally generate box prompts or a self-prompt module into SAM’s structure [87, 27, 72], are feasible but will introduce extra complexity into SAM’s structure, hampering our analysis on the effectiveness of fine-tuning strategies. Moreover, we believe our work offers a potentially effective strategy, i.e. feeding the generated prompts to the fine-tuned SAM instead of the original one, for this research direction.

In this work, we select ViT-B and ViT-H from SAM’s provided checkpoints because they represent the most efficient and effective versions of SAM’s encoder respectively. Additionally, we explore a more compact option: ViT-T(iny). This backbone is introduced in MobileSAM [94], where the authors distill the knowledge from SAM’s pre-trained ViT-H to ViT-T, a customized vision transformer with $1\%$ of the network parameters and offers an inference speed $50$ times faster compared to the original SAM. The lightweight nature can widen the application of SAM, ease the computational requirements, and, more importantly, raise interesting research questions, such as whether distillation can maintain or even lead to the superior performance of SAM during fine-tuning.

When selecting which set(s) of SAM’s parameters to update, the most straightforward way is to select one or multiple components from SAM. As discussed in Section 2.1.1, we use the default prompt embeddings and thus can keep the prompt encoder as is.

Some methods state the effectiveness of maintaining the encoder’s weights and only adapting the lightweight mask decoder of SAM’s architecture, as shown in Table 1. These approaches are grounded in the understanding that SAM’s encoder, trained on a comprehensive collection of natural images, is adequately powerful and knowledgeable in the general imaging domain with zero-shot ability to perform feature extraction, and updating the lightweight decoder can enhance performance in specific applications. Also, it is suggested that fine-tuning the encoder of SAM would demand considerably higher computational resources. However, due to the significant appeal difference between natural images and medical images, it remains unclear if SAM’s encoder could effectively capture the distinctive features of medical images without any modifications to the image encoder. Therefore, we consider these two choices: mask decoder only and mask decoder plus image encoder.

When conducting model fine-tuning or transfer learning, the easiest and most intuitive way is to update all model parameters involved when feeding medical domains of data, which is vanilla fine-tuning [6].

Recently, parameter-efficient fine-tuning (PEFT) techniques have gained an increasing amount of attention. Compared with vanilla fine-tuning, PEFT keeps most of the network parameters fixed and only updates a small percentage, i.e. less than $5\%$, of them. Such a strategy significantly reduces the computational cost and is beneficial in preventing catastrophic forgetting as well as over-fitting [23]. The simplest PEFT techniques involve selecting and fine-tuning a subset of the parameters from the original pre-trained model. These selective methods concentrate on identifying the optimal parameters for fine-tuning, aiming to strike an ideal balance between parameter efficiency and overall model performance. For instance, BitFit [93] proposed fine-tuning only the bias term in a transformer-based network, and Tinytl [12] shares a similar idea but focuses on CNN-based networks. Touvron et al. [79] have demonstrated the effectiveness of selectively fine-tuning only the attention layers in Vision Transformers (ViTs), whereas alternative approaches suggest that exclusively optimizing the normalization layers [7].

Rather than updating certain selected layers, recent research [54] indicates that incorporating extra modules into the large models appears to be more beneficial because it re-configures models towards a downstream task than solely updating network parameters. Among these additional module methods, Adapters, which were originally introduced for learning multi-domain representations [68, 69], have emerged as a particularly prevalent choice. In addition to the conventional Adapter block [35], Low-Rank Adaptation (LoRA), known for its efficient low-rank matrices added into self-attention layers that require even fewer parameters, has also gained popularity [36]. Additionally, AdaptFormer [15], which integrates a scalable Adapter into the multi-layer perceptron (MLP) layer of the attention block specifically for vision transformers, has shown that adding just 2% extra parameters can surpass the performance of fully fine-tuned networks across various natural imaging benchmarks. In our study, we aim to investigate the effectiveness of applying PEFT for fine-tuning SAM within the medical imaging domain. Considering Adpater and LoRA are the most popular PEFT methods adopted in LLMs and vision foundation transformers [91], we choose to include them as well as the simplest Vanilla for our experiments. In summary, we select the following methods:

1. 1.

   Vanilla - a “baseline” approach that updates all parameters.

2. 2.

   Adapter / LoRA - these approaches involve adding extra layers with original parameters fixed.

To sum up, we choose the following combinations to fine-tune SAM’s performance for specific datasets or tasks, as illustrated in Figure 2: (1) Variations in the image encoder’s architecture, including ViT-H, ViT-B, and ViT-T from MobileSAM [94]. (2) Different approaches to the fine-tuning scope of the network, encompassing both the encoder and decoder updates, as well as the decoder only. (3) Diverse parameter updating settings, comprising Vanilla, PEFT with Adapter blocks, and PEFT with LoRA blocks. By integrating these three variables, we present a total of $3\times 2\times 3=\textbf{18}$ distinct architectures for fine-tuning SAM on task-specific supervised training setup.

When conducting task-specific supervised training, in addition to finding the most effective fine-tuning strategy, we also want to compare the performance against conventional UNet-based algorithms. Therefore, we have chosen three well-known and frequently used automated segmentation models for comparison: UNet [70], nnUNet [40], and SwinUNet [13]:

1. 1.

   UNet serves as the baseline of the convolutional neural network-based segmentation algorithm.

2. 2.

   nnUNet configures a UNet-based algorithm according to the specific attributes of the input datasets and approximates the performance of UNet with the optimal design choice.

3. 3.

   SwinUNet replaces the encoder of UNet with a swin-transformer [58].

For PEFT with Adapter blocks, we implemented it based on [87], which added adapter blocks to SAM’s attention blocks. In the image encoder, adapter blocks were added into transformer blocks right after the multi-attention head or within the MLP residual block with a scalable parameter [15]. In the default setting, we added Adapter blocks at the first two and last two transformer blocks (layers) in the image encoder. In the mask decoder, Adapters were added after the multi-head attention head. For PEFT with LoRA blocks, we implemented it based on SAMed [96], which inserted LoRA layers into the “query” and “value” projection layers of each transformer block in the image encoder and mask decoder.

In the pre-processing step of our datasets, we removed all data lacking target objects, specifically those images where the corresponding masks only consist of the background, for both training and testing purposes. We also followed the same normalization mean and standard deviation as the original SAM. We utilized dice loss and cross-entropy loss with equal weights. For all experiments across various datasets, we adapted a basic learning rate of 1e-4 with a warmup of $200$ iterations and applied AdamW with $0.1$ weight decay for the remaining iterations. The batch size was set to $4$. Histogram equalization and color jitter were the augmentations used during training. The maximum epochs were set to $200$ epochs, incorporating an early stopping mechanism that halted training if there was no improvement in the highest validation DSC for $20$ consecutive epochs. For settings that require heavy computing resources, such as ViT-H with the vanilla setting, we adopted Distributed Data Parallel (DDP) combined with model parallel to split batches or the model into 2-3 GPUs. The specifics regarding the total number of parameters in the model, the number of parameters subject to updates for each combination, and the computing resources needed are detailed in Table 3.

Figure 4 shows the performance of various task-specific fine-tuning approaches across each dataset and the performance of various UNet-based algorithms. The detailed results of each architecture-component-method fine-tuning strategy over 17 datasets can be found in the Appendix. Note that we provide the numerical results for all figure-based results in the Appendix. By comparing the red and blue bars, we first observe an almost definite (except in the ViT-B + Vanilla setting) improvement in the average performance when changing from only updating the decoder parameters to that of both the encoder and decoder. This discovery indicates that SAM’s original image encoder, initially trained on a vast collection of natural images, is not well-suited for efficiently extracting features from medical images, aligned with previous studies [61]. Therefore, freezing it and directly adapting the mask decoder for downstream tasks in the medical domain results in sub-optimal performance.

When comparing the performances among Vanilla, Aapter, and LoRA, we observe little difference between the two PEFT methods (Adapter and LoRA), and the rank between Vanilla and PEFT methods depends on the choices of backbones. When using the standard-size backbone (ViT-H or ViT-B), we observe an absolute advantage of using PEFT techniques. For example, when comparing the middle three blue bars in Figure 4, we observe using PEFT with Adapter or LoRA achieves $77.35\%$ and $78.93\%$ DSC respectively, while vanilla fine-tuning yields only $65.39\%$ DSC. We also conduct an additional study to insert the Adapter block after each layer and observe little difference. However, the difference between vanilla fine-tuning and PEFT is significantly reduced when using ViT-T, the small backbone introduced by MobileSAM [94]. We hypothesize the large gap between Vanilla and PEFT methods with a standard-size backbone is caused by the model collapse and knowledge-forgotten issues when fine-tuning all parameters of a large model [23, 26]. To support our hypothesis, we track the change of validation DSC on all datasets, and the examples of each modality are shown in Figure 5. It reveals that ViT-H + En/Decoder + Vanilla would cause a model collapse at an early stage. For instance, when applied to the MRI-BrainDuke dataset, this configuration results in a significant reduction in performance, with a decrease of approximately $35\%$ in DSC.

Next, our results reveal that the choice of encoder network architecture has minimal impact on performance when other choices remain the same. For instance, under the same settings of updating both encoder and decoder with Adapter (EN/Decoder + Adapter), we observe similar DSC scores across different encoders: ViT-H, ViT-B, and ViT-T. This similarity in performance occurs despite ViT-H having a significantly larger encoder with $6.8$ times network parameters compared with ViT-B and $64$ times parameters compared with ViT-T, as listed in Table 3. We deduce that while ViT-H excels in foundational settings for multiple segmentation tasks, it may be too parameter-rich for singular medical imaging segmentation tasks, given the typically smaller dataset sizes in the medical imaging field. Such property can also make the training difficult. As shown in Figure 5 top right, it takes $140$ epochs to train the ViT-H + En/Decoder + Adapter setting, which is approximately $30$ times longer when conducting the same setting with ViT-B, yet the latter gives slightly better performance ($80.53\%$ DSC vs. $81.21\%$ DSC when using ViT-H and ViT-B as the backbone respectively). Given the uniform results across our 17 datasets, we deduce that using PEFT (Adapter or LoRA) to fine-tune SAM within a smaller framework offers more benefits. This strategy strikes an optimal balance between efficiency and performance, especially for medical image segmentation tasks where datasets typically range from a couple hundred to tens of thousands of images.

Lastly, when comparing the results of automated segmentation by directly training on UNets, we observe that fine-tuning SAM with the En/Decoder generally yields better average performance. To further study the optimal performance across individual datasets, we select the best SAM-based method (ViT-H+LoRA+En/Decoder) vs. the best UNet-based method (nnUNet) in Figure 4, right block. It appears that SAM excels in tasks where the objects of interest are complicated, such as MRI-Brain Sclerosis or CT-Colon. Conversely, nnUNet demonstrates higher precision in tasks with simpler targets, including Xray-Hip or MRI-Kidney.

When implementing task-expansive prompt-based learning, we utilized a collection of datasets and their corresponding masks. Our experimental setup focused on training with all MRI datasets listed in Table 2 except “MRI-Knee”. In this collection criteria, we include 13 MRI datasets including 45k training image and non-empty mask pairs, and 7.9k evaluation images and mask pairs. We selected ViT-B as the encoder architecture since (1) it achieved similar performance to ViT-H in the data-specific fine-tuning setting; (2) MobileSAM was a distilled version of SAM which was not provided by initial SAM and thus less representative.

Under this multi-dataset supervised training setting, we configured a batch size of 3, utilizing DDP across 3 A6000 GPUs for training. The training was stopped when the validation performance was not improved for 5 epochs. To train the network, we fixed the prompt encoder and updated the image encoder and mask decoder in a vanilla setting (same as SAM’s supervised training stage). We generated a box or point prompt for each image as described in Section 2.3. Additionally, in this setting, we incorporated MedSAM [59] as a parallel approach since its training algorithm was similar to the proposed task-expansive prompt-based learning but with a different database.

In this experimental setup, our methodology follows a two-phase procedure, depicted in Figure 1 Setup 2: (1) Phase 1: task-expansive prompt-based learning. (2) Phase 2: task-specific supervised training for each dataset, beginning from pre-trained weights. We name the weights after phase 1 as SupSAM. The implementation details in the task-specific setting were identical to those of the previous section (Section 4.1). In phase 2, we had 2 model component choices (En/Decoder or Decoder only) and 3 Vanilla/PEFT choices in the automated setting. We evaluated the final automated segmentation performance after phase 2.

The results are shown in Figure 6 and Figure 7, in which we list the performance of the original method when starting from SAM’s initial weights in gray. We further illustrate the difference, highlighted in green (when positive) or red (if negative), when changing the initial weights from the original SAM to those acquired from task-expansive prompt-based learning or MedSAM, respectively. When task-specific supervised training is applied solely to the decoder (the upper 3 blocks in both figures), we observe a consistent performance decrease in both scenarios. Since in this case, phase 2 fine-tuning is applied exclusively to the decoder, it appears that weights learning through supervised training fail to adequately prepare the image encoder with the medical imaging domain knowledge necessary for automated feature extraction. Thus, attempting to conserve computational resources while starting with a publicly available, prompt-based medical model, such as MedSAM, becomes unnecessary if the goal is to update only the small decoder for customized datasets.

When fine-tuning both the encoder and decoder during phase 2 task-specific supervised training (Figure 6 and 7, bottom 3 blocks), we observe either a negative or no effect when utilizing a different pre-trained weights when using PEFT methods and a positive effect when using Vanilla fine-tuning. However, the latter setting leads to worse overall performance and thus offers little practical value. Despite the decrement being alleviated when using MedSAM, the overall trend shows the ineffectiveness of prompt-based learning when the ultimate goal is automated segmentation (MedSAM has around 34 times larger training set yet results in no performance change at best). We infer that the observed performance enhancement in En/Decoder-Vanilla is attributed to supervised pre-training potentially nudging the weights closer to the medical domain, thereby reducing the likelihood of model collapse. Nonetheless, given that this configuration still results in the lowest segmentation performance compared to the PEFT setting, this improvement does not offer practical benefit.

Our research indicates that supervised pre-training across various labeled medical images, despite working well in the original setting (MedSAM [59] is verified in its paper and ours is verified in Section 4.2), fails to enhance automated segmentation performance. We believe the ineffectiveness of additional supervised pre-training is due to the difference in the tasks: (1) The weights learned from interactive (prompt-based) mode vs. automated mode might be different and could not commute with each other effectively. Additionally, extended pre-training might disrupt the original weights’ transferability. (2) Furthermore, training the network on multiple segmentation tasks simultaneously could lead to conflicts within the training process thus making the training more challenging.

In this stage, all MRI datasets in Table 2 were used during training. The mask information was ignored and only the images from the training set were utilized. Blank slices were further removed since they were not informative for the reconstruction task. This resulted in $159,104$ images for training.

Following the previous section, we only select ViT-B as the encoder architecture. To employ MAE’s objective, we introduced an additional decoder to reconstruct the inputs, which was discarded after training. The additional decoder’s structure, mask ratio, and learning rate schedule all follow that of the original paper. Since the shape of SAM’s input was restricted to $1024\times 1024$, we set the batch size to $2$. To accommodate for the small batch size, we set the accumulative iteration, i.e. the number of iterations to only compute gradients without updating the network parameters, to $16$. The training was conducted on 8 A6000 for $50$ epochs.

In this experimental setup, our methodology follows a similar two-phase procedure: (1) Phase 1: task-agnostic self-supervised learning. (2) Phase 2: task-specific supervised training for each dataset, beginning from pre-trained weights. We name the weights after phase 1 as SSLSAM, which is shared publicly at https://github.com/mazurowski-lab/finetune-SAM.

Figure 8 presents the results of using SSLSAM as the weights before task-specific supervised training. We observe an improvement in the average performance of all MRIs across all architecture combinations besides En/Decoder + LoRA. Noticeably, when selecting En/Decoder + Adapter as the fine-tuning strategy, we observe consistent improvements over all MRI datasets, and the model achieves an average performance of $78.09\%$, which surpasses the leading ViT-H based method by $1.08\%$ and the best UNet-based method (nnUNet) by $3.58\%$. Contradictory to task-expansive prompt-based learning, the features learned during SSL are task agnostic, and since we only pre-train the encoder, there is no dependence on the task condition, which we have shown to hurt the performance. In the non-MRI datasets base, we observe a minor decrease in the performance. This can be due to the specificity of the MRI.

One of the popular scenarios for adopting fine-tuning strategies is when the target dataset has a limited number of annotated images, known as few-shot fine-tuning. Few-shot fine-tuning gets even more attention in the field of medical imaging segmentation, given the difficulty in securing a substantial quantity of expert annotations. We explore the effectiveness of different fine-tuning procedures under the following setting.

Our study primarily investigates 5-shot learning, selecting five images and their corresponding annotations randomly from the training set and validation set respectively. In the case of datasets derived from 3D volumes, these samples were chosen from distinct volumes to ensure diversity. In addition to the constraint of using only 5 training/validation images, we also modified the training hyper-parameters better suited for the few-shot setting. Specifically, we set a maximum of 1000 epochs for training, a batch size of 2, and did not include any early stopping criteria.

In this setting, we also investigated the previous three setups within a few-shot framework, including (1) task-specific fine-tuning exclusively, and (2) two pre-trained setups: task-expansive prompt-based learning and task-agnostic self-supervised learning, each followed by additional task-specific fine-tuning. To note, we cannot utilize the weights learned from task-expansive prompt-based learning since its training pipeline requires all labeled images. Thus, we consider the weights from MedSAM as an alternative. In this section, we investigate three key questions: (1) whether additional pre-training offers superior initial weights for the few-shot setting, (2) which fine-tuning strategy proves most effective in a few-shot context, (3) whether fine-tuning foundation models still work better than UNets in the few-shot setting?

Figure 9 shows the results for 5-shot setting, comparing the fine-tuning of SAM’s models with direct training on UNets. The results demonstrate that using ViT-B + En/Decoder + Adapter on the initial SAM yields optimal performance across all configurations. Comparing the results with UNets’, especially with the best-performing nnUNet, fine-tuned SAM still manages to achieve a slight improvement. This observation aligns with the findings presented in Section 4.1.

Consistently, the observation that MedSAM is ineffective persists in this few-shot scenario. While SSLSAM has shown effectiveness in the full-dataset setting, its efficacy diminished and even became detrimental in the few-shot scenario. This observation aligns with the understanding that the task-agnostic training disrupts the model’s connection with the segmentation tasks. While beneficial as initial weights for full-dataset training, this approach may not adequately rebuild these connections with only five images, resulting in lower performance. In summary, we observe the ineffectiveness of additional pre-training (either using task-expansive prompt-based learning or task-agnostic self-supervised learning) with limited training data. This can be contradictory to the common belief that imposing medical knowledge to the network in advance will be beneficial in the few-shot learning scenario.

We considered two types of prompts: point or box. To generate a prompt, we utilized the same strategy as in task-expansive prompt-based learning. The generated prompts were fed to the prompt encoder to obtain the prompt embeddings during both training and inference. The prompt encoder was fixed. We also limited the evaluation data to MRI only since this task was not our main focus.

We included the performance of applying prompts to the original SAM and its oracle performance; the latter reflects the hypothetical performance when conducting an extra step to determine the best prediction from the network. Additionally, we include weights acquired from task-expansive prompt-based learning and MedSAM in the box-prompt setting. When applying task-specific supervised training to SAM, we followed Section 4.1 to include all strategies.

Figure 10 shows the performance of all task-specific supervised training strategies in the interactive mode, as well as that of the original SAM as a reference. First, we observe the performance of SAM with point prompts is not satisfying. Replacing the point prompts with box ones significantly improves the performance. In this setting, we also examine the weights acquired from task-expansive prompt-based learning and observe similar performance (less than $0.92\%$ DSC difference compared to that of the initial SAM). This observation confirms that the inefficiency of task-expansive prompt-based learning is inherently from having different tasks in mind.

To be noted, the best automated task-specific supervised training strategy yields an average performance of $77.01\%$ on all MRI datasets, which is better than the original SAM with point prompts and on par with that with box prompts.

Conducting task-specific supervised training is also shown to be effective in the interactive segmentation scenario, especially when points are used as the prompts. Moreover, we find the best task-specific supervised training strategy is ViT-T + En/Decoder + Adapter in both scenarios, achieving $83.83\%/88.17\%$ DSC when using a point/box as the prompt.