# Align then Refine: Text-Guided 3D Prostate Lesion Segmentation

arXiv: https://arxiv.org/abs/2604.18713 · HTML: https://mlsanity.com/paper/2604.18713/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 10 of its 18 references that are in the map.
- First posted on arXiv 20 Apr 2026.

## Builds on (most cited first, 10 of 10)

- [3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation](https://mlsanity.com/paper/1606.06650/) (2016)
- [Attention U-Net: Learning Where to Look for the Pancreas](https://mlsanity.com/paper/1804.03999/) (2018)
- [UNETR: Transformers for 3D Medical Image Segmentation](https://mlsanity.com/paper/2103.10504/) (2021)
- [Swin UNETR: Swin Transformers for Semantic Segmentation of Brain Tumors in MRI Images](https://mlsanity.com/paper/2201.01266/) (2022)
- [MedCLIP: Contrastive Learning from Unpaired Medical Images and Text](https://mlsanity.com/paper/2210.10163/) (2022)
- [BiomedCLIP: a multimodal biomedical foundation model pretrained from fifteen million scientific image-text pairs](https://mlsanity.com/paper/2303.00915/) (2023)
- [Medical SAM Adapter: Adapting Segment Anything Model for Medical Image Segmentation](https://mlsanity.com/paper/2304.12620/) (2023)
- [3D MRI brain tumor segmentation using autoencoder regularization](https://mlsanity.com/paper/1810.11654/) (2018)
- [SAM-Med3D: Towards General-purpose Segmentation Models for Volumetric Medical Images](https://mlsanity.com/paper/2310.15161/) (2023)
- [MedSAM2: Segment Anything in 3D Medical Images and Videos](https://mlsanity.com/paper/2504.03600/) (2025)

## Abstract (arXiv)

Automated 3D segmentation of prostate lesions from biparametric MRI (bp-MRI) is essential for reliable algorithmic analysis, but achieving high precision remains challenging. Volumetric methods must combine multiple modalities while ensuring anatomical consistency, but current models struggle to integrate cross-modal information reliably. While vision-language models (VLMs) are replacing the currently used architectural designs, they still lack the fine-grained, lesion-level semantics required for effective localized guidance. To address these limitations, we propose a new multi-encoder U-Net architecture incorporating three key innovations: (1) an alignment loss that enhances foreground text-image similarity to inject lesion semantics; (2) a heatmap loss that calibrates the similarity map and suppresses spurious background activations; and (3) a final-stage, confidence-gated multi-head cross-attention refiner that performs localized boundary edits in high-confidence regions. A phase-scheduled training regime stabilizes the optimization of these components. Our method consistently outperforms prior approaches, establishing a new state-of-the-art on the PI-CAI dataset through enhanced multi-modal fusion and localized text guidance. Our code is available at https://github.com/NUBagciLab/Prostate-Lesion-Segmentation.