Contrastive and Selective Hidden Embeddings for Medical Image Segmentation

TL;DR

This paper introduces contrastive and selective hidden embedding techniques for medical image segmentation, leveraging unlabeled data and uncertainty-aware feature selection to improve performance and reduce training data requirements.

Contribution

It proposes novel modules, PDCR and UAFS, integrated into existing architectures, to enhance segmentation accuracy and efficiency using contrastive learning and feature selection.

Findings

Achieves state-of-the-art results on 8 datasets from 6 domains.

Reduces training data needs to a quarter while maintaining performance.

Enhances segmentation with label-contained information extraction.

Abstract

Medical image segmentation has been widely recognized as a pivot procedure for clinical diagnosis, analysis, and treatment planning. However, the laborious and expensive annotation process lags down the speed of further advances. Contrastive learning-based weight pre-training provides an alternative by leveraging unlabeled data to learn a good representation. In this paper, we investigate how contrastive learning benefits the general supervised medical segmentation tasks. To this end, patch-dragsaw contrastive regularization (PDCR) is proposed to perform patch-level tugging and repulsing with the extent controlled by a continuous affinity score. And a new structure dubbed uncertainty-aware feature selection block (UAFS) is designed to perform the feature selection process, which can handle the learning target shift caused by minority features with high uncertainty. By plugging the proposed 2 modules into the existing segmentation architecture, we achieve state-of-the-art results across 8 public datasets from 6 domains. Newly designed modules further decrease the amount of training data to a quarter while achieving comparable, if not better, performances. From this perspective, we take the opposite direction of the original self/un-supervised contrastive learning by further excavating information contained within the label.