Multi-Prototype Embedding Refinement for Semi-Supervised Medical Image Segmentation
Yali Bi, Enyu Che, Yinan Chen, Yuanpeng He, Jingwei Qu
TL;DR
Addresses intra-class voxel variance and uncertainty in semi-supervised medical image segmentation. Proposes MPER, a multi-prototype embedding refinement framework that models intra-class structure by learning multiple prototypes per class. Integrates a consistency constraint between linear and prototype-based classifiers and adds a contrastive loss, with training conducted in a two-stage Mean Teacher framework. Demonstrates state-of-the-art performance on LA and ACDC benchmarks, especially when labeled data are scarce, highlighting the method's robustness and potential clinical utility.
Abstract
Medical image segmentation aims to identify anatomical structures at the voxel-level. Segmentation accuracy relies on distinguishing voxel differences. Compared to advancements achieved in studies of the inter-class variance, the intra-class variance receives less attention. Moreover, traditional linear classifiers, limited by a single learnable weight per class, struggle to capture this finer distinction. To address the above challenges, we propose a Multi-Prototype-based Embedding Refinement method for semi-supervised medical image segmentation. Specifically, we design a multi-prototype-based classification strategy, rethinking the segmentation from the perspective of structural relationships between voxel embeddings. The intra-class variations are explored by clustering voxels along the distribution of multiple prototypes in each class. Next, we introduce a consistency constraint to alleviate the limitation of linear classifiers. This constraint integrates different classification granularities from a linear classifier and the proposed prototype-based classifier. In the thorough evaluation on two popular benchmarks, our method achieves superior performance compared with state-of-the-art methods. Code is available at https://github.com/Briley-byl123/MPER.