[Paper Review] Using Multitask Learning to Improve 12-Lead Electrocardiogram Classification

J. Weston Hughes, Taylor Sittler+4 more|arXiv (Cornell University)|Dec 2, 2018

ECG Monitoring and Analysis4 citations

TL;DR

This paper proposes a multi-task convolutional neural network (CNN) that improves 12-lead ECG classification by jointly learning multiple cardiac diagnoses, significantly boosting performance on rare arrhythmias through shared representation learning. Using a 41,522-sample ECG dataset with 24 binary labels, the multi-headed model achieves higher F1 scores—especially for rare classes like 2nd Degree AV Block (Mobitz I), where performance improved from 0.58 to 0.73—demonstrating that supervision from related common classes enhances rare-class generalization.

ABSTRACT

We develop a multi-task convolutional neural network (CNN) to classify multiple diagnoses from 12-lead electrocardiograms (ECGs) using a dataset comprised of over 40,000 ECGs, with labels derived from cardiologist clinical interpretations. Since many clinically important classes can occur in low frequencies, approaches are needed to improve performance on rare classes. We compare the performance of several single-class classifiers on rare classes to a multi-headed classifier across all available classes. We demonstrate that the addition of common classes can significantly improve CNN performance on rarer classes when compared to a model trained on the rarer class in isolation. Using this method, we develop a model with high performance as measured by F1 score on multiple clinically relevant classes compared against the gold-standard cardiologist interpretation.

Motivation & Objective

• To address the challenge of low performance in classifying rare cardiac arrhythmias in 12-lead ECGs due to insufficient training data.
• To investigate whether multi-task learning can improve classification accuracy for rare ECG diagnoses by leveraging related, more frequent classes.
• To develop a high-performing, multi-headed deep learning model for 12-lead ECGs that generalizes across multiple clinically relevant diagnoses.
• To evaluate whether adding even a single related common class can significantly improve performance on rare classes, reducing the need for large-scale data collection for rare conditions.

Proposed method

• A 34-layer convolutional neural network (CNN) is trained on 12-lead ECGs sampled at 250Hz over 10 seconds, resulting in a 12×2500 input tensor.
• The model uses a multi-headed architecture with one output head per diagnostic class, enabling joint training on all 24 binary classification tasks simultaneously.
• Each head predicts the presence or absence of a specific cardiac diagnosis using binary cross-entropy loss, with gradients backpropagated through shared convolutional layers.
• The dataset was resampled to balance class frequencies, with minority classes up-sampled to 4,000 examples to mitigate class imbalance.
• Performance is evaluated using F1 score across all classes, with comparisons between single-headed and multi-headed models trained on varying subsets of classes.
• The model leverages hierarchical and correlated relationships between classes (e.g., AV blocks, tachycardias) to improve representation learning in shared feature layers.

Experimental results

Research questions

• RQ1 Can multi-task learning improve the F1 score of a deep learning model on rare 12-lead ECG diagnoses by sharing representations with more common classes?
• RQ2 Does the inclusion of a single related, common class (e.g., 1st Degree AV Block) provide a performance boost comparable to including many unrelated classes for a rare diagnosis like 2nd Degree AV Block (Mobitz I)?
• RQ3 How does the performance of a multi-headed model compare to single-headed models trained in isolation on rare classes, particularly when rare classes have fewer than 1,000 training examples?
• RQ4 To what extent does multi-task learning enhance feature learning in deeper layers of a CNN when applied to ECGs with heterogeneous, low-frequency diagnostic labels?

Key findings

• The multi-headed model outperformed single-headed models on all classes, with the most significant F1 score improvements observed for rare classes such as Complete Heart Block (F1: 0.62 vs. 0.09) and Atrioventricular Nodal Reentry Tachycardia (F1: 0.49 vs. 0.39).
• For 2nd Degree AV Block (Mobitz I), a rare class with 626 examples, the multi-headed model achieved an F1 score of 0.73, a substantial improvement over the single-headed model’s 0.58.
• Adding just one related common class (1st Degree AV Block) to the training of 2nd Degree AV Block (Mobitz I) yielded the same performance gain as adding all other classes, indicating that related supervision is highly effective.
• Common classes like Sinus Rhythm (F1: 0.91) and Atrial Fibrillation (F1: 0.85) achieved high performance in both single- and multi-headed models, but the multi-headed model still provided modest gains.
• The model demonstrated that shared representation learning across multiple tasks improves feature extraction, particularly in deeper layers, leading to better generalization on low-resource classes.
• The results suggest that multi-task learning can reduce the data efficiency barrier for rare ECG diagnoses, potentially enabling high-performance models even with fewer than 1,000 examples per rare class.