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

PromptM-UNet combines a LightM-UNet-derived Mamba backbone with frozen language encoding for text-guided 3D CT segmentation. The selected configuration uses MiniLM and bottleneck feature modulation, producing one binary mask for the requested spleen, liver, or kidney-union target.

• GitHub Repository: https://github.com/kiuyha/PromptM-UNet
• Target Organs: Spleen, Liver & Kidney (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.
• Compact Model: The selected deployed configuration contains 27.83M parameters, including the frozen MiniLM text encoder and excluding auxiliary training heads.
• Frozen Language Conditioning: CPU text encoding and cached embeddings reduce repeated language-model work; host memory and text-encoding costs remain part of deployment.
• Dual-Resolution Spatial Zoom: Coarse localization at 3 mm and organ-centered refinement at 1.5 mm, with sliding windows and coarse-mask gating.
• Measured Segmentation Performance: 82.78% DSC / 74.36% NSD on the default-prompt raw-space test evaluation; 93.64% crop proxy DSC / 91.83% crop proxy NSD in architectural validation. These protocols are not directly comparable.

---

📊 Model Comparison

Raw-space test evaluation uses 89 CT scans, with one binary target per text request: spleen, liver, or the union of both kidneys. TotalSegmentator is an automatic multi-structure reference. SegVol, VoxTell and PromptM-UNet use the same default wording, “Mark the {organ} region”, through their respective inference pipelines. Values are mean ± standard deviation; segmentation scores are percentages. NSD uses a 2 mm tolerance.

Default-prompt segmentation

Model	Spleen		Liver		Kidneys
	DSC	NSD	DSC	NSD	DSC	NSD
TotalSegmentator	98.79 ± 1.53	97.93 ± 5.15	96.75 ± 14.73	95.54 ± 14.86	92.44 ± 20.48	92.44 ± 20.30
SegVol	82.81 ± 28.08	72.28 ± 30.03	75.84 ± 29.80	46.51 ± 30.24	57.53 ± 30.20	56.09 ± 30.07
VoxTell	97.63 ± 2.30	95.49 ± 7.62	94.01 ± 14.66	82.60 ± 17.65	90.49 ± 18.78	88.46 ± 20.71
PromptM-UNet	90.77 ± 20.07	85.59 ± 23.32	82.81 ± 30.45	68.47 ± 28.78	74.76 ± 32.37	69.03 ± 32.98

All-case scores include absent targets: two empty masks score 100%, and exactly one empty mask scores 0%. Reference masks are empty in 24 spleen cases and 19 cases each for liver and kidneys. These results measure complete inference in the original CT grid and are separate from the crop-based ablations below.

Overall segmentation and prompt variation

The sweep uses all 54 held-out positive prompts, covering organ name (N), name and shape (NS), name and location (NL), and all three (NSL). Prompts are averaged within patient and organ before aggregation; prompt variants are not independent patients.

Model / variant	Default DSC (%)	Default NSD (%)	Prompt-sweep DSC (%)	Prompt-sweep NSD (%)
TotalSegmentator	95.99 ± 10.78	95.31 ± 10.91	—	—
SegVol	72.06 ± 20.35	58.29 ± 21.73	25.29 ± 22.24	24.23 ± 22.84
VoxTell	94.04 ± 9.11	88.85 ± 11.55	80.49 ± 11.73	75.74 ± 14.46
PromptM-UNet	82.78 ± 19.39	74.36 ± 21.23	82.06 ± 20.13	73.86 ± 21.90
SegVol, punctuation removed	—	—	41.29 ± 19.26	38.56 ± 20.59

The punctuation-removed SegVol run keeps the remaining prompt wording unchanged. PromptM-UNet exceeds VoxTell's mean sweep DSC by 1.57 percentage points, while VoxTell retains higher sweep NSD and higher default-prompt scores. These are descriptive comparisons; statistical superiority is not established. Stability across wording does not by itself demonstrate use of shape or location details.

Inference resources

Model	Jobs	Parameters (M)	Time (s/job)	Peak GPU (GB)	Counted GFLOPs/job
TotalSegmentator	89	492.33	23.538 ± 7.446	1.26 ± 0.53	5589.50 ± 3744.21
SegVol	267	180.91	1.420 ± 0.789	1.17 ± 0.08	2451.02 ± 1378.16
VoxTell	267	4461.80	7.410 ± 6.271	5.63 ± 0.12	61429.42 ± 53637.53
PromptM-UNet	267	27.83	6.219 ± 4.230	3.22 ± 0.54	1808.56 ± 1382.46

GPU memory uses decimal GB, converted from the archived MiB measurements. A TotalSegmentator job predicts multiple structures; a prompted job requests one target. Text encoding is offloaded to CPU, whose memory is excluded. Different timer boundaries and incomplete custom-kernel FLOP counting limit direct cost comparisons.

PromptM-UNet has approximately 6.50× fewer parameters than SegVol and 160.34× fewer than VoxTell. Its measured runtime and GPU memory are lower than VoxTell's, but higher than SegVol's. Efficiency here refers to compactness and measured tradeoffs rather than uniformly best speed, memory or accuracy.

---

🧪 Architectural Ablations

DSC and NSD are macro crop-validation proxies from each run's final Validate FULL output in notebooks/ablation_promptm_unet.ipynb. They use fixed organ-centered 96³ crops and the full validation prompt sweep, excluding anatomy outside the crops and coarse localization. These scores are separate from full-volume test results.

Total parameters are counted from each saved best_model.pth under Ablation/, excluding auxiliary deep-supervision heads, plus the matching frozen text encoder weights stored separately. Counts include the visual backbone, fusion layers, and frozen language encoder. Values are in millions (M); DSC and NSD are percentages.

1. Fusion Operation

CLIP text encoder with all-stage fusion.

Operation	DSC (%)	NSD proxy (%)	Total params (M)
Multiplication	92.82	90.97	68.43
Concatenation	91.66	89.32	68.60
Hybrid	92.50	90.51	76.16

2. Fusion Placement

Multiplicative fusion with the CLIP text encoder.

Placement	DSC (%)	NSD proxy (%)	Total params (M)
Bottleneck only	93.07	90.85	68.31
Last decoder only	92.22	90.02	68.20
Bottleneck and all decoders	92.82	90.97	68.43

3. Text Encoder

Multiplicative fusion at the bottleneck.

Text encoder	DSC (%)	NSD proxy (%)	Total params (M)
CLIP ViT-B/32	93.07	90.85	68.31
Bio_ClinicalBERT	91.86	90.07	113.52
Sentence-BERT / MiniLM	93.24	91.44	27.83

4. Visual Initialization

MiniLM text encoder with multiplicative fusion at the bottleneck.

Initialization	DSC (%)	NSD proxy (%)	Total params (M)
Random	93.24	91.44	27.83
LiTS-pretrained backbone	93.64	91.83	27.83

---

🚀 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 = ["kidney 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 organ mask shape: {binary_mask.shape}")

3. CLI Training & Evaluation

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 100

# 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 Per-Organ Binary Mask
                          (Spleen / Liver / Kidney)

---

📜 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, Liver and Kidney},
  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).