Improved performance and robustness of multi-task representation learning with consistency loss between pretexts for intracranial hemorrhage identification in head CT.

Researchers

Journal

Modalities

Models

Abstract

With the recent development of deep learning, the classification and segmentation tasks of computer-aided diagnosis (CAD) using non-contrast head computed tomography (NCCT) for intracranial hemorrhage (ICH) has become popular in emergency medical care. However, a few challenges remain, such as the difficulty of training due to the heterogeneity of ICH, the requirement for high performance in both sensitivity and specificity, patient-level predictions demanding excessive costs, and vulnerability to real-world external data. In this study, we proposed a supervised multi-task aiding representation transfer learning network (SMART-Net) for ICH to overcome these challenges. The proposed framework consists of upstream and downstream components. In the upstream, a weight-shared encoder of the model is trained as a robust feature extractor that captures global features by performing slice-level multi-pretext tasks (classification, segmentation, and reconstruction). Adding a consistency loss to regularize discrepancies between classification and segmentation heads has significantly improved representation and transferability. In the downstream, the transfer learning was conducted with a pre-trained encoder and 3D operator (classifier or segmenter) for volume-level tasks. Excessive ablation studies were conducted and the SMART-Net was developed with optimal multi-pretext task combinations and a 3D operator. Experimental results based on four test sets (one internal and two external test sets that reflect a natural incidence of ICH, and one public test set with a relatively small amount of ICH cases) indicate that SMART-Net has better robustness and performance in terms of volume-level ICH classification and segmentation over previous methods. All code is available at https://github.com/babbu3682/SMART-Net.Copyright © 2022 Elsevier B.V. All rights reserved.