[Paper Review] Boosting Semi-supervised Image Segmentation with Global and Local Mutual Information Regularization

Jizong Peng, Marco Pedersoli+1 more|arXiv (Cornell University)|Mar 8, 2021

Advanced Neural Network ApplicationsComputer Science59 references22 citations

TL;DR

This paper proposes a semi-supervised image segmentation method that enhances performance using mutual information (MI) regularization on both global and local feature representations. By maximizing MI between clustered feature embeddings via projection heads and enforcing consistency across network stages, the method achieves near-supervised accuracy with only 5% labeled data on medical imaging benchmarks.

ABSTRACT

The scarcity of labeled data often impedes the application of deep learning to the segmentation of medical images. Semi-supervised learning seeks to overcome this limitation by exploiting unlabeled examples in the learning process. In this paper, we present a novel semi-supervised segmentation method that leverages mutual information (MI) on categorical distributions to achieve both global representation invariance and local smoothness. In this method, we maximize the MI for intermediate feature embeddings that are taken from both the encoder and decoder of a segmentation network. We first propose a global MI loss constraining the encoder to learn an image representation that is invariant to geometric transformations. Instead of resorting to computationally-expensive techniques for estimating the MI on continuous feature embeddings, we use projection heads to map them to a discrete cluster assignment where MI can be computed efficiently. Our method also includes a local MI loss to promote spatial consistency in the feature maps of the decoder and provide a smoother segmentation. Since mutual information does not require a strict ordering of clusters in two different assignments, we incorporate a final consistency regularization loss on the output which helps align the cluster labels throughout the network. We evaluate the method on four challenging publicly-available datasets for medical image segmentation. Experimental results show our method to outperform recently-proposed approaches for semi-supervised segmentation and provide an accuracy near to full supervision while training with very few annotated images.

Motivation & Objective

• To address the challenge of limited annotated medical images in deep learning-based semantic segmentation.
• To improve semi-supervised learning by leveraging unlabeled data through representation invariance and spatial consistency.
• To overcome limitations of existing consistency-based and contrastive learning methods that require large batch sizes or struggle with high-dimensional continuous features.
• To develop a method that effectively regularizes both global and local feature representations using discrete clustering and mutual information.

Proposed method

• Maximizes mutual information (MI) between cluster assignments of transformed views of the same image to enforce global representation invariance in the encoder.
• Uses projection heads to map continuous feature embeddings to discrete cluster assignments, enabling efficient MI estimation without continuous density estimation.
• Introduces a local MI loss that encourages spatial consistency in decoder feature maps by maximizing MI between neighboring spatial features.
• Applies a final consistency regularization loss to align cluster labels across different network stages, ensuring coherent semantic grouping.
• Combines supervised loss with MI-based global and local regularization in a joint optimization framework, avoiding the two-stage pretraining-finetuning setup.
• Employs a symmetric MI estimation approach using entropy-based formulation to avoid trivial clustering solutions.

Experimental results

Research questions

• RQ1 Can mutual information regularization on discrete cluster assignments improve semi-supervised medical image segmentation?
• RQ2 How does combining global and local MI regularization enhance feature representation and segmentation accuracy?
• RQ3 Does joint optimization of supervised and MI-based losses outperform two-stage pretraining-fine-tuning strategies?
• RQ4 Can MI-based methods achieve near-supervised performance with minimal labeled data in medical image segmentation?
• RQ5 How does the proposed method compare to standard consistency-based and contrastive learning approaches in terms of robustness and efficiency?

Key findings

• The proposed method outperforms state-of-the-art semi-supervised methods such as Mean Teacher and Entropy minimization on four public medical image segmentation datasets.
• On the ACDC dataset, the method achieves a mean DSC of 86.7% with only 5% labeled data, approaching the performance of full-supervision.
• The joint optimization strategy improves segmentation accuracy by 5.10% over the two-stage pretraining-fine-tuning baseline, demonstrating the benefit of end-to-end MI regularization with labeled data.
• The method achieves stable and consistent performance across three independent runs, with low standard deviation, indicating robustness to random initialization.
• The ablation study confirms that both global and local MI losses contribute significantly to performance, with the global loss improving invariance and the local loss enhancing spatial smoothness.
• The use of discrete cluster assignments via projection heads enables efficient MI estimation and avoids the computational burden of continuous MI estimation.