J-RAS: Mutual Adaptation for Medical Image Segmentation via Contrastive Retrieval-Augmented Joint Optimization

Abstract

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.