Robust and Task-Aware Training of Deep Residual Networks for Varying-Lead ECG Classification

Hansheng Ren, Miao Xiong, Bryan Hooi
National University of Singapore


Abstract

For the task of ECG classification using varying-lead signals (12, 6, 3, or 2 leads), we develop an end-to-end deep learning model based on the ResNet model, accompanied by a task-aware loss function as well as improvements which aim to train a more robust model that generalizes better to different data distributions. To fully utilize the given training information, we first preprocess the data by splitting long samples into multiple short samples - at test time, we aggregate over these samples using majority vote to derive the final predictions. Our task-aware loss function aims to train a classifier in a way that aligns better with clinical realities as captured in the scoring metric - we do this by defining a continuous relaxation of the scoring metric and optimizing the resulting loss function directly in our end-to-end framework. Finally, we use additional modifications to improve the robustness of the trained model to allow it to generalize better to unseen distributions, including mixup training (which tends to induce smoother decision boundaries) as well as further tuning of decision thresholds using a data splitting approach.

Results: On the training set, our method achieves cross validation scores of 0.109, 0.156, 0.115, 0.104 for 12, 6, 3, and 2 leads respectively. On the validation set, our method achieves scores of 0.004, 0.010, 0.003 and 0.000 respectively, under the team name of DataLA_NUS.