Contrastive learning of global and local features for medical image segmentation with limited annotations (2006.10511v2)

Abstract: A key requirement for the success of supervised deep learning is a large labeled dataset - a condition that is difficult to meet in medical image analysis. Self-supervised learning (SSL) can help in this regard by providing a strategy to pre-train a neural network with unlabeled data, followed by fine-tuning for a downstream task with limited annotations. Contrastive learning, a particular variant of SSL, is a powerful technique for learning image-level representations. In this work, we propose strategies for extending the contrastive learning framework for segmentation of volumetric medical images in the semi-supervised setting with limited annotations, by leveraging domain-specific and problem-specific cues. Specifically, we propose (1) novel contrasting strategies that leverage structural similarity across volumetric medical images (domain-specific cue) and (2) a local version of the contrastive loss to learn distinctive representations of local regions that are useful for per-pixel segmentation (problem-specific cue). We carry out an extensive evaluation on three Magnetic Resonance Imaging (MRI) datasets. In the limited annotation setting, the proposed method yields substantial improvements compared to other self-supervision and semi-supervised learning techniques. When combined with a simple data augmentation technique, the proposed method reaches within 8% of benchmark performance using only two labeled MRI volumes for training, corresponding to only 4% (for ACDC) of the training data used to train the benchmark. The code is made public at https://github.com/krishnabits001/domain_specific_cl.

Contrastive Learning for Medical Image Segmentation with Limited Annotations

The paper "Contrastive Learning of Global and Local Features for Medical Image Segmentation with Limited Annotations" addresses a significant challenge in medical image analysis: the scarcity of labeled datasets required for supervised deep learning methods. The authors propose an innovative approach leveraging self-supervised learning (SSL) and, specifically, contrastive learning to improve segmentation accuracy in volumetric medical images when annotations are limited.

Methodology and Contributions

The paper presents two primary contributions to the contrastive learning framework within SSL:

1. Domain-Specific Global Contrastive Strategies: The authors introduce novel strategies that exploit structural similarities across volumetric medical images. By doing so, they redefine the notion of similarity and dissimilarity among image pairs. Instead of relying solely on image transformations, they consider alignment across volumes to enhance representation learning. This domain-specific strategy yields improved contrastive loss calculations by including partition-based image pair comparisons within MRI volumes.
2. Local Contrastive Loss for Per-Pixel Segmentation: To complement the global approach, the authors propose a local version of the contrastive loss. This loss focuses on learning robust local region representations within images, essential for detailed segmentation tasks. By encouraging similar representations across transformed versions of local regions and differentiating them from other regions within the same images, they better suit the requirements of pixel-wise prediction tasks.

Experimental Evaluation

The proposed methods were validated on three MRI datasets: ACDC, Prostate, and MMWHS. Results demonstrated significant improvements in segmentation accuracy, particularly in scenarios with very limited annotations. Through extensive evaluations, the authors highlight that their approach reaches within approximately 8% of benchmark performance using only two labeled MRI volumes for training in some cases.

Comparative Analysis

In comparison with various pretext-based SSL methods (e.g., rotation, inpainting) and other contemporary techniques like semi-supervised learning and data augmentation, the proposed method consistently achieves higher performance. Notably, the proposed contrastive learning strategy effectively complements other methods, further improving segmentation accuracy when combined.

Theoretical Implications and Practical Applications

Theoretically, this research underscores the potential of integrating domain knowledge into contrastive learning frameworks to fine-tune global and local representations beneficial for complex tasks like segmentation. Practically, the method offers a viable solution for medical image analysis, particularly beneficial in clinical settings where obtaining extensive annotated datasets is often infeasible.

Future Directions

The insights gained from this research pave the way for further exploration of domain-adapted SSL strategies in other medical imaging modalities and potentially different applications within image analysis tasks. Future work could investigate the extension of these strategies to other types of medical imaging data where similar structural or anatomical consistencies are present.

Overall, this paper contributes a substantial advancement in leveraging SSL for medical image segmentation, providing a promising direction for improving model performance under annotation constraints.