BASIC: Semi-supervised Multi-organ Segmentation with Balanced Subclass Regularization and Semantic-conflict Penalty

2025-01-07

AI Summary

To address severe class imbalance in semi-supervised multi-organ segmentation caused by large inter-organ size variations, this paper proposes a collaborative learning framework that jointly optimizes class balance and semantic consistency. Methodologically: (1) it introduces an auxiliary sub-class segmentation task coupled with sub-class rebalancing regularization to mitigate scale bias; (2) it incorporates a semantic conflict penalty term to explicitly enforce semantic consistency between sub-class predictions and parent-class ground-truth labels; (3) it unifies multi-task learning, consistency regularization, and semantic constraints within a Mean Teacher architecture. Evaluated on the WORD and FLARE 2022 benchmarks, our method significantly improves segmentation accuracy for small organs, achieving an average mDice gain of +3.2% over state-of-the-art methods—demonstrating its effectiveness in modeling and alleviating class imbalance.

Technology Category

• Machine Learning: Semi-Supervised Learning
• Computer Vision: Segmentation
• Search and Optimization: Learning to Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs

Abstract

Semi-supervised learning (SSL) has shown notable potential in relieving the heavy demand of dense prediction tasks on large-scale well-annotated datasets, especially for the challenging multi-organ segmentation (MoS). However, the prevailing class-imbalance problem in MoS caused by the substantial variations in organ size exacerbates the learning difficulty of the SSL network. To address this issue, in this paper, we propose an innovative semi-supervised network with BAlanced Subclass regularIzation and semantic-Conflict penalty mechanism (BASIC) to effectively learn the unbiased knowledge for semi-supervised MoS. Concretely, we construct a novel auxiliary subclass segmentation (SCS) task based on priorly generated balanced subclasses, thus deeply excavating the unbiased information for the main MoS task with the fashion of multi-task learning. Additionally, based on a mean teacher framework, we elaborately design a balanced subclass regularization to utilize the teacher predictions of SCS task to supervise the student predictions of MoS task, thus effectively transferring unbiased knowledge to the MoS subnetwork and alleviating the influence of the class-imbalance problem. Considering the similar semantic information inside the subclasses and their corresponding original classes (i.e., parent classes), we devise a semantic-conflict penalty mechanism to give heavier punishments to the conflicting SCS predictions with wrong parent classes and provide a more accurate constraint to the MoS predictions. Extensive experiments conducted on two publicly available datasets, i.e., the WORD dataset and the MICCAI FLARE 2022 dataset, have verified the superior performance of our proposed BASIC compared to other state-of-the-art methods.

Problem

• Semi-Supervised Learning
• Multi-organ Segmentation
• Class Imbalance

Innovation

• BASIC Network
• Auxiliary Task
• Balanced Subclass Regularization