Exploiting Unlabeled Structures through Task Consistency Training for Versatile Medical Image Segmentation

2025-09-04

AI Summary

In medical image multi-class segmentation, partially labeled datasets (PLDs) induce severe class imbalance, and existing pseudo-labeling methods rely on auxiliary models and are vulnerable to label noise. Method: We propose a consistency-based training framework that requires no auxiliary model: it enforces prediction consistency between a primary segmentation head and multiple auxiliary task heads to collaboratively mine structural information from unlabeled regions; further, we introduce a dynamic filtering strategy and a unified auxiliary uncertainty-weighted loss (UAUWL) to suppress noise propagation and mitigate task-dominance bias. Contribution/Results: Evaluated on an abdominal multi-organ dataset spanning eight clinical centers, our method significantly alleviates class imbalance and achieves state-of-the-art performance in multi-class segmentation, outperforming existing semi-supervised and weakly supervised approaches.

Technology Category

• Machine Learning: Multi-class/Multi-label Learning & Extreme Classification
• Computer Vision: Segmentation
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Responsible Web: Consent frameworks and practices on the web
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Versatile medical image segmentation (VMIS) targets the segmentation of multiple classes, while obtaining full annotations for all classes is often impractical due to the time and labor required. Leveraging partially labeled datasets (PLDs) presents a promising alternative; however, current VMIS approaches face significant class imbalance due to the unequal category distribution in PLDs. Existing methods attempt to address this by generating pseudo-full labels. Nevertheless, these typically require additional models and often result in potential performance degradation from label noise. In this work, we introduce a Task Consistency Training (TCT) framework to address class imbalance without requiring extra models. TCT includes a backbone network with a main segmentation head (MSH) for multi-channel predictions and multiple auxiliary task heads (ATHs) for task-specific predictions. By enforcing a consistency constraint between the MSH and ATH predictions, TCT effectively utilizes unlabeled anatomical structures. To avoid error propagation from low-consistency, potentially noisy data, we propose a filtering strategy to exclude such data. Additionally, we introduce a unified auxiliary uncertainty-weighted loss (UAUWL) to mitigate segmentation quality declines caused by the dominance of specific tasks. Extensive experiments on eight abdominal datasets from diverse clinical sites demonstrate our approach's effectiveness.

Problem

• Addressing class imbalance in versatile medical image segmentation
• Leveraging partially labeled datasets without requiring extra models
• Mitigating performance degradation from label noise and task dominance

Innovation

• Task Consistency Training framework without extra models
• Consistency constraint between main and auxiliary heads
• Unified auxiliary uncertainty-weighted loss for segmentation