Uncertainty-Aware Domain Adaptation for ECG Classification

Dawnlicity Charls, Vidhyasaharan Sethu, Beena Ahmed

Abstract: Automated electrocardiogram classification can facilitate cheaper and more accessible arrhythmia and cardiovascular diagnoses for patients. However, clinical machine learning models fail in the real world when the data of interest is collected under different conditions to the training data e.g. due to a different hospital, demographic group, equipment. These differing data collection conditions are known as the target and source domains, respectively. While it is time-consuming and costly to obtain labeled data from the target domain, unlabeled target domain data is more readily available. Unsupervised domain adaptation from a source domain to an unlabeled target domain can unlock more widespread clinical use of machine learning models, including for arrhythmia classification. In this paper, we show that uncertainty objectives can improve arrhythmia classification on new datasets. We apply a Dirichlet Prior Network to test two novel approaches: the first minimizes target domain uncertainty to encourage domain alignment, while the second aligns the class prediction uncertainty of source and target heartbeats. We demonstrate that these methods improve heartbeat classification of electrocardiograms collected at different hospitals using the MIT-BIH and St. Petersburg INCART arrhythmia datasets. Our first approach improves the baseline target F1 score by 7% for ternary arrhythmia classification, while our second approach improves state-of-the-art domain adaptation performance by 3%.