Joint Self-Supervised Image-Volume Representation Learning with Intra-inter Contrastive Clustering

Authors

  • Duy M. H. Nguyen German Research Centre for Artificial Intelligence University of Stuttgart
  • Hoang Nguyen University of Science, VNU-HCMUS
  • Truong T. N. Mai Dongguk University
  • Tri Cao University of Science, VNU-HCMUS
  • Binh T. Nguyen University of Science, VNU-HCMUS
  • Nhat Ho University of Texas at Austin
  • Paul Swoboda Max Planck Institute for Informatics
  • Shadi Albarqouni Helmholtz AI, Helmholtz Munich University of Bonn
  • Pengtao Xie University of California San Diego
  • Daniel Sonntag German Research Center for Artificial Intelligence Oldenburg University

DOI:

https://doi.org/10.1609/aaai.v37i12.26687

Keywords:

General

Abstract

Collecting large-scale medical datasets with fully annotated samples for training of deep networks is prohibitively expensive, especially for 3D volume data. Recent breakthroughs in self-supervised learning (SSL) offer the ability to overcome the lack of labeled training samples by learning feature representations from unlabeled data. However, most current SSL techniques in the medical field have been designed for either 2D images or 3D volumes. In practice, this restricts the capability to fully leverage unlabeled data from numerous sources, which may include both 2D and 3D data. Additionally, the use of these pre-trained networks is constrained to downstream tasks with compatible data dimensions. In this paper, we propose a novel framework for unsupervised joint learning on 2D and 3D data modalities. Given a set of 2D images or 2D slices extracted from 3D volumes, we construct an SSL task based on a 2D contrastive clustering problem for distinct classes. The 3D volumes are exploited by computing vectored embedding at each slice and then assembling a holistic feature through deformable self-attention mechanisms in Transformer, allowing incorporating long-range dependencies between slices inside 3D volumes. These holistic features are further utilized to define a novel 3D clustering agreement-based SSL task and masking embedding prediction inspired by pre-trained language models. Experiments on downstream tasks, such as 3D brain segmentation, lung nodule detection, 3D heart structures segmentation, and abnormal chest X-ray detection, demonstrate the effectiveness of our joint 2D and 3D SSL approach. We improve plain 2D Deep-ClusterV2 and SwAV by a significant margin and also surpass various modern 2D and 3D SSL approaches.

Downloads

Published

2023-06-26

How to Cite

Nguyen, D. M. H., Nguyen, H., Mai, T. T. N., Cao, T., Nguyen, B. T., Ho, N., Swoboda, P., Albarqouni, S., Xie, P., & Sonntag, D. (2023). Joint Self-Supervised Image-Volume Representation Learning with Intra-inter Contrastive Clustering. Proceedings of the AAAI Conference on Artificial Intelligence, 37(12), 14426-14435. https://doi.org/10.1609/aaai.v37i12.26687

Issue

Section

AAAI Special Track on AI for Social Impact