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

• Published in: Computational Mathematics Modeling in Cancer Analysis Vol. 13574; pp. 110 - 120
• Main Authors: Zhu, Jiayuan, Wang, Shujun, He, Jinzheng, Schönlieb, Carola-Bibiane, Yu, Lequan
• Format: Book Chapter
• Language: English
• Published: Switzerland Springer 2022; Springer Nature Switzerland
• Series: Lecture Notes in Computer Science
• Subjects: 3D MRI image classification; Brain tumor grading; Contrastive learning; Multi-task learning
• ISBN: 9783031172656; 3031172655
• ISSN: 0302-9743; 1611-3349
• DOI: 10.1007/978-3-031-17266-3_11

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% $$10.5\%$$ average AUC improvement compared with the basic model on two datasets.