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

TL;DR

This paper introduces UG-CEMT, a novel semi-supervised medical image segmentation framework that combines co-training, knowledge distillation, and uncertainty-guided regularization to achieve state-of-the-art results with limited labeled data.

Contribution

The work presents a new framework integrating cross-attention ensemble, uncertainty guidance, and vision transformer concepts for improved semi-supervised segmentation.

Findings

Achieves state-of-the-art results on prostate and cardiac MRI datasets.

Performs well with only 10% labeled data, nearing fully supervised performance.

Outperforms existing methods like Mean Teacher and Cross-pseudo Supervision.

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 \url{https://github.com/Meghnak13/UG-CEMT}