Multi-Label ECG Classification by Exploiting Prior Label Correlations

Yang Liu1, Qince Li1, Runnan He1, Yacong Li1, Henggui Zhang2
1Harbin Institute of Technology, 2University of Manchester


Abstract

Automatic ECG classification technology is in demand to cope with the rapid growth of ECG monitoring data. Considering the diversity of cardiac conditions and the complex correlations between them, the ECG classification is inherent a multi-label classification problem. This problem is challenging not only because of the huge number of possible label combinations, but also because of the imbalance between different conditions. Besides, the different configurations of ECG leads also need to be taken into account. In this work, we propose a novel loss function named Class-Correlation-Weighted-Loss (CCWL) to incorporate prior correlations between different classes into the loss calculation for the training of a deep neural network (DNN) which combines a residual convolutional network with an attention mechanism. Specifically, CCWL tackles the problems of label correlations and imbalance by assigning specific weights to the predictive loss for each class on each sample according to the correlations between the class and the reference labels. In addition, CCWL can also improve the robustness of the model against annotation errors, e.g., missing labels. The scores of our model are 0.525 (12 leads), 0.496 (6 leads), 0.504 (3 leads) and 0.477 (2 leads) in local 5-fold cross validations on the training set of 2021 PhysioNet/CinC Challenge. However, our online scored entry (team name: HIT-CS) is without the CCWL, and its online scores are 0.450 (12 leads), -0.062, (6 leads) 0.398 (3 leads) and 0.454 (2 leads), and its local 5-fold cross validation scores are 0.405 (12 leads), 0.359 (6 leads), 0.367 (3 leads) and 0.332 (2 leads). Therefore, CCWL is effective to improve the model’s performance for multi-label ECG classification under various lead configurations.