PromptM-UNet: Efficient Text-Prompted 3D Medical Image Segmentation with Mamba SSM

PromptM-UNet is an ultra-lightweight, multi-modal 3D medical image segmentation framework that integrates Residual Vision Mamba (SSM) with multi-stage language conditioning. It enables accurate, prompt-guided volumetric CT segmentation while maintaining linear (O(N)) computational complexity and a minimal memory footprint feasible for standard consumer GPUs.

• GitHub Repository: https://github.com/kiuyha/PromptM-UNet
• Target Organ: Spleen (TotalSegmentator CT Cohort)

---

🌟 Key Features

• Linear (O(N)) Volumetric Modeling: Built on 3D Residual Vision Mamba (RVM) blocks from LightM-UNet, providing global receptive field modeling without quadratic (O(N^2)) Transformer memory explosion.
• Ultra-Lightweight Footprint: Active visual backbone is only ~1.87 Million parameters ((16\times) smaller than standard 3D nnU-Net, (100\times) smaller than SegVol).
• Zero-Overhead Language Conditioning: Decoupled text embedding caching ensures frozen language models (CLIP, BioBERT, Sentence-BERT) consume 0 extra GPU memory during training/inference.
• Dual-Resolution Spatial Zoom: SegVol-inspired dual-scale pipeline (3mm global context + 1.5mm high-resolution target crop).
• Clinical Performance: Achieves (0.862+) Dice Similarity Coefficient (DSC) and (0.792+) Normalized Surface Distance (NSD) on TotalSegmentator CT scans with under 4.3 GB validation VRAM.

---

📊 Benchmark & Ablation Results

The following results are evaluated on the official TotalSegmentator spleen validation cohort across 30 training epochs (from notebooks/ablation_promptm_unet.ipynb):

1. Fusion Strategy Ablation

Evaluated using Multiplication operation, CLIP text encoder, and 50:50 BCE supervision:

Fusion Strategy	Peak Epoch	Peak DSC	Peak NSD	Final DSC (Epoch 30)	Final NSD (Epoch 30)	Peak Val VRAM
Early Fusion	19	0.8622	0.7928	0.8196	0.7514	< 4.3 GB
All-Stage Fusion	24	0.8619	0.7871	0.7737	0.6848	< 4.3 GB
Late Fusion	12	0.8576	0.7675	0.7347	0.6095	< 4.3 GB

2. Loss Function Ablation

Evaluated across compound loss formulations:

Loss Configuration	Peak Epoch	Peak DSC	Peak NSD	Final DSC (Epoch 30)	Final NSD (Epoch 30)
Focal Loss (50:50)	20	0.8627	0.7827	0.8045	0.7195
Standard BCE (50:50)	24	0.8619	0.7871	0.7737	0.6848
Dynamic Loss (80:20 (\to) 65:35)	27	0.8537	0.7668	0.8280	0.7299
Focal Loss (80:20)	20	0.7827	0.6821	0.7358	0.6393

3. Fusion Operation Ablation

Evaluated using 50:50 BCE supervision:

Operation	Peak Epoch	Peak DSC	Peak NSD	Final DSC (Epoch 30)	Final NSD (Epoch 30)
Multiplication	24	0.8619	0.7871	0.7737	0.6848
Concatenation	26	0.8381	0.7647	0.8120	0.7387
Hybrid	28	0.8152	0.7301	0.7962	0.7041

4. Text Encoder Ablation

Evaluated with Early Fusion:

Text Encoder	Peak Epoch	Peak DSC	Peak NSD	Final DSC (Epoch 30)	Final NSD (Epoch 30)
CLIP	19	0.8622	0.7928	0.8196	0.7514
BioBERT	19	0.8466	0.7584	0.7968	0.7148

5. Multi-Tier Clinical Prompt Robustness

Evaluated on Best Early Fusion + CLIP checkpoint across prompt complexity tiers:

Prompt Tier	Clinical Description	Peak DSC
Tier (N)	Organ Name (e.g., "spleen")	0.8620
Tier (NS)	Name + Synonym (e.g., "spleen, lien")	0.8621
Tier (NL)	Name + Location (e.g., "spleen in left upper quadrant")	0.8624
Tier (NSL)	Name + Synonym + Location (e.g., "spleen, lien in upper left abdomen")	0.8622

---

🚀 Quickstart & Inference

1. Installation

# Clone the repository
git clone https://github.com/kiuyha/PromptM-UNet.git
cd PromptM-UNet

# Install dependencies and package
pip install -e .

2. Python Inference

import torch
import yaml
from promptm_unet.models.PromptMUNet import PromptMUNet

# Load configuration
with open("configs/default.yml", "r") as f:
    config = yaml.safe_load(f)

# Instantiate model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = PromptMUNet(config).to(device)

# Load checkpoint
checkpoint = torch.load("best_model.pth", map_location=device)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()

# Dummy 3D input: (Batch, Channels, Depth, Height, Width)
dummy_ct = torch.randn(1, 1, 64, 128, 128).to(device)
prompt = ["spleen segmentation in abdominal CT scan"]

# Forward pass
with torch.no_grad():
    prediction = model(dummy_ct, prompt)  # Output: (1, 1, 64, 128, 128)
    probabilities = torch.sigmoid(prediction)
    binary_mask = (probabilities > 0.5).cpu().numpy()

print(f"Predicted spleen mask shape: {binary_mask.shape}")

3. CLI Training & Evaluation

# Run training with default 50:50 compound supervision
accelerate launch -m promptm_unet.cli train \
  --path.raw_data_dir "/path/to/totalsegmentator_raw" \
  --path.prepro_data_dir "/path/to/preprocessed_data" \
  --training.batch_size 2 \
  --training.epochs 30 \
  --loss.dice_weight 0.5 \
  --loss.bce_weight 0.5

# Run multi-tier evaluation
python -m promptm_unet.cli test \
  --checkpoint "./experiments/checkpoints/best_model.pth"

---

🏗️ Model Architecture

3D CT Volume (1.5mm / 3.0mm)
        │
        ▼
 ┌───────────────┐
 │ Visual Encoder│ (Residual Vision Mamba - RVM Blocks)
 └───────┬───────┘
         │
         ▼
 ┌───────────────┐        ┌──────────────────┐
 │  Bottleneck   │ ◄────► │ Text Projection  │ ◄── Frozen CLIP / BioBERT / SBERT
 └───────┬───────┘        └──────────────────┘     (512D / 768D / 384D)
         │                         ▲
         ▼                         │ (Multi-Stage Hadamard Product)
 ┌───────────────┐                 │
 │ Visual Decoder│ ────────────────┘
 └───────┬───────┘
         │
         ▼
 Deep Supervision Heads ──► High-Resolution Spleen Binary Mask

---

📜 Citation & License

This project is licensed under the Apache License 2.0.

If you find this work useful in your research, please cite our repository:

@misc{promptm_unet2026,
  title={PromptM-UNet: Efficient Text-Prompted 3D Medical Image Segmentation Using State-Space Model Mamba for Spleen},
  author={Shridhara, Ketut and Setyawan, Ivan Andika and Al-Habib, Hasanuddin},
  year={2026},
  publisher={Universitas Negeri Surabaya},
  howpublished={\url{https://github.com/kiuyha/PromptM-UNet}}
}

---

🤝 Acknowledgements

• LightM-UNet for the 3D Mamba visual backbone.
• VoxTell & SegVol for multimodal prompting inspirations.
• TotalSegmentator for CT dataset annotations.
• Research funded by LPPM Universitas Negeri Surabaya (UNESA).