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, while AI-based models automate this process but often falter under limited data and domain shifts. Inspired by how humans learn new tasks through guidance such as how children draw more accurately when shown reference images, we propose Joint Retrieval-Augmented Segmentation (J-RAS), a framework that 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 and cross-dataset generalization. Evaluations on two public benchmarks, ACDC and M&Ms, across multiple backbones (U-Net, TransUNet, SAM, and SegFormer) demonstrate the generality 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.