VizECGNet: Visual ECG Image Network for Cardiovascular Diseases Classification with Multi-Modal Training and Knowledge Distillation

The main goal of VizECGNet is to extract the correlations between different modalities and extract discriminative features via cross- and self-modal attention modules. To achieve this goal, we extract features from the 12-lead ECG signals $\mathbf{X}^{s}=\{\mathbf{x}^{s}_{1},\dots,\mathbf{x}^{s}_{12}\}$ and image $\mathbf{X}^{i}$ using CNN-based feature extractors $\mathbf{f}$ and $\mathbf{g}$, respectively. For each $l$-th single-lead signal $\mathbf{x}^{s}_{l}=[x^{s}_{l,1},\dots,x^{s}_{l,T}]$ with time length $T$, we use twelve different 1D CNN-based feature extractor $\mathbf{f}=\{f_{1},\dots,f_{12}\}$ with sharing parameters to efficiently fuse each signal features as follows:

$$\mathbf{z}^{s}_{\text{avg}}=\frac{1}{12}\sum_{l=1}^{12}f_{l}(\mathbf{x}^{s}_{l})\in\mathbb{R}^{C^{s}\times T^{s}}.$$

(1)

Similar to ECG signal $\mathbf{X}^{s}$, we extract the feature maps from 12-lead ECG image $\mathbf{X}^{i}$ using 2D CNN-based feature extractor $\mathbf{g}$ as follow:

$$\mathbf{z}^{i}=\mathbf{g}(\mathbf{X}^{i})\in\mathbb{R}^{C^{i}\times H^{i}\times W^{i}},$$

(2)

where $C^{s}=C^{i}=512$. Subsequently, we apply cross-modal attention module between two extracted features $\mathbf{z}^{s}_{\text{avg}}$ and $\mathbf{z}^{i}$. For clarify, let assume we use self-attention on modality $m$ based on modality $n$. Then, we can write CMAM as follows:

$$\bar{\mathbf{z}}^{m\leftarrow n}=\text{Softmax}\left(\mathbf{Q}_{n}\mathbf{K}^{T}_{n}\right)\mathbf{V}_{m}\\$$

(3)

where $\mathbf{Q}_{n}=W_{\mathbf{Q}_{n}}\mathbf{z}^{n},\mathbf{K}_{n}=W_{\mathbf{K}_{n}}\mathbf{z}^{n},\mathbf{V}_{m}=W_{\mathbf{V}_{m}}\mathbf{z}^{m}$. Then, each refined features $\bar{\mathbf{z}}^{s\leftarrow i}$ and $\bar{\mathbf{z}}^{i\leftarrow s}$ contain different modality information. Now, we apply SMAM to extract discriminative features for two features with mixed information in each modal stream as follows:

$$\begin{cases}&\hat{\mathbf{z}}^{s}=\textbf{Softmax}\left(\mathbf{Q}_{s}\mathbf{K}^{T}_{s}\right)\mathbf{V}_{s}\\ &\hat{\mathbf{z}}^{i}=\textbf{Softmax}\left(\mathbf{Q}_{i}\mathbf{K}^{T}_{i}\right)\mathbf{V}_{i}\end{cases}$$

(4)

where $\mathbf{X}_{m}=W_{\mathbf{X}_{m}}\bar{\mathbf{z}}^{m\leftarrow n}$ for $m\in\{s,i\}$ and $\textbf{X}\in\{\textbf{Q},\textbf{K},\textbf{V}\}$. Each final refined discriminative features $\hat{\mathbf{z}}^{s}$ and $\hat{\mathbf{z}}^{i}$ are forwarded into the fully-connected layers for each modality stream to classify abnormal ECG signal types.

Knowledge distillation is a technique commonly used in machine learning to transfer knowledge from a complex model (teacher model) to a simple model (student model). To utilize only printed ECG signal to classify abnormal signals during inference and acquire knowledge about 12-lead ECG signals, we adopt knowledge distillation from signal stream into image stream. First, predicting the probability distributions $p^{s}$ and $p^{i}$ for each class is performed using a classifier for each modality stream as follows:

$$p^{m}=\textbf{MLP}_{m}\left(\textbf{GAP}\left(\hat{\mathbf{z}}^{m}\right)\right)$$

(5)

where $\textbf{MLP}_{m}(\cdot)$ is a multi-layer perceptron (MLP) for each signal and image modality stream. For each prediction, the classification loss function $\mathcal{L}_{cls}$ for the same label $t$ is computed as follows:

$$\mathcal{L}_{cls}=\sum_{c=1}^{C}\left(\mathcal{L}_{BCE}\left(t_{c},p^{s}_{c}\right)+\mathcal{L}_{BCE}\left(t_{c},p^{i}_{c}\right)\right)$$

(6)

where $C$ is a number of classes and $\mathcal{L}_{BCE}$ is the binary cross-entropy loss function. Since the dataset used in this paper can have multiple diseases for a single ECG signal, a binary classification is performed for each class. Finally, we calculate knowledge distillation loss $\mathcal{L}_{KD}$ between two modality streams to reduce the difference in probability distribution $p^{s}$ and $p^{i}$ as follows:

$$\begin{split}\mathcal{L}_{KD}(p^{s},p^{i})=\sum_{c=1}^{C}\mathcal{L}_{KL}(p^{s}_{c}||p^{i}_{c})\\ =\sum_{c=1}^{C}\sum_{x\in\mathcal{X}}p^{s}_{c}(x)\text{log}\left(\frac{p^{s}_{c}(x)}{p^{i}_{c}(x)}\right)\end{split}$$

(7)

where $\mathcal{L}_{KL}$ is a Kullback-Leibler Divergence to calculate the difference between two probability distributions $p^{s}_{c}$ and $p^{i}_{c}$ for each class $c$. The final loss function $\mathcal{L}_{total}=\lambda_{1}\mathcal{L}_{cls}+\lambda_{2}\mathcal{L}_{KD}$ is used for updating parameters of each modality stream. To reduce the sensitivity of hyperparameters, we fix $\lambda_{1}=\lambda_{2}=1$.

We implemented VizECGNet in Pytorch 1.11 and Python 3.8. The large-scale 12-lead ECG signal dataset [15] used in this paper is multi-label (1dAVb, RBBB, LBBB, SB, AF, ST). The ECG signal of each lead is composed of 4096 time lengths. In this paper, to eliminate the trend of each signal, the average voltage for the entire time is set to zero mean, and we apply a detrending process. Finally, we convert the 12-lead ECG signals using the ecg_plot Python library for multi-modal learning on the images. We compared our VizECGNet with five signal-based models (InceptionTime [8], XResNet1D-101 [9], Transformer [10], ACNet [11], and 1D RANet [7]), two image-based models (ResNet18 [12] and MobileNetV3 [13]), and multi-modal models (MHM [14]). Since the use of default training settings generally performs poorly in ECG dataSet, for fair comparison, we optimized the parameters to work best on ECG dataset. We train all models in an end-to-end manner using the Adam optimizer. The initial learning rate starts from $10^{-3}$ and is decreased to $10^{-6}$ using the cosine annealing learning rate scheduler [16], and the training settings were set to a batch size of 16 and epochs of 300 till the loss functions of all models converged. For evaluation, we used three metrics (Precision, Recall, and macro-averaged F1-Score) to measure the performance of each model. To efficiently extract features from signals and image, we utilize ResNet18 as feature extractor.

As shown in Table 1, VizECGNet outperforms single- and multi-modality models on all performance evaluation metrics in both inference data types. When predicting using ECG signals, VizECGNet achieve 9.58%, and 19.68% higher F1-Score compared with 1D RANet, and MHM, respectively. Furthermore, when using printed ECG images, VizECGNet achieve 17.21%, and 24.06% higher F1-Score compared with ResNet18, and MHM, respectively. Note that MobileNetV3 and ResNet18 used only simple ECG images for training and evaluation. This training strategy makes two image-based models unable to understand the abnormal signal characteristics of each lead. However, although 1D RANet receives 12-lead ECG signal data with complex data structures, its performance is low. These reasons indicate that the interpretability of 1D RANet for complex data is still poor. On the other hand, VizECGNet achieves high classification performance because it exchanges information based on multi-modal learning and distills the knowledge of abnormal signals generated by subtle differences in 12-lead ECG signals into an image modality stream during learning.

We also examined extrapolation to actual ECG printed images to confirm the utility of the model (Fig. 2). In addition, the inference results of two image-based models (ResNet18 and MobileNetV3) are added. Fig. 2.(e) shows the disease probability for each model for each disease (AF, RBBB, LBBB, 1dAVB). Our model was positive for all diseases, but the other two models were negative for all but well-characterized diseases such as 1dAVB. These results suggest that VizECGNet is more practical than other models because it can be applied to real printed images even though it was trained on synthetic ECG images created from 12-lead ECG signals.

In this section, we analyze the effectiveness of two attention modules (CMAM and SMAM). In Table 2, the results of the ablation study on attention modules are listed with four configurations. Basically, attention is used to extract discriminative features from messy features. In this paper, important information can be additionally extracted by focusing on each modality or between modalities. It can be seen from the actual experimental results that when the mode is paid attention to, all performance evaluation indicators have achieved high performance.