Uncertainty-Guided Cross Attention Ensemble Mean Teacher for Semi-Supervised Medical Image Segmentation

Abstract

This work proposes a novel framework Uncertainty-Guided Cross Attention Ensemble Mean Teacher (UG-CEMT) for achieving state-of-the-art performance in semi-supervised medical image segmentation. UG-CEMT leverages the strengths of co-training and knowledge distillation by combining a Cross-attention Ensemble Mean Teacher framework (CEMT) inspired by Vision Transformers (ViT) with uncertainty-guided consistency regularization and Sharpness-Aware Minimization emphasizing uncertainty. UG-CEMT improves semi-supervised performance while maintaining a consistent network architecture and task setting by fostering high disparity between sub-networks. Experiments demonstrate significant advantages over existing methods like Mean Teacher and Cross-pseudo Supervision in terms of disparity domain generalization and medical image segmentation performance. UG-CEMT achieves state-of-the-art results on multi-center prostate MRI and cardiac MRI datasets where object segmentation is particularly challenging. Our results show that using only 10% labeled data UG-CEMT approaches the performance of fully supervised methods demonstrating its effectiveness in exploiting unlabeled data for robust medical image segmentation. The code is publicly available at https://github.com/Meghnak13/UG-CEMT