J-RAS: Mutual Adaptation for Medical Image Segmentation via Contrastive Retrieval-Augmented Joint Optimization
About
Manual medical image segmentation by clinicians, though accurate, is time-consuming and variable across experts, whereas AI-based models automate this process but often underperform with limited data and domain shifts. Inspired by how pathology trainees acquire disease recognition skills through guided comparison with expert-annotated slides and histopathology atlas reference images, we propose Joint Retrieval-Augmented Segmentation (J-RAS). This framework enables segmentation networks to learn with guidance. J-RAS jointly optimizes a segmentation model and a retrieval model through alternating contrastive and supervised learning, allowing the retrieval network to discover contextually relevant image-mask pairs that refine the segmentation model's anatomical reasoning. Unlike conventional retrieval-based augmentation that passively provides similar samples, J-RAS establishes a mutual adaptation and optimization loop where the retrieval model learns to emphasize segmentation-relevant cues, while the segmentation model leverages retrieved examples to improve boundary delineation, robustness to rare cases, and cross-dataset generalization. Evaluations on four public benchmarks spanning different imaging modalities, including ACDC and M&Ms (MRI), Breast Cancer Ultrasound, and lung and infection CT, across multiple backbones (U-Net, TransUNet, SAM, and SegFormer) demonstrate the generalizability and effectiveness of J-RAS. For instance, on ACDC, SegFormer improves from a mean Dice of 0.8708$\pm$0.042 and HD of 1.8130$\pm$2.49 to 0.9115$\pm$0.031 and 1.1489$\pm$0.30. These results highlight how retrieval-guided contrastive optimization bridges human-like guidance and machine-learned precision in medical image segmentation.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Medical Image Segmentation | ACDC (test) | Avg DSC91.18 | 180 | |
| Segmentation | ACDC (test) | Dice93.95 | 28 | |
| Medical Image Segmentation | M&Ms (test) | Average DSC88.47 | 11 | |
| Segmentation | COVID-19 (test) | Dice Score94.44 | 6 | |
| Segmentation | Breast Cancer (test) | Dice Coefficient79.06 | 2 |