SIM-ECG: A Signal Importance Mask-driven ECG Classification System

Heart disease is the leading cause of death in the US and worldwide for both men and women [13]. In the US, heart diseases have been the leading cause of death since 2015, and the number of deaths from heart disease increased by 4.8% from 2019 to 2020 [1]. The Center for Disease Control and Prevention notes that heart disease accounts for one in every four deaths in the United States each year. These deaths are attributed to many factors ranging from undetected heart diseases causing sudden death, late detection that may damage heart muscles and require repair, or even improper monitoring after successful heart surgery. Every year, more than 5 million Americans are affected by Heart Failures. Electrocardiogram (ECG), which records the heart’s electrical activity, has long been the preferred and trusted technique for doctors to detect and diagnose these heart conditions. Therefore, accurate ECG analysis is critical in early disease detection and saving patients’ lives. However, this is not the case quite frequently. A recent study found that 30% of myocardial infarction events were misclassified as low risk, with ECG misinterpretation responsible for half of the misclassifications [7]. Misdiagnosis is also a top concern expressed by cardiac patients [9]. It is therefore essential to develop mechanisms to improve analysis accuracy.

The automation of ECG reading has been a long-standing need. Trying to analyze the signals manually is mundane, requires extreme concentration, causes mental fatigue, and is not well reimbursed adequately, to the amount of effort, by the insurance companies [6]. Consequently, the number of people who can read ECG signals is shrinking, and experienced doctors don’t have enough time to scrutinize patients’ ECGs. However, commercial ECG machines only aid with basic preprocessing and analysis: they give information such as PR intervals, QRS duration, irregular rhythms, and a potential indication of abnormal ECG signals – these analyses are not very sophisticated and offer limited assistance in detection of more complex cardiac conditions [10, 29]. Furthermore, many arising personal ECG monitoring devices also provide basic analysis either locally or send data to more sophisticated processing and eventually, human reading [26].

The machine learning community observed this need, and numerous machine learning algorithms were proposed for disease detection. Convolutional Neural Network assisted in arrhythmia and myocardial detection [12]. Other algorithms such as decision trees, $k$-nearest neighbor, logistic regression, support vector machines [8], and Inception neural networks [20] have also been evaluated. While existing machine learning algorithms have succeeded in classifying basic cardiac conditions, classification of more complex cardiac events is still challenging, and existing solutions are no match for highly skilled human readers. Furthermore, existing solutions provide diagnosis in a black-box manner, requiring the medicinal personnel to carefully analyze the ECG again to validate the algorithm’s interpretations. We propose a system that can continuously collect ECG signal importance mask feedback from expert ECG readers and evolve its detection algorithms to surpass existing static solutions to address those challenges. Furthermore, the proposed system provides a visual interpretation of why the system decided and what portion of the signals were used in making the decision, allowing for quick validation by the reader. Our model and interpretability maps can save manual labor and provide readers with a precise region of the ECG signals where the disease manifestation is visible, which allows them to focus on validating the results instead of trying to find them in the first place.

Deep learning has recently gained popularity in heart disease detection from 12-lead ECG signals. Models like residual convolutional neural networks, convolutional neural networks (CNNs), bidirectional long short term memory (LSTM), CNN and LSTM, CNNs, and densely-connected CNNs, have been used for small scale ($\leq 12$-way) classification problem from 12-lead ECG signals  [21, 2, 16, 14, 11, 31]. Lack of trust by medical personnel in black-box approaches resulted in recent efforts on interpretability analysis of ECG classification models [32], by using convolutional and recurrent neural networks for arrhythmia classification from 12-lead signals and providing explanation from attention mechanism. Similarly, [34] performs interpretability analysis using the Shapley additive explanations method. Recurrent neural networks (RNNs) for atrial fibrillation (AF) detection from a single-lead ECG signal combined with interpretability maps from attention mechanism have been explored [17]. All of the above methods are used for small scale classification ($\leq 12$-way) and do not learn from doctor’s feedback. In contrast, our proposed solution simultaneously tackles large-scale (71-way) multi-label classification tasks, with a training procedure that considers doctors’ signal importance masks (a set of time windows in ECG signal that are most indicative of the diseases exhibits) and provides model interpretability analysis.

Auxiliary information such as bounding boxes in computer vision [5] and rationales in natural language processing [33, 24, 36] have been used to improve model accuracy and interpretability. Theoretical analysis on learning from auxiliary information beyond labels has also been explored [28, 19, 3]. Alignment of the attention maps generated by models to the attention maps labeled by experts to improve the accuracy and interpretability of CNNs has also been explored [15]. The addition of a regularization term in the training loss function that penalizes model gradients outside important regions in examples’ input features has also been explored [23, 4], and we utilize those ideas in our solution. Use of part localization information to train CNNs to improve the accuracy has been proposed in [35]. Recent use of explanation penalization terms to loss function that tries to match the model explanation with human explanation has improved the accuracy and interpretability of image and text classification systems [22]. Although using auxiliary information for improving the trustworthiness of image classification models and natural language models have been well studied in computer vision and natural language processing domains as discussed above, our work is the first that applies it for accurate and interpretable ECG signal-based heart disease prediction.

The goal of our system is to detect cardiac diseases with high accuracy and provide its rationale, i.e., identifying important parts of the input ECG responsible for the predicted disease to allow quick validation by medical personnel and remove the black box approach of just providing the diagnosis. Existing solutions take labeled data to train models and use the trained models to classify new data in the field. This approach is desired if there is no expert to evaluate the data, and the system has to make decisions independently. Applying this classic approach, the doctors are presented with a result of a detected cardiac condition that they can accept or reject. This decision process is a black-box approach, and the doctor does not know how the machine arrived at the given classification. Our process starts similarly, but the learning system and doctor repeatedly communicate with each other with the following protocol: (1) Upon receiving a patient’s ECG, the system detects the patient’s heart condition while also pinpoints parts of the ECG that are responsible for its prediction – for example, it highlights the difference in peaks in the ECG. This makes the decision process more trustworthy and acceptable to the medical community; (2) The doctors review the system’s prediction and explanation and make their treatment decisions. The doctor can also provide feedback to the system by either confirming the prediction and explanation or correcting the prediction and by providing signal importance masks, i.e., parts of the ECG they believe are most relevant to the patient’s condition.

We created a web application where doctors can highlight important regions in 12-lead ECG signals to provide the signal importance masks. Figure 1 shows our interface where doctors can highlight certain regions as well as provide feedback in the form of natural language explanation. The web application was built using Flask with a PostgreSQL database. The doctors highlight all important regions for the given labels, and the feedback is stored in the database, which is subsequently used by our model training algorithm. Although we collect doctors’ text explanations feedback in our system, we do not utilize them in our model training; we leave improving ECG classification models using text feedback as interesting future work.

Our model training algorithm is inspired by recent works on learning by incorporating interpretable auxiliary information [23]. Recall that we would like to train classification models for 71-way heart disease classification from 12-lead ECG signals with signal importance masks being part of training data. Specifically, the training data are represented as a set of tuples $\mathinner{\left\{(x_{i},y_{i})\right\}}_{i=1}^{n}$, where for each example $i$, $x_{i}\in\mathbb{R}^{d}$ is its 12-lead ECG signal representation; in our ECG dataset, $d=12\times 300$. $y_{i}\in\mathinner{\left\{-1,1\right\}}^{K}$ is the class labeling, where $K=71$, for each coordinate $j\in\mathinner{\left\{1,\ldots,K\right\}}$, $y_{j}=-1$ and $+1$ indicate that label $j$ is present and absent, respectively. In our training data, the doctor has provided signal importance masks $M_{i}\subset\mathinner{\left\{1,\ldots,d\right\}}$ to some of the examples – we denote by $E$ the index set of such examples.

Our model is represented by a function $f(x;\theta)$, where $x$ represents the ECG signal, and $\theta$ represents the model parameter. Given $x,\theta$, the model output $f(x;\theta)$ lies in $\mathbb{R}^{K}$, and we use $\operatorname{sign}(f(x;\theta))\mathrel{\mathop{\mathchar 58\relax}}=\mathinner{\left(\operatorname{sign}(f(x;\theta)_{1}),\ldots,\operatorname{sign}(f(x;\theta)_{K})\right)}\in\mathinner{\left\{-1,1\right\}}^{K}$ to denote its multi-label classification result, where $\operatorname{sign}(z)=1$ if $z>0$ and $\operatorname{sign}(z)=-1$ otherwise. Our training process will hinge on finding $\theta$ that has a small average multi-label logistic loss on the training examples; here, the multi-label logistic loss of model $f(\cdot;\theta)$ on a multi-label example $(x,y)$ is defined as $\ell_{\operatorname{logistic}}(\theta,(x,y))=\sum_{j=1}^{K}\ln\mathinner{\left(1+\exp(-y_{j}f_{j}(x;\theta))\right)}$.

Our training objective function follows [23] that proposes a regularized loss objective that takes advantage of signal importance masks, defined as:

$$\min_{\theta}\mathinner{\left(\sum_{i\in E}\ell_{\lambda}(\theta,(x_{i},y_{i},M_{i}))+\sum_{i\notin E}\ell_{\operatorname{logistic}}(\theta,(x_{i},y_{i}))\right)},$$

where

$$\ell_{\lambda}(\theta,(x,y,M))\mathrel{\mathop{\mathchar 58\relax}}=\ell_{\operatorname{logistic}}(\theta,(x,y))+\lambda\sum_{j\in[d]\setminus M}\left(\dfrac{\partial{}\ell_{\operatorname{logistic}}(\theta,(x,y))}{\partial{{x^{j}}}}\right)^{2},$$

(1)

and $\lambda>0$ is a tuning parameter. Each example $(x,y)$ can be viewed as introducing two parts of losses that contribute to the training objective: the first part is the standard multi-label logistic loss $\ell_{\operatorname{logistic}}(\theta,(x,y))$; the second part $\sum_{j\in[d]\setminus M}\left(\dfrac{\partial{}\ell_{\operatorname{logistic}}(\theta,(x,y))}{\partial{{x^{j}}}}\right)^{2}$, only applies to examples in $E$ whose signal importance mask $M$ is available. This term regularizes the sensitivity of the model with respect to the input using the signal importance masks: it penalizes the model from being too sensitive to parts of the ECGs outside their corresponding signal importance masks.

We evaluate our algorithm using the large-scale 71-way multi-label heart disease classification dataset consisting of 12-lead ECG signals provided by [30]. We follow the previously developed pipeline [27] as a baseline system with folds 1-8 as the training set, fold 9 as the validation set, and fold 10 as the test set. The dataset consists of 10-second recordings of the ECG signals with 17441 training examples, 2193 validation examples, and 2203 test examples. Similarly, we follow the previous evaluation metrics to evaluate and compare the models, based on Fmax and macro-AUC scores [27]. Macro AUC score is used to accommodate for class imbalance, as the number of normal ECG signals is quite high [30]. We use the inception1d architecture to perform our experiments, which was shown to perform the best among all architectures [27].

We use PyTorch [18] for all our experiments. Prior approach [27] uses multiple sliding windows of 2.5 seconds and aggregates the predictions from these windows to make a final prediction which acts as a data augmentation during the training phase. We use different training and inference to avoid the long time it takes for doctors to provide feedback for the full 10-second signal. We only use the first 3 seconds of the signal in training set for the feedback. To make a fair comparison between our model training and the baseline system, we use the first 3 seconds for training, but we still use full 10-seconds signals for validation and test.

We collected signal importance mask feedback from doctors for 1359 training examples out of all 17441 training examples. We use Adam optimizer with an initial learning rate of $0.002$ and batch size of 64 for all of our experiments. We repeated our experiment five times.

Figure 2 shows the normal training pipeline (Normal for short), while Figure 3 shows our training pipeline (Feedback for short).

We have built a system that can collect signal importance masks for 12-lead ECG signals from doctors. We have proposed and evaluated an ECG classification algorithm that utilizes signal importance mask feedback from doctors for accurate and interpretable predictions of 71 heart diseases. Our experiments demonstrate that collecting signal importance mask feedback for 1359 out of 17441 (7.7%) training examples already shows good improvement over the normal training baseline. We conjecture that training with larger feedback set can lead to even better models. In addition, we have collected natural language explanation feedback from doctors, which can be utilized in future works to train better models. We leave these as interesting avenues for future investigation.