Data Augmentation via Wasserstein Geodesic Perturbation for Robust Electrocardiogram Prediction

There has been an increased interest in applying deep neural networks to automatically interpret and analyze the 12-lead electrocardiogram (ECG). However, the imbalance and heterogeneity of real-world datasets place obstacles to the efficient training of neural networks. Moreover, deep learning classifiers could be vulnerable to adversarial examples and perturbations and could lead to catastrophic outcomes for clinical trials and insurance claims.In this paper, we propose a physiologically-inspired data augmentation to improve the performance, generalization, and to increase the robustness of ECG prediction models. We obtain augmented samples by perturbing the data distribution towards other classes along the geodesic in Wasserstein space. To better utilize the domain knowledge, we design a ground metric that recognizes the difference between ECG signals based on physiological features. Learning from 12-lead ECG signals, our model is able to distinguish five categories of cardiac conditions. Our results demonstrate improvements in accuracy and robustness reflecting the effectiveness of our data augmentation method.