A Probabilistic Segment Anything Model for Ambiguity-Aware Medical Image Segmentation

Abdullah Imran; Tyler Ward

REVIEW 3 major objections 4 minor 15 references

A Probabilistic Segment Anything Model for Ambiguity-Aware Medical Image Segmentation

T0 review · 3 major / 4 minor · reviewed 2026-08-05 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read Probabilistic SAM beats U-Net on ambiguous lung CT segmentation

desk verdict A clean CVAE-on-SAM integration whose evaluation doesn't isolate the latent mechanism from SAM's backbone, so the diversity claim overreaches. read the letter →

arxiv 2509.05809 v1 pith:34Q3VCD7 submitted 2025-09-06 cs.CV

Tyler Ward , Abdullah Imran This is my paper

classification cs.CV

keywords probabilisticsegmentationSegmentAnythingModelconditionalvariationalautoencoderuncertaintymedicalimageLIDC-IDRIgeneralizedenergydistanceannotatorvariability

verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper tries to extend the Segment Anything Model (SAM) so that, given one image and one prompt, it produces a distribution of plausible segmentation masks instead of a single deterministic one. The proposed model, Probabilistic SAM, injects a conditional latent variable into SAM's pipeline: during training a posterior network sees the ground-truth mask, while a prior network sees only the image, and the sampled latent vector is added to the prompt embeddings before SAM's frozen mask decoder generates the mask. Trained on LIDC-IDRI lung CT slices with four expert annotators, the model reports a GED of 0.2910, DSC of 0.8255, and IoU of 0.7849, outperforming Probabilistic U-Net, Dropout U-Net, and Dropout SAM on these metrics. If correct, this gives clinicians and downstream models a way to see how much a lesion boundary is in dispute and to sample several plausible outlines from a single prompt.

What carries the argument

The central mechanism is the learned latent variable z with two Gaussian encoders: a posterior q(z|x,y) trained with the ground-truth mask and a prior p(z|x) trained with the image alone. A latent sample is projected by an MLP and added to SAM's sparse prompt embeddings; the modified embeddings plus image embeddings go into SAM's frozen mask decoder. At inference, drawing z from the prior is the only source of stochasticity, so repeating the draw yields diverse masks. The training objective combines BCE and Dice reconstruction loss with a KL term (weighted by β=10) that aligns posterior and prior.

What would settle it

Run deterministic SAM on the same LIDC-IDRI test set with the same box prompts and evaluation protocol (no latent sampling) and compute GED, DSC, and IoU. If deterministic SAM reaches or beats GED 0.2910, the claim that the latent space is what captures ambiguity is falsified.

Watch

Extended reading notes

Core claim

Probabilistic SAM claims that segmentation ambiguity can be modeled by conditioning a latent variable on the image and prompt, then using that latent code to modulate SAM's prompt embeddings before decoding. During training, a posterior network q(z|x,y) takes the image and ground-truth mask to define a Gaussian over z, while a prior network p(z|x) uses only the image; the two are aligned with a KL term, and the mask is reconstructed through SAM's decoder using BCE plus Dice loss. At inference, z is drawn from the prior, projected through an MLP, and added to the sparse prompt embedding, making the decoder's output stochastic. On the LIDC-IDRI test set the authors report GED 0.2910, DSC 0.825

Load-bearing premise

The paper assumes the reported improvements come from its latent mechanism, yet no deterministic SAM is evaluated, so a stronger pretrained backbone alone could be responsible for part or all of the gains.

Editorial extensions

If this is right

From one box or point prompt, a user can sample many plausible masks and inspect agreement across samples, making segmentation uncertainty visible.
Because only the prior network, MLP, and latent addition are new at inference, the added cost over deterministic SAM is small.
The lower GED reported means the sampled masks as a set are closer to the four-expert annotation distribution than Probabilistic U-Net's samples on this dataset.
The latent modulation is not tied to a specific prompt type or organ, so the same recipe can attach to any SAM prompt and, in principle, other promptable backbones.
The approach directly targets annotation ambiguity rather than only aleatoric noise, since diversity is trained against multiple expert masks.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

Because the experiments include no deterministic SAM arm, part of the reported gain may come from SAM's stronger pretrained encoder rather than the latent mechanism; a deterministic SAM baseline with identical prompts would isolate the latent's contribution.
GED rewards mean accuracy as well as diversity, so a model that simply outputs a sharper consensus mask could lower GED without genuinely capturing annotator variability; checking per-expert likelihoods or the spread of samples would test this.
If the latent space is well structured, latent samples should separate by expert tendency; clustering the inferred z vectors by annotator identity is a testable extension.
The same CVAE modulation could be applied to other promptable segmentation models, and the mechanism suggests a route to controllable ambiguity where the latent dimension corresponds to boundary tightness or style.

Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Desk Editor's Note

read the letter

Honestly, this is a straightforward transplant of Probabilistic U-Net's CVAE onto SAM, and the paper would be fine if the experiments actually isolated what the latent variable adds. As submitted, they don't.

The new bit is real: sampling a latent from an image-conditioned prior and adding it to SAM's sparse prompt embeddings is a neat arrangement, and the LIDC evaluation is new. The authors cite Probabilistic U-Net properly, make code public, and the ELBO is standard. If you want a quick reference for "SAM with a latent-variable uncertainty module," this is a usable template.

The problem is that Table 1 compares Probabilistic SAM against Probabilistic U-Net, Dropout U-Net, and Dropout SAM, but not against a deterministic SAM trained on the same split. GED rewards closeness to the target distribution in expectation, so a sharp consensus mask can get a low GED without any diversity. The reported gain over Probabilistic U-Net could be mostly SAM's stronger backbone. That may be useful in practice, but it doesn't support the claim that the latent mechanism captures annotator variability.

Also, the GED/DSC comparisons switch baselines (GED vs. PU-Net, DSC vs. Dropout SAM), there's no per-seed variance, and several implementation details are missing: latent dimension, posterior mask encoding, freeze/fine-tune status, GED sampling protocol. Beta=10 is hand-set with no sensitivity analysis. These are fixable but they matter for replication.

I don't think there's anything dishonest here; the method is clear and the qualitative examples show diversity. But the evidence for the central claim is weaker than the abstract suggests.

This is for people who want a starting point for probabilistic SAM, not for a deep analysis of uncertainty quantification. I'd send it to peer review, but a serious referee should push for a matched-backbone control and a diversity-only metric. With those, it could be a solid paper. Without them, the conclusion overreaches.

Referee Report

3 major / 4 minor

Summary. The paper proposes Probabilistic SAM, a CVAE-style extension of the Segment Anything Model (SAM) for ambiguity-aware medical image segmentation. A latent variable z is sampled from a posterior network during training and from an image-conditioned prior during inference; the sampled latent is projected and added to the sparse prompt embedding before SAM's mask decoder generates a mask. The model is trained with a reconstruction loss (BCE + Dice) plus a β-weighted KL term, with β=10. The method is evaluated on LIDC-IDRI lung nodule slices against Dropout U-Net, Probabilistic U-Net, and a Dropout SAM baseline, reporting generalized energy distance (GED), Dice, and IoU. The headline result is GED 0.2910 versus 0.3349 for Probabilistic U-Net, with notably higher DSC/IoU. The central claim is that the latent mechanism captures inter-expert variability and that the model outperforms existing probabilistic baselines on uncertainty-aware metrics.

Significance. If the reported gains are attributable to the probabilistic latent mechanism, the paper offers a simple and practical recipe for adding ambiguity modeling to a promptable foundation model, with code released and evaluation on a held-out test split against four external expert masks. The variational formulation is standard, and the architectural intervention—modulating prompt embeddings with a latent sample—is clean and easy to replicate. However, the current experimental design does not isolate the mechanism responsible for the gains: no deterministic SAM arm is reported, and GED conflates mean accuracy with sample diversity. The significance of the paper depends on closing this identification gap, because the conclusion that the probabilistic machinery 'captures annotator variability' is not yet supported by the evidence.

major comments (3)

[Section 3.3, Table 1] The experimental design does not include a deterministic SAM baseline trained or fine-tuned on the same split and prompt protocol. The abstract and conclusions attribute the GED/DSC/IoU improvements to the CVAE latent space, but the only SAM-based comparators are Dropout SAM and Probabilistic SAM. A deterministic SAM that emits a single accurate consensus mask could achieve a low GED because GED rewards closeness to the target distribution in expectation. Without a matched-backbone control, the observed gains are equally compatible with an explanation based on SAM's strong pretrained image encoder and mask decoder rather than on the prior/posterior networks or prompt-embedding modulation. Please add a deterministic SAM arm, or an ablation in which the latent is removed (e.g., fixed to its prior mean) while keeping the same backbone and training setup.
[Section 3.3, GED interpretation] The diversity claim is not supported by GED alone. GED (generalized energy distance) includes a term that rewards closeness between the model's samples and the target annotation distribution, so a model with high mean accuracy but zero diversity can still obtain a competitive GED. Probabilistic SAM's DSC/IoU are much higher than Probabilistic U-Net's (0.8255/0.7849 vs 0.5818/0.5557), so the GED advantage may reflect accuracy differences rather than the ability to reproduce annotator variability. Please report a diversity-only metric—for example, the expected pairwise Dice/IoU among sampled masks, or the E||X−X'|| term of GED separately—and, if possible, disaggregate results per expert target or within accuracy-matched bins.
[Section 3.2 and Section 3.3] The evaluation protocol for the headline metrics is underspecified. The paper does not define the GED estimator, state the number of Monte Carlo samples used to approximate it, describe the distance function used in GED (e.g., squared Euclidean distance on flattened masks), or explain how DSC/IoU are aggregated across the 1,996 test slices and across sampled masks and expert annotations. The paired t-test details (unit of pairing, correction for multiple comparisons) are also absent. Because the claimed GED improvement is only 4.39%, this specification is necessary for reproducibility and for interpreting whether the improvement is meaningful.

minor comments (4)

[Section 2, Eq. (2)] The BCE loss in Eq. (2) has an unbalanced parenthesis: the bracket is not closed. Also, the text writes N(µ, σ) while the implementations typically use variance; please clarify whether σ denotes standard deviation or variance.
[Section 3.2] Implementation details are incomplete: latent dimension, MLP size, optimizer, learning rate, batch size, number of training epochs, and whether the SAM image/prompt encoders are frozen or fine-tuned are not reported. These details matter for reproducibility and for judging whether the comparison is fair.
[Section 3.3] The one-tailed paired t-test is mentioned but no test statistic, degrees of freedom, or correction for multiple comparisons are provided. It is also unclear whether the pairing is per slice or per patient; please specify.
[Section 2, Eq. (4)] The choice β=10 is stated without justification or sensitivity analysis. Since the KL weight controls the trade-off between reconstruction fidelity and diversity, a small sensitivity study would strengthen the paper.

Circularity Check

0 steps flagged · score 0.0 of 10

No significant circularity; the CVAE derivation and held-out evaluation are self-contained.

full rationale

The paper's central derivation is a standard conditional variational autoencoder objective: Eq. (4) combines a BCE/Dice reconstruction loss with a KL term between posterior q(z|x,y) and prior p(z|x), and at inference latent samples from the prior are added to SAM's prompt embeddings. This is a conventional CVAE construction, not a definition that presupposes the target result. All quantitative claims in Table 1 (GED, DSC, IoU) are computed on a held-out LIDC-IDRI test split against external four-expert annotations; no fitted constant or trained parameter for the test set enters the metrics. The hand-set beta=10 is a hyperparameter, and the absence of a deterministic-SAM control is an experimental identification concern about attributing gains to the latent mechanism, not circularity. The self-citations (refs [4,5,7]) appear only in the introduction as related-work context for automating prompt generation and are not load-bearing for the probabilistic segmentation claim. No self-defined quantities, fitted inputs renamed as predictions, or uniqueness arguments imported from the authors' prior work are present. Therefore, no circular step can be exhibited.

Assumptions & free parameters 3 free parameters · 4 assumptions · 0 invented entities

The latent variable z is inherited from CVAE literature (refs [11, 14]); the paper introduces no new entity such as a particle, force, or conserved quantity. The main burden is carried by the hand-set beta, the unspecified latent dimension, and the assumption that prompt-embedding modulation suffices to express annotator diversity.

free parameters (3)

beta (KL weight) = 10
Hand-set in Eq. (4). It controls the trade-off between reconstruction quality and posterior-prior alignment, and with beta = 10 the KL term dominates; no sensitivity analysis is reported.
Latent dimension of z = unspecified
Design choice that sets the diversity capacity of the model; the paper never states it (Section 2).
Monte Carlo sample count for GED = unspecified
GED estimates depend on how many masks are sampled per image at inference; the paper does not state the protocol (Section 3.3).

assumptions (4)

domain assumption The four LIDC-IDRI expert annotations are exchangeable samples from the conditional segmentation distribution p(y|x).
GED treats expert masks as i.i.d. draws from the target distribution (Section 3.3); if experts are systematically different, the target distribution has structure that a Gaussian prior cannot express.
domain assumption Training with the beta-weighted ELBO makes the prior p(z|x) a faithful approximate posterior sampler over plausible masks, without posterior collapse.
Inference samples from the prior only (Section 2); if the KL term forces the posterior toward a low-variance prior, diversity vanishes and any GED gain reduces to mean accuracy.
domain assumption Adding the projected latent to the sparse prompt embedding is sufficient for the SAM mask decoder to express every plausible mask.
The mechanism changes only the prompt pathway (Fig. 2); the paper provides no ablation showing the decoder can realize the full annotator range this way.
ad hoc to paper The beta-weighted objective with beta = 10 is a valid bound for the segmentation CVAE.
Eq. (4) uses L_recon + 10 * KL; a large KL weight is not the standard ELBO and its effect on the learned prior is not characterized.

how reviews work

0 comments

Cite this review

Pith. "Pith review of A Probabilistic Segment Anything Model for Ambiguity-Aware Medical Image Segmentation." pith.science (2026). https://pith.science/paper/34Q3VCD7

@misc{pith2026250905809,
  author       = {Pith},
  title        = {Pith review of: A Probabilistic Segment Anything Model for Ambiguity-Aware Medical Image Segmentation},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/34Q3VCD7}},
  note         = {Machine review of arXiv:2509.05809}
}

read the original abstract

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/.

Figures

Figures reproduced from arXiv: 2509.05809 by the authors.

**Figure 2.** The sampling process of Probabilistic SAM. A prior network maps image embeddings to a Gaussian latent [PITH_FULL_IMAGE:figures/full_fig_p003_2.png] view at source ↗

**Figure 3.** Qualitative results of Probabilistic SAM vs. Probabilistic U-Net and the ground truth annotations from four [PITH_FULL_IMAGE:figures/full_fig_p004_3.png] view at source ↗

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

15 extracted references · 14 canonical work pages

[1]

Annotating Ambiguous Images: General Annotation Strategy for High-Quality Data with Real-World Biomedical Validation

Schmarje, L., Grossmann, V., Zelenka, C., Br¨ unger, J., and Koch, R., “Annotating ambiguous images: General annotation strategy for high-quality data with real-world biomedical validation,”arXiv preprint arXiv:2306.12189(2023)

work page Pith review arXiv 2023
[2]

Is segmentation uncertainty useful?,

Czolbe, S., Arnavaz, K., Krause, O., and Feragen, A., “Is segmentation uncertainty useful?,” in [Interna- tional Conference on Information Processing in Medical Imaging], 715–726, Springer (2021)

work page 2021
[3]

Segment anything,

Kirillov, A., Mintun, E., Ravi, N., Mao, H., Rolland, C., Gustafson, L., Xiao, T., Whitehead, S., Berg, A. C., Lo, W.-Y., et al., “Segment anything,” in [Proceedings of the IEEE/CVF International Conference on Computer Vision], 4015–4026 (2023)

work page 2023
[4]

Annotation-efficient task guidance for medical Segment Anything,

Ward, T. and Imran, A. A. Z., “Annotation-efficient task guidance for medical Segment Anything,” in [2025 IEEE 22nd International Symposium on Biomedical Imaging (ISBI)], 1–4, IEEE (2025)

work page 2025
[5]

Detection of Breast Cancer Lumpectomy Margin with SAM-incorporated Forward-Forward Contrastive Learning

Ward, T., Wang, X., McFarland, B., Ahamed, M. A., Nozad, S., Arshad, T., Nebbache, H., Chen, J., and Imran, A., “Detection of breast cancer lumpectomy margin with sam-incorporated forward-forward contrastive learning,”arXiv preprint arXiv:2506.21006(2025)

work page Pith review arXiv 2025
[6]

Autoprosam: Automated prompting sam for 3d multi-organ segmentation,

Li, C., Sultan, R. I., Khanduri, P., Qiang, Y., Indrin, C., and Zhu, D., “Autoprosam: Automated prompting sam for 3d multi-organ segmentation,” in [2025 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV)], 3570–3580, IEEE (2025)

work page 2025
[7]

Autoadaptive medical Segment Anything Model,

Ward, T., Owen, M. K., Coleman, O., Noehren, B., and Imran, A.-A.-Z., “Autoadaptive medical Segment Anything Model,”arXiv preprint arXiv:2507.01828(2025)

work page arXiv 2025
[8]

Segment anything in medical images,

Ma, J., He, Y., Li, F., Han, L., You, C., and Wang, B., “Segment anything in medical images,”Nature Communications15(1), 654 (2024)

work page 2024

Show all 15 references

[9]

Flaws can be applause: Unleashing potential of segmenting ambiguous objects in SAM,

Li, C., Li, W., Liu, H., Liu, X., Xu, Q., Chen, Z., Huang, Y., and Yuan, Y., “Flaws can be applause: Unleashing potential of segmenting ambiguous objects in SAM,”Advances in Neural Information Processing Systems37, 45578–45599 (2024)

2024
[10]

Trustworthy clinical AI solutions: A unified review of uncertainty quantification in deep learning models for medical image analysis.,

Lambert, B., Forbes, F., Doyle, S., Dehaene, H., and Dojat, M., “Trustworthy clinical AI solutions: A unified review of uncertainty quantification in deep learning models for medical image analysis.,”Artificial Intelligence in Medicine150, 102830 (2024)

2024
[11]

Learning structured output representation using deep conditional generative models,

Sohn, K., Lee, H., and Yan, X., “Learning structured output representation using deep conditional generative models,”Advances in neural information processing systems28(2015)

2015
[12]

Auto-encoding variational bayes,

Kingma, D. P. and Welling, M., “Auto-encoding variational bayes,”arXiv preprint arXiv:1312.6114(2013)

2013 arXiv
[13]

The Lung Image Database Consortium (LIDC) and Image Database Resource Initiative (IDRI): A completed reference database of lung nodules on CT scans,

Armato III, S. G., McLennan, G., Bidaut, L., McNitt-Gray, M. F., Meyer, C. R., Reeves, A. P., Zhao, B., Aberle, D. R., Henschke, C. I., Hoffman, E. A., et al., “The Lung Image Database Consortium (LIDC) and Image Database Resource Initiative (IDRI): A completed reference datab...

2011
[14]

A probabilistic U-Net for segmentation of ambiguous im- ages,

Kohl, S., Romera-Paredes, B., Meyer, C., De Fauw, J., Ledsam, J. R., Maier-Hein, K., Eslami, S., Jimenez Rezende, D., and Ronneberger, O., “A probabilistic U-Net for segmentation of ambiguous im- ages,”Advances in Neural Information Processing Systems31(2018)

2018
[15]

Energy statistics: A class of statistics based on distances,

Sz´ ekely, G. J. and Rizzo, M. L., “Energy statistics: A class of statistics based on distances,”Journal of statistical planning and inference143(8), 1249–1272 (2013)

2013

Pith tools

Reviewed August 5, 2026 · model on record in the stance chip above.

[1] [1]

Annotating Ambiguous Images: General Annotation Strategy for High-Quality Data with Real-World Biomedical Validation

Schmarje, L., Grossmann, V., Zelenka, C., Br¨ unger, J., and Koch, R., “Annotating ambiguous images: General annotation strategy for high-quality data with real-world biomedical validation,”arXiv preprint arXiv:2306.12189(2023)

work page Pith review arXiv 2023

[2] [2]

Is segmentation uncertainty useful?,

Czolbe, S., Arnavaz, K., Krause, O., and Feragen, A., “Is segmentation uncertainty useful?,” in [Interna- tional Conference on Information Processing in Medical Imaging], 715–726, Springer (2021)

work page 2021

[3] [3]

Segment anything,

Kirillov, A., Mintun, E., Ravi, N., Mao, H., Rolland, C., Gustafson, L., Xiao, T., Whitehead, S., Berg, A. C., Lo, W.-Y., et al., “Segment anything,” in [Proceedings of the IEEE/CVF International Conference on Computer Vision], 4015–4026 (2023)

work page 2023

[4] [4]

Annotation-efficient task guidance for medical Segment Anything,

Ward, T. and Imran, A. A. Z., “Annotation-efficient task guidance for medical Segment Anything,” in [2025 IEEE 22nd International Symposium on Biomedical Imaging (ISBI)], 1–4, IEEE (2025)

work page 2025

[5] [5]

Detection of Breast Cancer Lumpectomy Margin with SAM-incorporated Forward-Forward Contrastive Learning

Ward, T., Wang, X., McFarland, B., Ahamed, M. A., Nozad, S., Arshad, T., Nebbache, H., Chen, J., and Imran, A., “Detection of breast cancer lumpectomy margin with sam-incorporated forward-forward contrastive learning,”arXiv preprint arXiv:2506.21006(2025)

work page Pith review arXiv 2025

[6] [6]

Autoprosam: Automated prompting sam for 3d multi-organ segmentation,

Li, C., Sultan, R. I., Khanduri, P., Qiang, Y., Indrin, C., and Zhu, D., “Autoprosam: Automated prompting sam for 3d multi-organ segmentation,” in [2025 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV)], 3570–3580, IEEE (2025)

work page 2025

[7] [7]

Autoadaptive medical Segment Anything Model,

Ward, T., Owen, M. K., Coleman, O., Noehren, B., and Imran, A.-A.-Z., “Autoadaptive medical Segment Anything Model,”arXiv preprint arXiv:2507.01828(2025)

work page arXiv 2025

[8] [8]

Segment anything in medical images,

Ma, J., He, Y., Li, F., Han, L., You, C., and Wang, B., “Segment anything in medical images,”Nature Communications15(1), 654 (2024)

work page 2024