Honours Dissertation · University of Glasgow · 2022
Explainable ML for Arrhythmia Detection from ECG
Can a deep-learning model that flags irregular heartbeats also explain itself well enough to be trusted in a clinical setting? This project builds and tests models that try to do both.
Overview
Cardiovascular disease is one of the leading causes of death worldwide, and reading ECGs manually can be slow. Machine learning can screen heartbeats far faster, but deep models tend to behave like black boxes and are difficult to be trusted by clinicians, who need to know why a decision was made.
This dissertation tackles that issue directly. It proposes four convolutional models that build explainability into training itself, then measures both how accurately they classify heartbeats and how clearly they explain their own decisions.
Approach
Everything was trained on the MIT-BIH arrhythmia dataset, classifying individual beats into eight clinical categories. Because normal beats largely outnumber the rest, the data was resampled to balance the classes before training.
- Two baselines: a support vector machine and a one-dimensional convolutional neural network.
- Four novel CNN variants that add gradient-based attribution priors to the loss function using integrated gradients and expected gradients, each in an L1 and an L2 form.
- A qualitative comparison of three explainability methods (integrated gradients, expected gradients and GradCAM) via heatmaps over the beat waveform.
- Two evaluation regimes: a beat-holdout test with five repeats of stratified 10-fold cross-validation, and a stricter patient holdout test that keeps whole patients out of training to avoid data leakage.
Results
- The baseline CNN was the strongest classifier overall in the beat holdout experiment, with significance confirmed through Friedman and Nemenyi tests.
- Among the novel models, the expected-gradients L2 variant performed best by edging past the baseline CNN on every metric in the patient-holdout test (~0.99 accuracy), though that run wasn't repeatable, so the edge can't be called statistically significant.
- The L2 attribution priors consistently beat their L1 counterparts.
- Integrated gradients produced the clearest, least noisy explanations of the three methods, making it easiest to see why a model got a beat wrong.
Takeaways
These findings revealed that expected gradients as an attribution prior is a viable, practical way to train ECG classifiers that stay explainable and it delivers competitive accuracy at a fraction of the training cost of integrated gradients.
It also revealed clear next steps: designing a penalty function better suited to time-series than the image-borrowed one used here, and running the patient holdout test on a larger dataset so its results can be validated statistically.