Unsupervised Feature Clustering Improves Contrastive Representation Learning for Medical Image Segmentation

TL;DR

This paper introduces an unsupervised feature clustering approach to enhance contrastive learning for medical image segmentation, improving feature representation by better selecting positive and negative samples without relying on additional data or assumptions.

Contribution

The novel method uses hierarchical clustering of features to select positives and negatives, addressing limitations of existing contrastive learning approaches in medical imaging.

Findings

Outperforms state-of-the-art contrastive methods on skin dermoscopic image segmentation.

Achieves superior results on 3D multi-class whole heart CT segmentation.

Demonstrates improved feature representations for medical image analysis.

Abstract

Self-supervised instance discrimination is an effective contrastive pretext task to learn feature representations and address limited medical image annotations. The idea is to make features of transformed versions of the same images similar while forcing all other augmented images' representations to contrast. However, this instance-based contrastive learning leaves performance on the table by failing to maximize feature affinity between images with similar content while counter-productively pushing their representations apart. Recent improvements on this paradigm (e.g., leveraging multi-modal data, different images in longitudinal studies, spatial correspondences) either relied on additional views or made stringent assumptions about data properties, which can sacrifice generalizability and applicability. To address this challenge, we propose a new self-supervised contrastive learning method that uses unsupervised feature clustering to better select positive and negative image samples. More specifically, we produce pseudo-classes by hierarchically clustering features obtained by an auto-encoder in an unsupervised manner, and prevent destructive interference during contrastive learning by avoiding the selection of negatives from the same pseudo-class. Experiments on 2D skin dermoscopic image segmentation and 3D multi-class whole heart CT segmentation demonstrate that our method outperforms state-of-the-art self-supervised contrastive techniques on these tasks.