MedSAM-U: Uncertainty-Guided Auto Multi-Prompt Adaptation for Reliable MedSAM

Traditional deep learning methods [14, 2, 15] are mostly designed for specific tasks and often face challenges in transferring to other domains. SAM [6] is the pioneering large foundation model for segmentation, consists of three primary components: an image encoder, a prompt encoder, and a mask decoder. The image encoder is based on a standard Vision Transformer (ViT) that has been pretrained using Masked Autoencoders (MAE). The prompt encoder can operate in either a sparse (e.g.,boxes) or dense (e.g., masks) manner. The mask decoder is a Transformer decoder block adapted to include a dynamic mask prediction head. This decoder employs two-way cross-attention mechanisms to capture the interactions between the prompt and image embeddings. Subsequently, SAM upsamples the image embedding, subsequently, SAM upsamples the image embedding, and a Multi-Layer Perceptron (MLP) maps the resulting token to a dynamic linear classifier that predicts the ground truth (GT) for the given image I. Due to its strong zero-shot performance, SAM marks a major breakthrough in the field of image segmentation.

To improve the unsatisfactory performance of SAM on medical image segmentation tasks, some approaches are to fine-tune SAM on medical images, including full fine-tuning and parameter-efficient fine-tuning [16, 17, 18, 19]. Recently, MedSAM [7] has investigated SAM’s application to medical image segmentation, exploring its performance in various contexts such as endoscopic surgery [20], tumor segmentation [21], polyp segmentation [22]. However, the existing MedSAM-based methods rarely investigate the sensitivity of different prompts.

Prompts-based methods have gained traction in medical image segmentation due to their ability to provide flexible and adaptive guidance during interactive segmentation tasks. While effective in some cases, current adaptations of SAM rely heavily on high-quality, standard prompts (such as points, boxes, and masks) to deliver satisfactory performance in medical image segmentation tasks. Wu et al. [23] introduce Space-Depth Transpose for adapting 2D SAM to 3D medical images and Hyper-Prompting Adapter for prompt-conditioned adaptation. Deng et al. [24] propose a multi-box prompt-triggered uncertainty estimation for SAM that uses Monte Carlo methods and test-time augmentation to enhance performance and provide pixel-level reliability for segmented lesions or tissues. Li et al. [25] propose a 3D medical image segmentation model using a single point prompt with SAM’s pretrained vision transformer and lightweight adapters, featuring a hybrid network and boundary-aware loss for precise results. Wu et al. [26] presents One-Prompt Segmentation method that combines one-shot and interactive approaches to handle unseen tasks with a single prompt in one pass. Currently, most MedSAM-based methods are limited to using a single type of prompt rather than multiple types of prompts [27, 28, 24, 25]. This reduces the available information for the model and may result in insufficient segmentation performance.

Uncertainty estimation [29] in segmentation has become increasingly important, as it offers a way to assess the confidence in model predictions. This is crucial in high-stakes applications such as medical image analysis, where mistakes can lead to serious repercussions. There are two primary types of uncertainty—aleatoric, which originates from the inherent noise in the data [30], and epistemic, which is due to the model’s limitations—is the key to enhancing the reliability of deep neural networks [31]. For instance, Saad et al. [32] utilized shape and appearance priors to quantify uncertainty in probabilistic medical image segmentation. Parisot et al. [33] leveraged segmentation uncertainty to inform adaptive sampling strategies for the simultaneous segmentation and registration of brain tumors. Additionally, Prassni et al. [34] developed a method to visualize uncertainty in random walker-based segmentation, which was then applied to volumetric segmentation of brain MRI and abdominal CT images. Zou et al. [35]. proposed a trusted brain tumor segmentation network that generates robust segmentation results and reliable uncertainty estimations by leveraging subjective logic theory to model uncertainty and parameterize class probabilities as a Dirichlet distribution.

Additionally, uncertainty estimation in large models, such as SAM, is critical for improving the reliability of segmentation outputs [24]. Yao et al. [27] propose a test-phase prompt augmentation technique for SAM that integrates multi-box augmentation with an aleatoric uncertainty-based FN and FP correction strategy to improve medical image segmentation. Zhang et al. [28] propose UR-SAM, an uncertainty-rectified framework that enhances SAM’s reliability in medical image segmentation by combining prompt augmentation with uncertainty-based rectification. Despite these advancements, research on integrating uncertainty estimation with prompt-based MedSAM remains insufficient, particularly in how to utilize uncertainty to guide reliable prompts.

MedSAM [7] primarily relies on box prompts as the initial input for segmentation tasks, without incorporating point prompts during its training process. This limitation means that MedSAM’s segmentation capabilities are optimized for scenarios where box annotations are available, but it may not fully exploit the potential benefits of point-based inputs. Formally, in MedSAM, the relationship between the inputs and the prediction mask $\boldsymbol{y}$ can be formally expressed as:

the function $f_{\text{MSAM}}$ can be denoted as:

where I denotes the input image, $b$ represents the box prompt, $\mathcal{F}_{D}$, $\mathcal{F}_{E}$ and $\mathcal{F_{P}}$ represent the Decoder, Image Encoder, Prompt Encoder modules of MedSAM, respectively, and $\boldsymbol{y}$ is the resulting segmentation mask of MedSAM.

Moreover, as detailed in [8, 19], the effect of point prompts on segmentation performance is demonstrated, particularly positive points. Accordingly, to improve the sparse encoder, this study proposes modifications that refine the encoding of points and boxes. These adjustments aim to optimize the integration of positional encoding with learned embeddings for each prompt type. We propose MPA that is designed to handle and adapt multiple types of prompt inputs for MedSAM, including the combination of points and boxes prompts.

Our goal is to extend the capabilities of MedSAM by adapting it to handle multi-prompt through a fine-tuning approach, called MPA-MedSAM, designed to enhance MedSAM’s flexibility without the need for full re-training. Instead of adjusting all parameters, we keep the pre-trained MedSAM parameters frozen except for the prompt encoder, develop an adapter module, and integrate it into the image encoder designated positions. Structurally, the adapter serves as a bottleneck model, consisting of three key components: a down-projection, ReLU activation, and up-projection sequentially, as illustrated in Fig. 2 (a), similar to Med-SA [23]. Given that MedSAM does not show improvements when multi-type prompts are incorporated, we unfreeze the prompt encoder to expand its capabilities. To break the limitation, we have unfrozen the prompt encoder to expand its capabilities. This adjustment allows MedSAM to effectively handle multi-type prompts, addressing the limitations encountered in earlier versions. For interactive segmentation, both point prompts and Bounding box (BBox) prompts are utilized during training. The prompts are generated by randomly selecting points and applying jitter to the Bbox derived from the segmentation mask. This simulates varying levels of user inaccuracy, making the model more robust to real-world variations and enhancing its adaptability in practical applications. The relationship between the inputs and the prediction mask $\boldsymbol{y}$ in MPA-MedSAM can be formally expressed as:

where I denotes the input image, $p$ represents the point prompt, $b$ represents the Bbox prompt, and $\boldsymbol{y}$ is the resulting segmentation mask of MPA-MedSAM. The function $f_{\text{MPA-MedSAM}}$ encapsulates the MPA-MedSAM, which processes the input image, point prompt, and box prompt to produce the corresponding mask.

Given the variability of prompts from users with differing levels of experience, MedSAM’s segmentation performance exhibits significant sensitivity to the position of Bboxes, potentially leading to inference errors. To address this, we adopt a strategy similar to the one used by SAM-U [24]. In this study, we introduce UGMP to incorporate multiple prompts to improve the accuracy of MPA-MedSAM. Assume giving an initial Bbox $b^{\text{init}}$, by applying a simple random augmentation strategy to the $b^{\text{init}}$, we generate $N$ multiple Bbox prompts as follows:

where $\delta_{i}\sim\mathcal{N}(\mu,\sigma)$, $b^{\text{init}}$ represents the initial Bbox, these random boxes are generated by applying adjustments $\delta_{i}$ to $b^{\text{init}}$, $\mathcal{N}$ denotes the normal distribution, $\mu$ denotes the coordinates of the $b^{\text{init}}$, $\delta$ represent the degree of perturbation applied to the $b^{\text{init}}$. The operation generate a set of multiple Bbox prompts $\mathbb{B}=\{b_{1},b_{2},\dots,b_{N}\}$ , and a set of $M$ multiple point prompts $\mathbb{P}=\{p_{1},p_{2},\dots,p_{M}\}$ that depends on the user’s or clinician’s choice, $\mathbb{P}=\emptyset$ if no points are used. To quantify the uncertainty arising from the use of multiple prompts, with $N$ box prompts, $M$ point prompts and input image I, MPA-MedSAM predict a set of results $\mathbb{Y}=\{{\boldsymbol{y}}_{1},{\boldsymbol{y}}_{2},\cdots,{\boldsymbol{y}_{N}}\}$, where $\boldsymbol{y}_{i}$ is predicted as follows:

where $b_{i}$ represents the $i$-th box of the set $\mathbb{B}$, $y_{i}$ represents the segmentation result obtained by inputting $b_{i}$ into MedSAM.

The aggregation of these predictions leads to enhanced segmentation accuracy and a reduction in uncertainty. The final combined prediction is formulated as follows:

$\overline{\boldsymbol{y}}$ represents the average segmentation result obtained by applying the MPA-MedSAM model across $N$ instances.

By ultilizing UGMP, the aleatoric uncertainty from a single given image $I$ , instead of being estimated from each individual prediction $\boldsymbol{y}_{i}$, is now directly estimated from the average prediction $\overline{\boldsymbol{y}}$, described by the entropy [36] :

where $f_{\text{UGMP}}\left(\overline{\boldsymbol{y}}\right)=-\int p(\overline{\boldsymbol{y}}|\textit{I})\log{p(\overline{\boldsymbol{y}}|\textit{I})}d\overline{\boldsymbol{y}}$. This allows us to estimate the uncertainty distribution based on the aggregated prediction $\overline{\boldsymbol{y}}$, rather than calculating it for each individual prediction $\boldsymbol{y}_{i}$.

In summary, the MPA-MedSAM with UGPM function can be expressed as follows:

where I represents the input image, $P_{b}$ are the multiple Bbox prompts $\{b_{1},b_{2},\dots,b_{N}\}$, $P_{p}$ are the points prompts $\{p_{1},p_{2},\dots,p_{M}\}$, also depending on the user’s or clinician’s choice. $\overline{\boldsymbol{y}}$ is the resulting segmentation mask based on the combination of predictions from multiple prompts, and $\overline{\textbf{{U}}}$ is the associated uncertainty map.

Providing prompts that produce desired results can be a difficult process that often requires the users or experts to go through tedious trial-and-error experimentation. Due to the uncertainty introduced by different prompts during the inference process, we investigated whether this knowledge could be utilized to refine the initial prompts, thereby producing more predictable and precise segmentation mask outputs. In an effort to address this issue, We introduce UGPA over a series of sample prompts to refine them, through UGPA could improve the model’s segmentation through automatic processes. Fig. 2 (b) shows the schematic view of our proposed UGPA.

In the UGPA process, firstly, we select the Top-K Bboxes based on the edges of the U, which highlight regions of high uncertainty. These boxes are then subjected to slight adjustments, simulating expert adjustments to refine their accuracy. Secondly, following this, we select the Top-K points with the highest uncertainty values from the U, representing the most uncertain areas within the image. and then these selected points are then combined with the adjusted Bboxes to create refined prompts. To provide a clear and intuitive description of the process, we represent each step using the following formulas:

where:

$\text{Edge}(\textbf{{U}})$ refers to identifying the coordinates of the edges of the uncertainty map, $(w_{j},h_{j})$ refers to the coordinates of the edge points. Specifically, $(w_{\min},h_{\min})$ represents the coordinates of the top-left corner, and $(w_{\max},h_{\max})$ represents the coordinates of the bottom-right corner of the bounding box that encloses these edge points. We select the Top-1 Bbox based on the edges of the U, which has the largest area and fully contains the edge of U. Then, we also apply Eq. (4) to simulate varying levels of expert knowledge, to generate a set of refined box prompts $\mathbb{B}^{*}=\{b_{1},b_{2},\dots,b_{k}\}$ . Then, we select the Top-K points with the highest uncertainty values from the U:

$\mathcal{S}$ denotes the sort function, $\{p_{1},p_{2},\dots,p_{k}\}$ represents a set of the selected $k$ points $\mathbb{P}^{*}$. we combine the adjusted Bboxes with the selected points to create the refined prompts set: $\{\mathbb{B}^{*},\mathbb{P}^{*}\}$.

Subsequently, these refined prompts are re-as-input into MPA-MedSAM, where they serve as crucial inputs to guide further refinement of the segmentation output. By leveraging these improved prompts, our method can enhance the accuracy of the segmentation, narrowing down errors and producing more precise results. This process ensures that the model benefits from the enhanced guidance provided by the refined prompts, ultimately leading to superior segmentation performance. The refined prompts then input into MPA-MedSAM with UGPM function can be expressed as follows:

$\boldsymbol{y}^{*}$ is the segmentation mask based on the UPGA, and $\textbf{{U}}^{*}$ is the associated uncertainty map.

Finally, the refined process is informed by Active Learning principles: we assess the output uncertainty following the prompt optimizations and update the segmentation results only if these optimizations lead to reduced uncertainty estimates. Importantly, the final output is not considered definitive until it is rigorously verified against the uncertainty map to ensure that performance has indeed improved. It can be denoted by:

The detials of the proposed method MedSAM-U are outlined in Algorithm. 1.

To validate the overall performance of our proposed method, we compared it with the SOTA segmentation foundation model across five different modalities. As shown in Table I, include comparisons with SAM [6], and fully fine-tuned SAM in medical (MedSAM) [7], evaluated using Dice Score and Iou Score. The results were shown in Table I.

Firstly, when comparing the performance of the same model under different Bboxes qualities (ratio from 0.5 to 0.75), a notable improvement is observed in SAM across all imaging modalities as the Bboxes ratio increases from 0.5 to 0.75. MedSAM exhibits a similar trend, with performance enhancements corresponding to the higher Bboxes quality. It reveals that the quality of the Bboxes significantly impacts segmentation performance, indicating that the models are sensitive to the quality of the Box prompts.

Furthermore, for our proposed method, represented as MedSAM-U 3B (0.5), and 3 B (0.75), shown in the $5^{th}$ row and the $6^{th}$ row, achieved SOTA performance in Dermoscopy, Colonoscopy, Ultrasound and CT, with a final Avg Dice of 90.9 % , 90.3 % , 95.8% and 78.3% with 3B (0.75), surpassing MedSAM by 4.8%, 2.3%, 2.7% and 2% respectively. When the Bboxes overlap ratio is 0.5, our method outperforms surpassing MedSAM by 20.5 %. The results shown in Table I demonstrate that, even when the initial Bboxes quality is low, significant improvements in segmentation performance are observed. This demonstrates the effectiveness of our approach in enhancing segmentation accuracy, particularly in scenarios where the initial Bx prompts are suboptimal. These results reveal that our MedSAM-U demonstrates significant accuracy across various medical segmentation tasks with the use of low-quality Bboxes, eliminating the need for manual automantic adjustments to achieve satisfactory results.

Fig. 3 visualized Some examples’ segmentation results for each method in different modalities, with Bboxes simulating varying degrees of rough approximations simulating expert annotations. SAM and MedSAM’s segmentations are based on a single step, and when the initial Bbox prompts provided by experts are not perfectly accurate, the segmentations may suffer from under-segmentation or over-segmentation errors. In contrast, MedSAM-U can accurately segment various targets across different imaging conditions.