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: This manuscript contributes a neural architecture to classify arrhythmia type from ECG data. The signal treated as 1D, and the architecture performs joint segmentation-classification detecting the abnormal beats and then classifying them as a function of their origine. It uses U-nets for segmentation and, for classification CNN and one fully-connected layer. The unet segmentation generates weights that are considered as an attention map and multipled with the original time series after pooling on a window (which amounts to smoothing).