A Probabilistic Segment Anything Model for Ambiguity-Aware Medical Image Segmentation
This addresses the problem of segmentation ambiguity in medical imaging for clinicians and researchers, though it is an incremental improvement over SAM.
The paper tackles the problem of deterministic segmentation models failing to capture ambiguity in medical imaging by introducing Probabilistic SAM, a probabilistic extension of SAM that models a distribution over segmentations, and demonstrates its ability to produce diverse outputs aligning with expert disagreement on the LIDC-IDRI lung nodule dataset, outperforming existing probabilistic baselines on uncertainty-aware metrics.
Recent advances in promptable segmentation, such as the Segment Anything Model (SAM), have enabled flexible, high-quality mask generation across a wide range of visual domains. However, SAM and similar models remain fundamentally deterministic, producing a single segmentation per object per prompt, and fail to capture the inherent ambiguity present in many real-world tasks. This limitation is particularly troublesome in medical imaging, where multiple plausible segmentations may exist due to annotation uncertainty or inter-expert variability. In this paper, we introduce Probabilistic SAM, a probabilistic extension of SAM that models a distribution over segmentations conditioned on both the input image and prompt. By incorporating a latent variable space and training with a variational objective, our model learns to generate diverse and plausible segmentation masks reflecting the variability in human annotations. The architecture integrates a prior and posterior network into the SAM framework, allowing latent codes to modulate the prompt embeddings during inference. The latent space allows for efficient sampling during inference, enabling uncertainty-aware outputs with minimal overhead. We evaluate Probabilistic SAM on the public LIDC-IDRI lung nodule dataset and demonstrate its ability to produce diverse outputs that align with expert disagreement, outperforming existing probabilistic baselines on uncertainty-aware metrics. Our code is available at: https://github.com/tbwa233/Probabilistic-SAM/.