Multi-task Learning-Driven Volume and Slice Level Contrastive Learning for 3D Medical Image Classification

Abstract

Automatic 3D medical image classification,e.g., brain tumor grading from 3D MRI images, is important in clinical practice. However, direct tumor grading from 3D MRI images is quite challenging due to the unknown tumor location and relatively small size of abnormal regions. One key point to deal with this problem is to learn more representative and distinctive features. Contrastive learning has shown its effectiveness with representative feature learning in both natural and medical image analysis tasks. However, for 3D medical images, where slices are continuous, simply performing contrastive learning at the volume-level may lead to inferior performance due to the ineffective use of spatial information and distinctive knowledge. To overcome this limitation, we present a novel contrastive learning framework from synergistic 3D and 2D perspectives for 3D medical image classification within a multi-task learning paradigm. We formulate the 3D medical image classification as a Multiple Instance Learning (MIL) problem and introduce an attention-based MIL module to integrate the 2D instance features of each slice into the 3D feature for classification. Then, we simultaneously consider volume-based and slice-based contrastive learning as the auxiliary tasks, aiming to enhance the global distinctive knowledge learning and explore the correspondence relationship among different slice clusters. We conducted experiments on two 3D MRI image classification datasets for brain tumor grading. The results demonstrate that the proposed volume- and slice-level contrastive learning scheme largely boost the main classification task by implicit network regularization during the model optimization, leading to a $10.5\%$ average AUC improvement compared with the basic model on two datasets.

J. Zhu and S. Wang—Equal contribution.