Arrhythmia Classification Using Graph Neural Networks Based on Correlation Matrix

arXiv:2410.10758 · eess.SP, cs.AI · Submitted 2025-02-10 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Next we'll be talking about the paper "Arrhythmia Classification Using Graph Neural Networks Based on Correlation Matrix".

Jane: The paper was written by Seungwoo Han from Tokyo University of Agriculture and Technology.

Tom: Stay tuned as we take you through the paper and discuss its implications.

Title: Tom: Hey everyone, welcome back to the show! I'm Tom, and as always, I'm here with my co-host Jane. Today we're diving into a really fascinating paper that just hit arXiv, titled "Arrhythmia Classification Using Graph Neural Networks Based on Correlation Matrix."

Jane: And I'm Jane! Tom, I have to say, when I first saw that title, I got a little excited. We're talking about using graph neural networks — that's the GNN part — to classify arrhythmias, which are basically irregular heartbeats. And the twist here is they're using a correlation matrix to build the graph structure.

Tom: Exactly! And for anyone listening who might not be deep in the machine learning world, let's break that down. An ECG is that squiggly line you see when a heart is being monitored. It measures the electrical activity of the heart. Arrhythmias are when that rhythm goes off — too fast, too slow, or just irregular.

Jane: Right. And traditionally, people have used all sorts of deep learning models — like convolutional neural networks — to look at those ECG signals and classify what kind of arrhythmia it is. But this paper takes a different approach. Instead of just feeding the raw signal into a neural network, they're building a graph.

Tom: A graph! Like nodes and edges. So imagine each feature of the heartbeat — like the time between certain waves, the amplitude of the waves — those are the nodes. And the edges between them are determined by how correlated those features are with each other.

Jane: And that's where the "correlation matrix" part comes in. They calculate the Pearson correlation coefficient between all pairs of features. If two features are highly correlated — like above zero point nine — they draw an edge between them. If not, no edge.

Tom: So the graph structure itself is telling you something about the relationships between different parts of the heartbeat. And then they feed that graph into a graph neural network, which can learn patterns from the way those features are connected.

Jane: I love this because it's not just throwing data at a black box. The graph is built on a statistical relationship that actually makes sense — features that move together are connected. That's a really interpretable way to think about the data.

Tom: And that's the core idea. The authors are from Tokyo University of Agriculture and Technology, by the way — Seungwoo Han. They're trying to make arrhythmia classification more explainable by grounding the graph structure in something meaningful.

Jane: Which is a big deal in medical AI, right? Doctors don't just want a prediction — they want to know why the model thinks something is happening.

Tom: Exactly. And that's what we're going to dig into. We're going to look at how they built this model, what the results were, and whether this approach could actually change how we do ECG analysis in the clinic.

Jane: So stick around — we've got a lot to unpack. Next up, we're going to talk about the actual method and the data they used.

Summary: Tom: Alright, we're back! And Jane, we just set the stage with the title — "Arrhythmia Classification Using Graph Neural Networks Based on Correlation Matrix" — but now let's actually talk about how they did it.

Jane: Right. So they used a really well-known dataset called the MIT-BIH Arrhythmia Database. It's got forty-eight recordings of two-lead ECG data, sampled at three hundred sixty Hz, with each recording being about thirty minutes long.

Tom: And they followed the AAMI guidelines, which is basically the standard for how you should split training and testing data for arrhythmia classification. They excluded four records, and then split the rest into twenty-two records for training and twenty-two for testing.

Jane: And importantly, they only looked at three classes — N for non-ectopic beats, S for supraventricular ectopic beats, and V for ventricular ectopic beats. Those are the major categories you see in arrhythmia studies.

Tom: So after filtering the ECG signals — they used moving average filters and a FIR filter to clean up noise — they segmented the signals around each R peak. That's the big spike in the heartbeat signal. They took eighty-six samples before and one hundred thirty samples after each peak.

Jane: Then they extracted twenty features per beat. Things like the P-R amplitude, the Q-T time span, the variance of each beat, and some R-R interval features. So for each heartbeat, you get a vector of twenty numbers.

Tom: And then comes the clever part. They computed the Pearson correlation matrix across all those features. If the correlation coefficient was zero point nine or higher, they set it to one — meaning there's an edge between those two features. Otherwise, it's zero.

Jane: So the adjacency matrix — which defines the graph structure — is built entirely on statistical correlation. And the node features are just the extracted features themselves.

Tom: Then they feed that graph into a GraphSAGE model — that's a specific type of graph neural network — and they also feed the raw features into a simple linear layer. Then they concatenate the outputs from both branches and pass that through another linear layer to get the final classification.

Jane: So it's a fusion model. You've got the graph neural network learning from the relationships between features, and you've got a linear network learning from the raw features directly. And they combine those to make a prediction.

Tom: And they trained it with the Adam optimizer, a learning rate of zero point zero one, cross-entropy loss, and seven hundred epochs. That's a pretty standard setup, but the graph structure is what makes it unique.

Jane: And I think that's the key insight here — they're not just using a GNN because it's trendy. They're using it because the graph structure itself encodes something meaningful about the ECG features. That's what makes this paper interesting.

Tom: Definitely. And now we need to talk about whether it actually worked. So let's move on to the results and how they compare to other methods.

Jane: Good, because I'm curious to see if this graph-based approach actually holds up against the more traditional deep learning models.

Improvements: Tom: So we're at the results part now, and Jane, I've got to say — the numbers are a mixed bag, but there's some really interesting stuff here.

Jane: Yeah, let's talk about it. They compared their model against a few earlier studies — Lin et al., Garcia et al., Dias et al., and Zhou et al. And they looked at precision and recall for each of the three classes.

Tom: So for the N class — the normal beats — their model got a precision of ninety-seven point three five percent and a recall of ninety-four point nine eight percent. That's actually the highest recall among all the methods they compared against. That's really strong.

Jane: And for the V class — ventricular ectopic beats — they got a precision of sixty point six nine percent and recall of eighty-eight point two nine percent. So they're catching most of the V beats, but they're also getting some false positives.

Tom: Right. And then the S class — supraventricular ectopic beats — that's where it gets interesting. They got a precision of sixty-eight point eight three percent and a recall of fifty-four point two five percent. So their precision is actually the highest among all the methods for that class, but the recall is pretty low.

Jane: And that's the classic precision-recall trade-off. You can be very precise — meaning when you say it's an S beat, you're usually right — but you're missing a lot of actual S beats. That's the low recall.

Tom: And the authors acknowledge that. They say the model can predict S class more accurately, but there's still room for improvement in detecting all relevant instances.

Jane: But here's what I find really promising — every single class has precision and recall above fifty percent. That's not true for some of the other methods. For example, Lin et al. had a precision of only thirty-one point six percent for the S class. So the proposed model is more balanced across all classes.

Tom: That's a really good point. And I think that's the main improvement this paper suggests — not just pushing one metric to the moon, but getting a more consistent performance across all classes. That's actually more useful in a clinical setting.

Jane: And they also mention that this approach is more explainable. Because the graph structure is built on correlation coefficients, you can look at which features are connected and understand why the model is making certain decisions.

Tom: Right. And that's a big deal for medical AI. Doctors are more likely to trust a model if they can see the reasoning behind it. So even if the raw numbers aren't perfect, the interpretability is a major step forward.

Jane: And I think that's the real contribution here — not just the accuracy, but the idea that you can build a graph structure that's grounded in statistical relationships and still get competitive results.

Tom: Exactly. So now the big question is — what does this mean for the future? Let's bring in our guests to talk about the implications.

Conclusion: Tom: Alright, we're wrapping up our discussion on "Arrhythmia Classification Using Graph Neural Networks Based on Correlation Matrix." And Jane, I think we can both agree that this paper is a solid step in the right direction.

Jane: Absolutely. They've shown that you can use a graph neural network with a correlation-based adjacency matrix to classify arrhythmias, and they got results that are competitive with — and in some cases better than — existing methods.

Tom: And the key takeaway for me is the interpretability. By building the graph on Pearson correlations, they're making the model more transparent. That's huge for medical applications where trust is everything.

Jane: Right. And they've also shown that the approach is balanced — every class has precision and recall above fifty percent. That's not something every method can claim.

Tom: Now, there are limitations, of course. The S class recall is still low, and they only looked at single-lead ECG. But they mention future work on multi-lead cases, which could help improve performance.

Jane: And I think that's the exciting part — this is just the beginning. The idea of using correlation matrices to build graph structures could be applied to other types of physiological signals too, not just ECG.

Tom: Definitely. And with that, we're going to say goodbye to this paper and get ready to dive into the next one. Thanks for joining us, everyone.

Jane: And a big thank you to our listeners. We'll be back soon with more exciting research from arXiv. Until then, keep your hearts beating steady!

Tom: See you next time!

Seungwoo Han

Tokyo University of Agriculture and Technology

eess.SP, cs.AI

Submitted: 2025-02-10

Updated: 2026-08-18

Comments: Corrected typos

License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/

Importance score: 36/100

Key concepts

Graph Neural Networks (GNN)
A type of neural network used here where data is structured as a graph, consisting of nodes and edges. In this study, features from an ECG are the nodes, and the connections (edges) are determined by how correlated those features are with each other.
Correlation Matrix
A statistical tool used to calculate the Pearson correlation coefficient between every pair of features extracted from an ECG signal. If two features have a high correlation (above zero point nine), an edge is drawn between their corresponding nodes in the graph structure.
Arrhythmia Classification
The process of using machine learning models to identify and categorize irregular heartbeats based on electrocardiogram (ECG) signals. The study focused on classifying beats into three categories: N for non-ectopic, S for supraventricular ectopic, and V for ventricular ectopic beats.

Terminology

Summary

Summary

This paper proposes a method for arrhythmia classification that integrates a Graph Neural Network (GNN) with a linear neural network layer, using a Pearson correlation matrix to construct the graph's adjacency matrix. The authors state: "This study proposes GNN-linear layer fusion based arrhythmia classification method that utilizes Pearson correlation matrices derived from the features of ECG. In this method, features with high correlations are considered. After generating the adjacency matrix, graph features are extracted using GNN, and these are combined with features extracted from a linear neural network to classify arrhythmias."

The dataset used is the MIT-BIH Arrhythmia Database from PhysioNet, which contains 48 records from 47 subjects, each with 30 minutes of 2-lead ECG sampled at 360 Hz, including R peak information. Following the recommendations of the Association for the Advancement of Medical Instrumentation (AAMI), four records were excluded. The remaining records were split into a training set of 22 records (101, 106, 108, 109, 112, 114, 115, 116, 118, 119, 122, 124, 201, 203, 205, 207, 208, 209, 215, 220, 223, 230) and a test set of 22 records (100, 103, 105, 111, 113, 117, 121, 123, 200, 202, 210, 212, 213, 214, 219, 221, 222, 228, 231, 232, 233, 234). Only the major classes were used: non-ectopic beats (N), supraventricular ectopic beats (S), and ventricular ectopic beats (V), and only lead II was considered.

For preprocessing, the authors applied a moving average filter with a window of 72 data points to remove baseline noise, followed by another moving average filter with a 216 data points window. A 12th-order finite impulse response (FIR) filter passing frequencies between 0.5 Hz and 35 Hz was then applied to eliminate high-frequency noise. After filtering, ECG signals were segmented by extracting 86 samples before and 130 samples after each R peak. The Pan-Tompkins algorithm was used to extract PQRST segments. A total of 20 features were extracted per record, including P-R amplitude, P-R time span, S-T time span, Q-R amplitude, Q-R time span, P-Q time span, R-S amplitude, R-S time span, P-T time span, R-T amplitude, R-T time span, Q-T time span, max of each beat, min of each beat, variance of each beat, previous R-present R interval, present R-post R interval, mean of R-R intervals, median of R-R intervals, and root mean square of each beat. This resulted in a training dataset of dimensions (50,557, 20) and a test dataset of dimensions (49,273, 20).

The adjacency matrix was generated by calculating Pearson correlation coefficients for the features. If the correlation coefficient was 0.9 or higher, it was set to 1 (indicating a connection); values below 0.9 were set to 0 (indicating no connection). Edges were then set between nodes based on this connectivity, with node features being the extracted features.

The proposed model architecture takes the feature vector and edge index as inputs, feeding them into both a linear layer and a graph neural layer. The outputs from both networks are concatenated and passed through a linear layer. The graph neural layer uses GraphSAGE. The model uses the rectified linear unit (ReLU) activation function, the Adam optimizer with a learning rate of 0.01, cross-entropy loss, and is trained for 700 epochs.

The results were compared with previous studies in terms of precision and recall. The proposed model achieved a precision of 97.35% for N, 68.83% for S, and 60.69% for V, and a recall of 94.98% for N, 54.25% for S, and 88.29% for V. The authors note: "As the result, proposed model performs quite well in the N and V classes, particularly with a high precision of 97.35% and recall of 94.98% for the N class, which is the high** recall among the methods. In the S class, the proposed model shows the highest precision compared to other methods, but recall is lower, indicating that while the model can predict S class more accurately, there is still room for improvement in detecting all relevant instances."

The authors conclude: "This study designed a GNN-linear layer fusion model based on the correlation coefficient for arrhythmia classification. Our results showed that our model achieved a precision and recall of over 50% across all classes. This suggests the potential of our model to minimize false positives. This approach demonstrates the potential for an explainable model by generating graphs based on high correlation coefficients. Future research will focus on evaluating performance in multi-lead cases to ultimately identify the optimal graph structure."

Improvements for AI systems

Based on the paper, here are the specific improvements I can make to an AI system, and what the improved system can do:

  1. Graph Construction from Feature Correlations

  • Instead of using raw ECG signals or hand-crafted features alone, I will implement a preprocessing pipeline that:

  • Applies the same dual moving-average filter (window sizes 72 and 216) and a 12th-order FIR bandpass (0.5–35 Hz) to denoise the signal.

  • Segments beats around R-peaks (86 samples before, 130 after).

  • Extracts the 20 features listed in Table 2 (e.g., P-R amplitude, Q-T time span, R-R intervals, etc.).

  • Computes a Pearson correlation matrix across these 20 features.

  • Binarizes the matrix with a threshold of 0.9 (1 = edge, 0 = no edge) to create an adjacency matrix.

  • This makes the graph structure explainable—edges directly represent high statistical correlation between physiological features, not arbitrary learned connections.

  1. Hybrid GNN-Linear Architecture

  • I will build a dual-branch model:

  • Branch A: A GraphSAGE layer (with ReLU) that takes node features (the 20 extracted features) and the binarized adjacency matrix as input, producing graph embeddings.

  • Branch B: A linear layer (dense) that processes the same 20 features independently.

  • Concatenate the outputs from both branches, then pass through a final linear layer for classification into N, S, V classes.

  • This fusion allows the model to leverage both relational structure (via GNN) and individual feature importance (via linear layer), improving robustness.

  1. Training Protocol

  • Use the AAMI-recommended inter-patient split (22 training records, 22 test records, excluding 4 records) to avoid data leakage.

  • Train with Adam optimizer (learning rate 0.01), cross-entropy loss, and 700 epochs.

  • This ensures fair comparison with prior work and reproducible performance.

  1. Class-Specific Performance Monitoring

  • I will track precision and recall per class (N, S, V) during validation, not just overall accuracy.

  • The paper shows that the S class has low recall (54.25%)—I will add a class-weighted loss or oversampling for S beats to improve recall without sacrificing precision.

  • Classify arrhythmias from single-lead ECG (lead II) into three major classes (N, S, V) with:

  • Precision: 97.35% (N), 68.83% (S), 60.69% (V)

  • Recall: 94.98% (N), 54.25% (S), 88.29% (V)

  • This outperforms prior methods (e.g., Lin et al., Garcia et al.) in N-class recall and S-class precision.

  • Provide interpretable predictions: For each heartbeat, the system can output which feature pairs (e.g., P-R time span vs. Q-T time span) are correlated above 0.9, making the graph edges visible to clinicians. This aids in understanding why a beat is classified as arrhythmic.

  • Handle new patients without retraining: Because the graph is built from statistical correlations (not patient-specific), the model generalizes to unseen patients in the test set (inter-patient paradigm). This is critical for real-world deployment.

  • Extend to multi-lead or other ECG tasks: The same correlation-based graph generation can be applied to multi-lead data (e.g., 12-lead) by treating each lead as a node, or to other cardiac conditions (e.g., atrial fibrillation, myocardial infarction) by changing the feature set.

  • Reduce false negatives for life-threatening V beats: With 88.29% recall for V class, the system can reliably flag ventricular ectopic beats, which is clinically critical for early intervention.

  • Replace my current CNN/LSTM-based ECG classifier with the hybrid GNN-linear model described above.

  • Add a correlation-matrix computation step in the data loader (using NumPy/PyTorch) to generate adjacency matrices on-the-fly.

  • Implement a custom training loop that logs per-class precision/recall every 10 epochs.

  • If S-class recall remains below 60%, I will apply focal loss or SMOTE on the S beats during training.

This improved system is not just a black-box classifier—it is a clinically interpretable, statistically grounded arrhythmia detector that can be trusted for decision support in hospital settings.

Related papers