Skip to content

Diffusion-empowered AutoPrompt MedSAM

šŸ•’ Published (v1): 2025-02-05 03:08 UTC Ā· Source: Arxiv Ā· link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

šŸ’¬ Ask ChatGPTāœ¦ Ask Claude

TL;DR

AutoMedSAM replaces MedSAM's manual prompt encoder with a diffusion-based class prompt encoder that takes organ/lesion class indices as sole input, enabling fully automated end-to-end segmentation. It also injects semantic labels into output masks and introduces an uncertainty-aware joint optimization strategy that adaptively balances four heterogeneous loss functions via learnable uncertainty weights.

Problem

MedSAM requires clinicians to supply precise explicit prompts (points, bounding boxes, scribbles) per segmentation target, which is impractical in busy clinical settings and sensitive to prompt placement accuracy. Generated masks also lack semantic labels, limiting utility for non-experts. Existing auto-prompting methods use organ-specific detector models that do not generalize and still lack class-level semantic grounding.

Method

Architecture: AutoMedSAM retains MedSAM's frozen ViT image encoder and mask decoder, replacing only the prompt encoder with a diffusion-based class prompt encoder (EP).

Forward conditional diffusion: A class index c is projected via a linear layer to match image embedding dimensions, then added to Gaussian noise at each diffusion timestep t (noise scale Ļƒt = 1/(t+1)), producing a class-conditioned noisy embedding Ft = FI + Ļµt + cexpand.

Two-branch reverse diffusion: A shared convolutional encoder processes Ft; two independent decoder branches generate: - Dense prompt embeddings via element-wise attention (capturing local fine-grained structure) - Sparse prompt embeddings via channel-wise attention with global average pooling (capturing global class information)

Skip connections and per-branch class re-injection reinforce class specificity throughout decoding.

Uncertainty-aware joint optimization: Four loss functions (MSE on prompt embeddings for MedSAM knowledge transfer, cross-entropy, Dice, and Shape Distance) are combined as: L = Ī£ [1/(2Ī»jĀ²) Lj + log(1 + Ī»jĀ²)] where Ī»j are learnable uncertainty parameters that adaptively weight each loss component without manual hyperparameter tuning. Only EP and the mask decoder are updated; the image encoder is frozen.

Key Contributions

  • End-to-end class-prompt-driven segmentation model that eliminates all explicit manual prompts from MedSAM
  • Diffusion-based class prompt encoder with dual-branch (dense/sparse) reverse diffusion for jointly capturing local and global class-specific features
  • Uncertainty-aware joint optimization combining pixel-, region-, and distribution-based losses with learnable adaptive weighting
  • Semantic label propagation: output masks are associated with class indices, enabling interpretation by non-experts

Results

All metrics are DSC (%) / NSD (%) unless noted; baselines include SAM, SAM2, MedSAM, Med2D, U-MedSAM, SAMed, H-SAM, AutoSAM, SurgicalSAM, nnU-Net, Swin-UNet, MedFormer.

  • AbdomenCT-1K (multi-organ CT): 94.580 / 95.148 ā€” best overall; next best is U-MedSAM at 92.979 / 91.158
  • BraTS (brain tumor MR): 91.057 / 92.661 ā€” best; next best H-SAM at 90.615 / 90.288
  • Kvasir-SEG (polyp endoscopy): 96.828 / 98.729 ā€” best; next best MedSAM at 95.803 / 97.829
  • Chest-XML (lung X-ray): 96.941 / 97.367 ā€” best; next best U-MedSAM at 96.813 / 97.310
  • Cross-dataset generalization (AbdomenCT-1K ā†’ AMOS): 71.141 / 77.282 vs SurgicalSAM 56.930 / 64.451
  • Cross-modality generalization (BraTS-FLAIR ā†’ BraTS-T1CE): 52.513 / 48.101 vs SurgicalSAM 36.843 / 39.825 (+15.67% DSC)
  • Ablation ā€” diffusion disabled (Chest-XML): DSC drops 5.379%, NSD drops 5.173%
  • Ablation ā€” uncertainty weighting replaced by fixed weights: DSC drops 3.582%, NSD drops 3.254%
  • Ablation ā€” MSE loss removed (knowledge transfer): DSC drops 4.368%, largest single-loss impact

Limitations

  • Class index must still be specified by the user; the model does not autonomously detect what is present in the image
  • Evaluated on only four datasets spanning CT, MR, endoscopy, and X-ray ā€” fine-grained or rare anatomy not assessed
  • Cross-dataset experiments compare only against SurgicalSAM, not the full baseline suite
  • No user study to substantiate the claimed usability improvement for non-expert users
  • Single-GPU training (RTX A40); inference overhead of the diffusion encoder relative to standard prompt encoding is not characterized

Relevance to Foundation Models in Medicine

AutoMedSAM demonstrates a practical strategy for adapting a large medical segmentation foundation model (MedSAM) to fully automated clinical deployment by surgically replacing only the prompt interface while preserving pretrained image encoder knowledge. The use of a diffusion process to generate prompt embeddings ā€” rather than conditioning on spatial coordinates ā€” provides a template for converting interactive foundation models into class-driven systems applicable to multi-organ and multi-modality settings. The uncertainty-aware loss weighting addresses a recurring challenge in medical foundation model fine-tuning: balancing heterogeneous supervision signals across highly varied imaging domains. This work sits within a growing line of research on removing the human-in-the-loop bottleneck from SAM-derived medical foundation models.