MedContext: Learning Contextual Cues for Efficient Volumetric Medical Segmentation

TL;DR

MedContext tackles data scarcity in 3D medical image segmentation by jointly optimizing supervised voxel-wise segmentation with self-supervised masked-input reconstruction in a single training stage. It uses a architecture-agnostic student-teacher framework where the masked input guides reconstruction in segmentation space, aided by a moving-average teacher to provide soft targets, and an EMA-based optimization to avoid mode collapse. Across Synapse, ACDC, and BraTS, MedContext yields consistent Dice and boundary improvements for multiple architectures, with notable gains in few-shot scenarios and when compared to pretraining-based baselines. The approach is plug-and-play, data-efficient, and demonstrates the practical impact of embedding contextual cues into 3D medical segmentation models without large-scale external data.

Abstract

Volumetric medical segmentation is a critical component of 3D medical image analysis that delineates different semantic regions. Deep neural networks have significantly improved volumetric medical segmentation, but they generally require large-scale annotated data to achieve better performance, which can be expensive and prohibitive to obtain. To address this limitation, existing works typically perform transfer learning or design dedicated pretraining-finetuning stages to learn representative features. However, the mismatch between the source and target domain can make it challenging to learn optimal representation for volumetric data, while the multi-stage training demands higher compute as well as careful selection of stage-specific design choices. In contrast, we propose a universal training framework called MedContext that is architecture-agnostic and can be incorporated into any existing training framework for 3D medical segmentation. Our approach effectively learns self supervised contextual cues jointly with the supervised voxel segmentation task without requiring large-scale annotated volumetric medical data or dedicated pretraining-finetuning stages. The proposed approach induces contextual knowledge in the network by learning to reconstruct the missing organ or parts of an organ in the output segmentation space. The effectiveness of MedContext is validated across multiple 3D medical datasets and four state-of-the-art model architectures. Our approach demonstrates consistent gains in segmentation performance across datasets and different architectures even in few-shot data scenarios. Our code and pretrained models are available at https://github.com/hananshafi/MedContext

Figures (6)

• Figure 1: Comparison, in term of Dice scores (%), when integrating our approach into UNETR UNETR, SwinUNETR SWIN_UNETR and nnFormer nnFormer for medical segmentation on Synapse dataset (Sec. \ref{['sec:experiments']}) using conventional setting (Left) and few-shot setting (5 samples only, Right). Without any modification to the model architecture or its training pipeline, our proposed universal approach complements the supervised voxel-wise segmentation and enhances the performance of state-of-the-art architectures.
• Figure 2: Qualitative Comparison between the baseline nnFormer nnFormer and our proposed MedContext integrated with nnFormer. The examples display different abdominal organs (Synapse) (Left) and regions of the heart (ACDC) (Right), with their corresponding labels in the legend below. The baseline nnFormer struggles to accurately segment the organs and heart regions. In certain cases, it gives false segmentation results highlighted in red boxes. Best viewed zoomed in. Refer to supplementary material for additional qualitative comparisons.
• Figure 3: Overview of our MedContext approach: The original 3D volume is masked and fed to the student model (top-row) along with the original input. The teacher model (bottom-row) is only fed with the original volume. The difference between the semantic voxelwise predictions for the masked and original inputs corresponding to the student and teacher networks respectively is minimized to guide the reconstruction of masked regions in the output segmentation space. Our approach induces contextual consistency by enabling the model to reconstruct and segment the missing organs or organ parts and therefore yields more precise and accurate segmentation results.
• Figure 4: DSC (%) on Synapse across different models. Left: Distillation from Teacher. We demonstrate the importance of knowledge distillation through teacher for effectively leveraging contextual cues. Right: Student vs Teacher. We show that utilizing student weights during inference benefits overall performance.
• Figure 5: Qualitative comparison on multi-organ synapse dataset: We showcase the benefit of our MedContext framework implemented on the UNETR architecture. The examples display various abdominal organs, with their corresponding labels in the legend below. The existing baseline method struggles to accurately segment the organs as can be seen from the red boxes. Best viewed in zoom.