[Paper Review] Correlation-Aware Mutual Learning for Semi-supervised Medical Image Segmentation

Shengbo Gao, Ziji Zhang+3 more|arXiv (Cornell University)|Jul 12, 2023

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes Correlation-Aware Mutual Learning (CAML), a novel semi-supervised medical image segmentation framework that enhances model performance by explicitly leveraging labeled data to guide feature learning from unlabeled data. CAML employs a co-teaching scheme with two modules—Cross-sample Mutual Attention (CMA) for dense cross-sample attention and Omni-Correlation Consistency (OCC) for consistency regularization across labeled and unlabeled features—achieving state-of-the-art performance, especially under low-label settings, with a Dice score of 87.34% at 5% labeled data and approaching fully supervised performance at 20% labeled data.

ABSTRACT

Semi-supervised learning has become increasingly popular in medical image segmentation due to its ability to leverage large amounts of unlabeled data to extract additional information. However, most existing semi-supervised segmentation methods only focus on extracting information from unlabeled data, disregarding the potential of labeled data to further improve the performance of the model. In this paper, we propose a novel Correlation Aware Mutual Learning (CAML) framework that leverages labeled data to guide the extraction of information from unlabeled data. Our approach is based on a mutual learning strategy that incorporates two modules: the Cross-sample Mutual Attention Module (CMA) and the Omni-Correlation Consistency Module (OCC). The CMA module establishes dense cross-sample correlations among a group of samples, enabling the transfer of label prior knowledge to unlabeled data. The OCC module constructs omni-correlations between the unlabeled and labeled datasets and regularizes dual models by constraining the omni-correlation matrix of each sub-model to be consistent. Experiments on the Atrial Segmentation Challenge dataset demonstrate that our proposed approach outperforms state-of-the-art methods, highlighting the effectiveness of our framework in medical image segmentation tasks. The codes, pre-trained weights, and data are publicly available.

Motivation & Objective

• To address the underutilization of labeled data in existing semi-supervised medical image segmentation methods, which typically treat it only as an initial supervision signal.
• To improve model generalization by explicitly modeling cross-sample correlations between labeled and unlabeled data to transfer prior knowledge.
• To develop a mutual learning framework that enhances feature representation through joint optimization of labeled and unlabeled data.
• To achieve robust and consistent performance even when only a small number of labeled images are available.
• To reduce the reliance on large-scale annotated datasets in clinical applications by maximizing information extraction from limited labeled data.

Proposed method

• CAML employs a co-teaching framework with two student networks trained via mutual pseudo-supervision to encourage consistency and knowledge distillation.
• The Cross-sample Mutual Attention (CMA) module establishes dense, bidirectional attention among a group of samples, enabling co-salient feature learning across labeled and unlabeled data.
• The Omni-Correlation Consistency (OCC) module constructs a memory bank of labeled features and computes omni-correlations between unlabeled features and this reference set to model similarity at the feature level.
• OCC enforces consistency in the omni-correlation matrices of both sub-models, regularizing the learning process and promoting a unified feature distribution in unlabeled data.
• The framework integrates both modules within a mutual learning loop, synchronously updating both networks while propagating knowledge from labeled to unlabeled samples.
• The method is end-to-end trainable and does not incur additional inference cost, making it suitable for clinical deployment.

Experimental results

Research questions

• RQ1 Can labeled data be effectively leveraged to improve representation learning from unlabeled medical images in semi-supervised segmentation?
• RQ2 How can cross-sample correlations between labeled and unlabeled data be modeled to enhance feature transfer and consistency?
• RQ3 Does introducing mutual attention across multiple samples improve segmentation performance compared to standard consistency regularization?
• RQ4 Can omni-correlation consistency between labeled and unlabeled features lead to more robust and generalizable feature learning?
• RQ5 How does the proposed CAML framework perform under extreme label scarcity, such as 5% or 10% labeled data?

Key findings

• At 5% labeled data, CAML achieves a Dice score of 87.34% with a standard deviation of 0.05, significantly outperforming the baseline and demonstrating high stability.
• With 10% labeled data, CAML achieves a Dice score of 89.62%, approaching the performance of a fully supervised model (upper bound: 90.98%).
• At 20% labeled data, CAML achieves a Dice score of 90.78%, showing minimal gap from the fully supervised upper bound.
• The ablation study confirms that both CMA and OCC modules are essential: their combination yields a 6.42% absolute improvement in Dice score at 5% labeled data.
• The method maintains low standard deviation across runs, indicating high robustness and generalization, especially under low-label regimes.
• Qualitative results show that CAML produces more precise and complete segmentation masks than all compared methods, particularly in boundary and small structure regions.