Attention Based Joint Learning for Supervised Electrocardiogram Arrhythmia Differentiation with Unsupervised Abnormal Beat Segmentation

1 INTRODUCTION . Arrhythmia in electrocardiogram ( ECG ) is a reflection of heart conduction abnormality and occurs randomly among normal beats . Deep learning based methods have demonstrated strong power in classifying different types of arrhythmia . There are plenty of works on classifying a single beat , involving convolutional neural networks ( CNN ) ( Acharya et al. , 2017b ; Zubair et al. , 2016 ) , long short-term memory ( LSTM ) ( Yildirim , 2018 ) , and generative adversarial networks ( GAN ) ( Golany & Radinsky , 2019 ) . For these methods to work in clinical setting , however , a good segmenter is needed to accurately extract a single beat from an ECG segment , which may be hard when abnormal beats are present . Alternatively , other works ( Acharya et al. , 2017a ; Hannun et al. , 2019 ) try to directly identify the genres of arrhythmia present in an ECG segment . The limitation of these works is that they work as a black-box and fail to provide cardiologists with any clue on how the prediction is made such as the location of the associated abnormal beats . In terms of ECG segmentation , there are different tasks such as segmenting ECG records into beats or into P wave , QRS complexity , and T wave . On one hand , some existing works take advantage of signal processing techniques to locate some fiducial points of PQRST complex so that the ECG signals can be divided . For example , Pan-Tompkins algorithm ( Pan & Tompkins , 1985 ) uses a combination of filters , squaring , and moving window integration to detect QRS complexity . The shortcomings of these methods are that handcraft selection of filter parameters and threshold is needed . More importantly , they are unable to distinguish abnormal heartbeats from normal ones . To address these issues , Moskalenko et al . ( 2019 ) ; Oh et al . ( 2019 ) deploy CNNs for automatic beat segmentation . However , the quality of these methods highly depends on the labels for fiducial points of ECG signals , the annotation process of which can be laborious and sometimes very hard . Besides , due to the high morphological variation of arrhythmia , strong variations exist even between annotations from experienced cardiologists . As such , unsupervised learning based approaches might be a better choice . Inspired by human ’ s perception of ECG signals , our proposed framework firstly locates the abnormal beats in an ECG segment in the form of attention map and then does abnormal beats classification by focusing on these abnormal beats . Thus , the framework not only differentiates arrhythmia types but also identifies the location of the associated abnormal beats for better interpretability of the result . It is worth noting that , in our workflow , we only make use of annotation for the type of abnormality in each ECG segment without abnormal beat localization information during training , given the difficulty and tedious effort in obtaining the latter . We validate our methods on two datasets from different sources . The first one contains 508 12-lead ECG records of Premature Ventricular Contraction patients , which are categorized into different classes by the origin of premature contraction ( e.g. , left ventricle ( LV ) or right ventricle ( RV ) ) . For the other dataset , we process signals in the MIT-BIH Arrhythmia dataset into segments of standard length . This dataset includes various types of abnormal beats , and we select 2627 segments with PVC present and 356 segemnts with Atrial Premature Beat ( APB ) present . Experiments on both two dataset show quantitative evidence that introducing the segmentation of abnormal beats through an attention map , although unsupervised , can in fact benefit the arrhythmia classification performance as measured by accuracy , sensitivity , specificity , and area under Receiver Operating Characteristic ( ROC ) curve . At the same time , a grade study by experts qualitatively demonstrates our method ’ s promising capability to segment abnormal beats among normal ones , which can provide useful insight into the classification result . Our code and dataset , which is the first for the challenging PVC differentiation problem , will be released to the public . 2 RELATED WORKS . Multitask learning There are many works devoted to training one deep learning models for multitasks rather than one specific task , like simultaneous segmentation and classification . ( Yang et al. , 2017 ) solves skin lesion segmentation and classification at the same time by utilizing similarities and differences across tasks . In the area of ECG signals , ( Oh et al. , 2019 ) modifies UNet to output the localization of r peaks and arrhythmia prediction simultaneously . What those two works have in common is that different tasks share certain layers in feature extraction . In contrast , our segmenter and classifier are independent models and there is no layer sharing between them . As can be seen in Figure 1 , we use attention maps as a bridge connecting the two models . ( Mehta et al. , 2018 ) segments different types of issues in breast biopsy images with a UNet and apply a discriminative map generated by a subbranch of the UNet to the segmentation result as input to a MLP for diagnosis . However , their segmentation and classification tasks are not trained end-to-end . ( Zhou et al. , 2019 ) proposes a method for collaborative learning of disease grading and lesion segmentation . They first perform a traditional semantic segmentation task with a small portion of annotated labels , and then they jointly train the segmenter and classifier for fine-tuning with an attention mechanism , which is applied on the latent features in the classification model , different from our method . Another difference is that for most existing multitask learning works , labels for each task are necessary , i.e. , all tasks are supervised . Our method , on the other hand , only requires the labels of one task ( classification ) , leading to a joint supervised/unsupervised scheme . Attention mechanism After firstly proposed for machine translation ( Bahdanau et al. , 2014 ) , attention model became a prevalent concept in deep learning and leads to improved performance in various tasks in natural language processing and computer visions . ( Vaswani et al. , 2017 ) exploits self-attention in their encoder-decoder architecture to draw dependency between input and output sentences . ( Wang et al. , 2017 ) builds a very deep network with attention modules which generates attention-aware features for image classification and ( Oktay et al. , 2018 ) integrates attention gates into U-Net ( Ronneberger et al. , 2015 ) to highlight latent channels informative for segmentation task . When it comes to ECG , ( Hong et al. , 2019 ) proposes a multilevel knowledge guided attention network to discriminate Atrial fibrillation ( AF ) patients , making the learned models explainable at beat level , rhythm level , and frequency level , which is highly related to our work . Our method and theirs however are quite different in the way attention weights are derived and applied , as well as the output of attention network . First , in that work , the attention weights are obtained from the outputs of hidden layers , while ours are directly from the input . Second , domain knowledge about AF is needed to help the attention extraction , so the process is weakly supervised , while ours do not use any external information and is fully unsupervised . Third , their attention weights are applied to latent features in that work while ours are applied to the input for better interpretability . Finally , in that work , the input ECG segment is divided into equal-length segments in advance and the attention network output only indicates which segment contains the target arrhythmia . The quality highly depends on how the segment is divided , and it does not provide the exact locations of abnormal beats . On the other hand , our method directly locates the abnormal beats on the entire input ECG , offering potentially better interpretability and robustness . 3 METHOD . 3.1 OVERVIEW OF THE FRAMEWORK . Here we briefly introduce the workflow of our joint learning frameworks for supervised classification and unsupervised segmentation . Firstly , in this work , we choose to model the input signal as a one-dimensional signal D ∈ RM×N , where M is the number of leads and N is the length of the input ECG segment ( number of samples over time ) . We then use a one-dimensional ( 1D ) fully convolutional network called segmenter S to output a feature map L = S ( D ) ∈ RM×N , After that , we apply a pooling layer to generate window-style element-wise attention A ∈ RM×N , containing weights directly for every sample in the input ECG . The after-attention signalX = A D ∈ RM×N , where represents element-wise production , is then fed into a multi-layer CNN called classifier C , in which the outermost fully connected layer gives the prediction of the arrhythmia types . After training , the abnormal areas are highlighted in X , thus achieving the goal of segmenting abnormal beats from normal ones . Moreover , x , which indicates those beats that are highly associated with the differentiation task , also serves as an explanation for C ’ s decision . The architecture of our framework is illustrated in Fig 1 . 3.2 SEGMENTER AND CLASSIFIER . In most existing works , the attention map is fused with the deep features in a neural network . However , for our specific purposes of enhancing intepretability of the classification results as well as unsupervised segmentation , the best result would be obtained by directly applying it to the input signal D. In order to generate attention weights of the same length as D , we choose to utilize UNet ( Ronneberger et al. , 2015 ) , a fully convolutional network highlighted by the skip connection on different stages . Encoding path extract features recursively and decoding path reconstruct the data as instructed by loss function . Note that the output of S has only 1 channel and we expand it channel-wise so that it matches the channel dimension of the ECG signal and at the same time each channel gets the same attention . The reason is that the 12 leads are measured synchronously and the abnormal beats occur at the same time across all the leads . Both recurrent neural networks ( RNN ) and CNN are candidate architectures for many arrhythmia classification works . RNN takes an ECG signal as sequential data and is good at dealing with the temporal relationship . CNN focuses on recognition of shapes and patterns in ECG , thus is less sensitive to the relative position of abnormal beats with respect to normal ones . Because abnormal beats may occur randomly among normal beats , we decide to use CNN as the backbone of our classifier . The detailed implementation of C is shown in 1 ( b ) . 3.3 POOLING FOR WINDOW-STYLE ATTENTION . We do not use the output of the segmenter L as the attention map directly but instead perform a pooling with large kernel size first . This is out of considerations for both interpretability and performance . Regarding interpretability , it is desirable that each abnormal beat is uniformly highlighted , i.e. , the attention weights should be almost constant and smooth for all the samples within each abnormal beat . Regarding performance , it is desirable that the attention map A does not distort the shape of abnormal beats after it is applied to the input X . Pooling layer is the easiest way to achieve this goal , functioning as a sliding window over multiple samples in an ECG signal for global information extraction . Max pooling outputs the same value around a local maximum , and average pooling reduces fluctuation by averaging over multiple samples . The kernel size can not be too large , which may fuse sharp changes from neighboring areas and lead to the loss of local information . Therefore , deciding the proper pooling kernel size is essentially finding a balance between local information and global information preservation . Through experiments to be shown in Section 5.3 , we find that setting the kernel size as nearly half the length of a normal beat yields the best balance between performance and interpretability . Padding of zeros on both sides of the segmenter output L is implemented to keep the length of the resulting attention map A after pooling to be the same as the input X . Meanwhile , the polarization of QRS complex is a critical feature for ECG signal , while the commonly used pooling layers , like max pooling and average pooling , fail to control the sign of output , leading to differentiation performance degradation . Rectified linear unit ( ReLu ) σ ( lc , m ) = max ( 0 , lc , m ) , where c and m denote channel number and spatial position in L respectively , is usually performed as an activation function to add non-linearity to neural network for stronger representation ability . In this work , we can apply ReLU on L before pooling so that the all the weights in A generated by the following max pooling are positive . Alternatively , we replace average pooling with L2 norm pooling that takes square root of the L2 norm of input . In that case , ReLU is not needed . The two pooling implementations can be expressed as : ac , m max = PMAX ( L ) c , m = max { σ ( lc , m ) , σ ( lc , m+1 ) , ... , σ ( lc , m+k ) } ( 1 ) ac , m L2 = PL2 ( L ) c , m = √√√√m+k∑ i=m l2c , i ( 2 ) where , ac , m is the mth data point in channel c of A and k is the kernel size for the pooling .

Summary: The paper proposes a framework for the classification of arrhythmias in electrocardiogram (ECG) data. The proposed approach performs segmentation and classification of the ECG signal. The segmenter performs segmentation of the signal (also called attention map) even though the term segmentation is not quite correct. This attention-modulated signal is then classified to identify the origin of Premature Ventricular Contraction (PVC).  The proposed approach is evaluated on a dataset from a single machine consisting of 508 segments (I am not sure what “segments” means in this context). The results seem ok, but it is not clear to me what level of performance is required in order to achieve a similar level of performance as an expert.