Efficient Prototype Consistency Learning in Medical Image Segmentation via Joint Uncertainty and Data Augmentation

TL;DR

This work tackles limited annotations in 3D medical image segmentation by proposing EPCL-JUDA, which unifies memory-efficient prototype learning with joint uncertainty quantification in a Mean-Teacher framework. It enriches prototype expressiveness through CutMix-based augmentation of labeled data, separates and stabilizes unlabeled data via JUQ, and fuses labeled and unlabeled prototypes into high-quality global prototypes that drive prototype-to-feature consistency. A prototype network reduces memory demands while maintaining or boosting performance, and a Temporal Ensembling scheme further refines prototypes over training. Across LA, Pancreas-NIH, and TBAD benchmarks, EPCL-JUDA achieves state-of-the-art results, especially under scarce-label regimes, demonstrating robust, memory-efficient semi-supervised performance for complex 3D segmentation tasks.

Abstract

Recently, prototype learning has emerged in semi-supervised medical image segmentation and achieved remarkable performance. However, the scarcity of labeled data limits the expressiveness of prototypes in previous methods, potentially hindering the complete representation of prototypes for class embedding. To overcome this issue, we propose an efficient prototype consistency learning via joint uncertainty quantification and data augmentation (EPCL-JUDA) to enhance the semantic expression of prototypes based on the framework of Mean-Teacher. The concatenation of original and augmented labeled data is fed into student network to generate expressive prototypes. Then, a joint uncertainty quantification method is devised to optimize pseudo-labels and generate reliable prototypes for original and augmented unlabeled data separately. High-quality global prototypes for each class are formed by fusing labeled and unlabeled prototypes, which are utilized to generate prototype-to-features to conduct consistency learning. Notably, a prototype network is proposed to reduce high memory requirements brought by the introduction of augmented data. Extensive experiments on Left Atrium, Pancreas-NIH, Type B Aortic Dissection datasets demonstrate EPCL-JUDA's superiority over previous state-of-the-art approaches, confirming the effectiveness of our framework. The code will be released soon.

Figures (3)

• Figure 1: The overview of the proposed EPCL-JUDA framework. For training process, labeled data is augmented using techniques like CutMix and combined with original data, fed into the teacher network to generate expressive prototypes $p^c_l$. And, a prototype network is introduced to generate prototype features, which effectively reduces memory consumption and improves segmentation performance. Additionally, a Joint uncertainty quantification, combining entropy, variance, and mean, is employed to optimize pseudo-labels and generate reliable prototypes $p^c_{u1}$ and $p^c_{u2}$ for original and augmented unlabeled data. These prototypes are fused to form high-quality global prototypes $p^c$, which are used to generate multiple prototype-to-feature similarity maps which are utilized to conduct consistency learning with labels and pseudo-labels, with different prototype features.
• Figure 2: The visualizations of experimental results on Left Atrium dataset. EPCL-JUDA denotes the proposed mixed prototypes consistency learning. GT denotes the ground truth labels. Baseline represents the method of URPC. ITK-SNAP py06nimg is the tool to visualize 3D medical image.
• Figure 3: The reliability map generated by Entropy and JUQ