{'fraction_non_alphanumeric': 0.03280234231551784, 'fraction_numerical': 0.007642300630241675, 'mean_word_length': 4.9638946433856175, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Optimal Transport based Data Augmentation for Heart Disease Diagnosis and Prediction

In this paper, we focus on a new method of data augmentation to solve the data imbalance problem within imbalanced ECG datasets to improve the robustness and accuracy of heart disease detection. By using Optimal Transport, we augment the ECG disease data from normal ECG beats to balance the data among different categories. We build a Multi-Feature Transformer (MF-Transformer) as our classification model, where different features are extracted from both time and frequency domains to diagnose various heart conditions. Learning from 12-lead ECG signals, our model is able to distinguish five categories of cardiac conditions. Our results demonstrate 1) the classification models' ability to make competitive predictions on five ECG categories; 2) improvements in accuracy and robustness reflecting the effectiveness of our data augmentation method.


I. INTRODUCTION
The 12-lead Electrocardiogram (ECG) is the foundation for much of cardiology and electrophysiology. It provides unique information about the structure and electrical activity of the heart and also systemic conditions, through changes in timing and morphology of the recorded waveforms. Computer generated interpretations are standardly provided following ECG acquisition, utilizing predefined rules and algorithmic pattern recognition. However, current approaches miss a lot of the specialized insights and nuances that practiced cardiologists/electrophysiologists can see. Depending on experience, physician reads can also be variable and inconsistent. Achievement of reliable ECG reading would be a significant achievement, akin to reliable autonomous vehicles, not least with respect to safety, such that critical and timely ECG interpretations of acute cardiac conditions can lead to efficient and cost-effective intervention.
A complete ECG analysis would consist of being able to identify the separate waveforms and their relationships, measure out the intervals between the waveforms representing specific electrical events in different areas of the heart, diagnose those conditions leading to the abnormalities noted on ECGs, understand the biology/physiology responsible for the perturbations in rhythm, generally localize the site of the perturbation anatomically, and make predictions of the future cardiac events based on current findings. The computer reads at the present time do an excellent job of the first two items -picking out the PQRST elements and measuring their duration and intervals -and are reasonably successful in diagnosing the more apparent conditions, but largely fail beyond that.
With the development of machine learning and deep learning methods, it may be possible to identify additional previously unrecognized signatures of disease. Many models have been applied on diagnosing physiological signals, i.e., EEG, ECG, EMG, etc [1]- [3]. One of the main issues is that organizing physiological datasets is extremely difficult, which is constrained by domain knowledge, time, and privacy regulations. There are not so many publicly available ECG datasets, and the available data is limited by specific diseases categories, making it even harder to apply a complex classification model to learn the patterns, as it requires lots of labeled training samples. Another problem is that the ECG data is mostly imbalanced, where the number of labeled ECG signals for a certain condition is very small, so the training samples contains many healthy ECG signals, making it difficult to classify the ECG signals with diseases due to the introduced imbalance in the original data.
In this paper, we proposed a new data augmentation method based on Optimal Transport (OT) to solve the data imbalance issue within the ECG signals. We augmented ECG disease categories, where the number of available data is small, from the healthy signals. We build the Multi-Feature Transformer model as our classifier to evaluate the performance of our proposed method.

II. RELATED WORK
Traditionally, when facing an imbalanced data problem, data augmentation is required before training, aiming to eliminate the effect caused by the imbalance. Traditional methods include sampling, cost-sensitive methods, kernelbased methods, active learning methods, and one-class learning or novelty detection methods [4]. Among them, sampling methods are mostly used, including random oversampling and undersampling, informed undersampling, synthetic sampling with data generation, adaptive synthetic sampling, sampling with data cleaning techniques, cluster-based sampling method, and integration of sampling and boosting. But traditional methods may introduce their own set of problematic consequences that can potentially hinder learning [5]- [7], which can cause the classifier to miss important concepts arXiv:2202.00567v1 [eess.SP] 25 Jan 2022 pertaining to the majority class, or lead to overfitting [4], [6], making the classification performance on the unseen testing data generally far worse.
The ECG data imbalance issue has been a long-standing problem. Martin et al. tried to use oversampling method to augment the imbalanced data [8]. Virgeniya et al. [9] also addressed the ECG data imbalance problem, where instead of using synthetic models such as synthetic minority oversampling technique (SMOTE), SMOTEBoost, or DataBoostIM, they tried to feed the data into the adaptive synthetic (ADASYN) [10] based sampling model, which utilized a weighted distribution for different minority class samples depending upon the learning stages of difficulty. Liu et al. [11] augmented the ECG data by using bandpass filter, noise addition, time-frequency transform and data selection. The methods above showed that balanced dataset performance is superior than unbalanced one.
Optimal Transport (OT) is a field of mathematics that studies the geometry of probability spaces [12]. The theoretical importance of OT is that it defines the Wasserstein metric between probability distributions. It reveals a canonical geometric structure with rich properties to be exploited. The earliest contribution to OT originated from Monge in the eighteenth century. Kantorovich rediscovered it under a different formalism, namely the Linear Programming formulation of OT. With the development of scalable solvers, OT is widely applied to many real-world problems [13], [14].
With the development in machine learning, many models have been applied to ECG disease detection [15]- [20]. Al-Zaiti et al. predicted acute myocardial ischemia in patients with chest pain with a fusion voting method [21]. Acharya et al. proposed a nine-layer deep convolutional neural network (CNN) to classify heartbeats in the MIT-BIH Arrhythmia database [22], [23]. Shanmugam et al. estimate a patient's risk of cardiovascular death after an acute coronary syndrome by a multiple instance learning framework [2]. Recently, Smigiel et al. proposed models based on SincNet [24] and used entropy-based features for cardiovascular diseases classification [25].
ECG signal can be considered as one type of sequential data, and Seq2seq models [26] are widely used in time series tasks. Since the attention mechanism was proposed [27], the Seq2seq model with attention has been improved in various tasks, which outperformed previous methods. Then Transformer model [28] was proposed to solve the problem in the Seq2Seq model, replacing Long Short-Term Memory (LSTM) models with an attention structure, which achieved better results in translation tasks. The transformer model has also recently been adopted in several ECG applications, i.e., arrhythmia classification, abnormalities detection, stress detection, etc [29]- [34]. But those models take only ECG temporal features as input and haven't considered the frequency domain features. To take advantage of multiple features across time and frequency domains, we proposed a Multi-Feature Transformer as our classification model to predict the heart diseases with 12-lead ECG signals.

A. Overall Pipeline
The overall pipeline of our method is shown in Fig. 1. Given the raw ECG data, we first use Optimal Transport to augment the data of minority categories to solve the data imbalance issue. Then we perform pre-processing on the raw ECG data as well as the augmented data, which includes data denoising, ECG temporal segmentation, and feature extraction. We extract multiple ECG features from both the time domain and frequency domain. Then we use the MF-Transformer as our classification model for performance prediction and evaluation. The details of each part are introduced in Section III-B, Section III-C, and Section IV respectively.

B. Multi-Feature Transformer
For the classification model, we take advantage of the transformer encoder [28], and proposed a Multi-Feature Transformer (MF-Transformer) model. The transformer is based on the attention mechanism [28] and outperforms previous models in accuracy and performance. The original transformer model is composed of an encoder and a decoder. The encoder maps an input sequence into a latent representation, and the decoder uses the representation along with other inputs to generate a target sequence. Our model is mostly based on the encoder, since we aim at learning the representations of ECG features, instead of decoding it to another sequence.
The input for the Multi-Feature Transformer is composed of three parts, including ECG raw features, time-domain features, and frequency domain features. The detailed feature pre-processing steps are introduced in Section IV. First, we feed out the input into an embedding layer, which is a learned vector representation of each ECG feature by mapping each ECG feature to a vector with continuous values. Then we inject positional information into the embeddings by: The attention model contains two sub-modules, a multiheaded attention model and a fully connected network. The multi-headed attention computes the attention weights for the input and produces an output vector with encoded information on how each feature should attend to all other features in the sequence. There are residual connections around each of the two sub-layers followed by a layer normalization, where the residual connection means adding the multi-headed attention output vector to the original positional input embedding, which helps the network train by allowing gradients to flow through the networks directly. Multi-headed attention applies a self-attention mechanism, where the input goes into three distinct fully connected layers to create the query, key, and value vectors. The output of the residual connection goes through a layer normalization.
In our model, our attention model contains N = 5 same layers, and each layer contains two sub-layers, which are a multi-head self-attention model and a fully connected feedforward network. Residual connection and normalization are added in each sub-layer. So the output of the sub-layer can be expressed as: For the Multi-head self-attention module, the attention can be expressed as: where multi-head attention uses h different linear transformations to project query, key, and value, which are Q, K, and V , respectively, and finally concatenate different attention results: where the projections are parameter matrices: where the computation of attention adopted scaled dotproduct: For the output, we use a 1D convolutional layer and softmax layer to calculate the final output.

C. Optimal Transport Based Data Augmentation
We use optimal transport to push forward samples from the distribution of a majority class to a minority class. We expect optimal transport to exploit global geometric information so that the synthetic samples match the real samples. In specific, we denote the data from a majority class to be X s = {x s,1 , ..., x s,ns } ∈ Ω s and the minority class data to be X t = {x t,1 , ..., x t,nt } ∈ Ω t . We assume that they are subject to distributions X s ∼ µ s and X t ∼ ν t , respectively, and we associate empirical measures to data samples: where δ x is the Dirac function at location x and p i are the probabilities masses associated to the samples. Solving the optimal transport objective give us the coupling: where C i,j = x i − x j 2 2 is a cost matrix, γ is a coefficient, and H(π) = π i,j log π i,j is the negative entropy regularization that enable us to employ the celebrated Sinkhorn algorithm [35]. The solution to problem (9) actually express the barycentric mappinĝ where x s,i is source sample andx s,i is the resulting mapped sample. When using l 2 norm as cost function, the barycenter has a convenient format that maps the source samples into the convex hull of target samples asX s = n s π * X t .

IV. DATASET AND REPROCESSING
We carried out the experiments on the PTB-XL dataset [36], which contains clinical 12-lead ECG signals of 10second length. There are five conditions in total, which include Normal ECG (NORM), Myocardial Infarction (MI), ST/T Change (STTC), Conduction Disturbance (CD), and Hypertrophy (HYP). The waveform files are stored in Wave-Form DataBase (WFDB) format with 16-bit precision at a resolution of 1µV/LSB and a sampling frequency of 100Hz.
First, we read the raw data by wfdb library 1 and perform Fast Fourier transform (fft) to process the time series data into the spectrum, which is shown in Fig. 3. Then we perform n-points window filtering to filter the noise and adopt notch processing to filter power frequency interference (noise frequency: 50Hz, quality factor: 30), where the filtered result is shown in Fig. 4.  We then detect the R peaks of each signal by ECG detectors 2 , so the data can be sliced at the fixed-sized interval on both sides to obtain individual beats. The examples of detecting R peaks in ECG signals and divided pieces are shown in Fig. 5 and Fig. 6, respectively.  To reduce the dimension of ECG features, we downsample the processed ECG signals to 50Hz. Then we extract more time domain features and frequency domain features to better represent the ECG signals. The time-domain features include: maximum, minimum, range, mean, median, mode, standard deviation, root mean square, mean square, k-order 1 https://pypi.org/project/wfdb/ 2 https://pypi.org/project/py-ecg-detectors/ moment and skewness, kurtosis, kurtosis factor, waveform factor, pulse factor, margin factor. The frequency-domain features include: fft mean, fft variance, fft entropy, fft energy, fft skew, fft kurt, fft shape mean, fft shape std, fft shape skew, fft kurt, which are shown in Table I.
Like many other datasets, the imbalance issue is a problem and needs to be solved before further steps. There are five categories in total, including NORM, MI, STTC, CD, and HYP. In a balanced dataset, each category should occupy the same proportion. In the original dataset, the number of patients in the NORM category is much larger than the others. After dividing the ECG signals into individual beats, the portion of each category changed due to heartbeat variance among people. However, if we count the segmented ECG beats and compare different categories' data, the imbalance issue still exists, which is shown in Table II. From Table  II, we can find out that NORM category and CD category is much larger than the other three categories, making the dataset unbalanced.

A. Data Augmentation by Optimal Transport
As discussed in Section (III-C), the augmented data generated by Optimal Transport is used to eliminate the data imbalance among different categories. In specific, (1) We always use NORM individual beats as the source and transport the samples from the NORM into each other minor categories. (2) In the augmentation procedure, we randomly Fig. 7: Examples of the original ECG signals and the augmented ECG signals within different conditions. The top row shows the 10-second 12-lead ECG signals of different heart condition categories. The bottom row shows the corresponding transported samples obtained from our Optimal Transport based data augmentation method. sample a batch of ECG signal from both the source and target categories and then use formulation in Equation (10) to get the barycentric mapping samples. The label of augmented samples are set to be the target categories. (3) We mix the original data and augmented data and then process them for the MF-Transformer following Fig. 1. We segment the ECG into individual beats, then we obtain the time and frequency statistical features with the method introduced in Section IV. After that we concatenate the ECG signals with all the features as the input for the MF-Transformer model. Examples of augmented data is shown in Fig. 7. From the results, we can find the augmented data preserves the semiperiodic nature, and the results of each lead fit well with the ECG pattern compared with original ECG signals by domain knowledge.

B. Heart Disease Detection
To evaluate our methods, we performed experiments on the PTB-XL dataset to predict the category of ECG signals.
First, we trained the MF-Transformer model with the original PTB-XL data to obtain the baseline performance for different categories. Second, we used the oversampling strategy to augment the ECG signals for the minority categories, then we trained the MF-Transformer model from scratch to obtain the performance by oversampling data augmentation method. Third, we augmented the data with our OT-based data augmentation method, and trained the MF-Transformer model from scratch again to evaluate the performance of our method. Note that the augmented data is only used for training, and the testing set remains the same as for all the experiments, which only contain the real-world ECG signals to have a fair evaluation of the proposed method.
The training and testing splitting strategy is the same as in [20], [36]. The experiments are carried out on four Nvidia Tesla V100 GPUs. The results of different approaches are shown in Fig. 8, which shows the corresponding accuracy pf each category in the confusion matrix.
From Fig. 8(a) we can see that due to the imbalance within the dataset, the accuracies of NORM and CD categories are much higher than the other three, meaning the model learns more patterns from those two categories compared with the others, which could have a negative impact since the other three heart disease categories can only achieve around 50% ∼ 60% classification accuracy. So the data imbalance issue needs to be resolved to improve the model's performance and robustness to different heart conditions. Fig. 8(b) shows by oversampling data from minority categories to make the data more balanced, the classification accuracy increased, especially for the minority categories. But the improvement is not high enough, and could easily lead to overfitting. Fig. 8(c) shows the results by learning from both raw data and our OT-augmented data. Compared with the oversampling results shown in Fig. 8(b), we can see that not only the classification accuracy of each category has improved, but the average classification result has also increased from 71.80% (original) and 72.05% (oversampling) to 75.82% (ours). Each category's performance comes to be more balanced, showing the robustness improvement compared with the baseline results and oversampling results in Fig. 8(a) and Fig. 8 (b).

VI. CONCLUSIONS AND FUTURE WORK
In this paper, we proposed a new method to deal with the ECG data imbalance problem. We augmented the minority category from the majority category with Optimal Transport to make the data balanced, which can handle the overfitting issue introduced by the traditional sampling method. We also proposed an MF-Transformer as our classification model to predict the heart conditions from ECG signals. We showed that after data augmentation, there are both accuracy and robustness improvements on the classification results over five ECG categories, which demonstrate the effectiveness of our method.
Our future work will incorporate standard ECG intervals, axis, and amplitudes (PR interval, P wave amplitude, etc.) to provide a solution with more interpretability.