SNN for ECG Arrhythmia Classification

End-to-end arrhythmia classification with a Spiking Neural Network (SNN) on MIT-BIH heartbeat data (AAMI 5-class setup).

Project Goal

The objective is to classify each heartbeat into one of five classes:

• N: Normal
• S: Supraventricular
• V: Ventricular
• F: Fusion
• Q: Unknown

The project combines signal preprocessing, spike encoding, and a trainable LIF-based neural model, then evaluates performance with clinically relevant multi-class metrics.

Data and Problem Characteristics

Dataset source:

• MIT-BIH heartbeat CSV splits (mitbih_train.csv, mitbih_test.csv)

Why this task is challenging:

• The dataset is heavily imbalanced (majority class N dominates)
• Minority classes (S, V, F) are harder to learn but important

Class Imbalance

This figure shows the severe skew in class frequencies:

ECG Signal Diversity Across Classes

Representative beats by class:

Average morphology and variability per class:

Preprocessing and Spike Encoding

Processing pipeline:

1. Normalize each beat to [-1, 1]
2. Create train/validation split from training data (stratified)
3. Apply class-weighted loss to reduce imbalance bias
4. Convert normalized ECG samples into spike trains using latency encoding

Latency coding principle:

• Higher amplitude values are converted to earlier spike times
• Time dimension uses 50 simulation steps

ECG to Spikes Illustration

Model Architecture and Modules

Main libraries/modules:

• PyTorch (torch, torch.nn, torch.optim)
• snnTorch (snn.Leaky, surrogate gradients)
• scikit-learn (split and evaluation metrics)
• NumPy, Pandas, Matplotlib, Seaborn

Model used: ECGSpikingNet

Architecture:

• Input layer: beat vector (BEAT_LEN)
• Hidden block 1: Linear(BEAT_LEN, 256) + LIF (beta=0.9) + Dropout (0.3)
• Hidden block 2: Linear(256, 128) + LIF (beta=0.9) + Dropout (0.3)
• Output block: Linear(128, 5) + output LIF neuron
• Surrogate gradient: fast sigmoid (slope=25)

Why this model choice:

• ECG is temporal; LIF neurons model temporal membrane dynamics naturally
• Latency encoding preserves timing structure from waveform amplitude
• Compact architecture is stable to train and avoids over-complexity
• Macro-F1 checkpoint selection promotes balanced performance, not only majority-class accuracy

Training Strategy

• Loss: weighted cross-entropy
• Optimizer: Adam (lr=1e-3, weight_decay=1e-4)
• LR scheduler: cosine annealing (T_max=30)
• Gradient clipping: 1.0
• Best model selected by validation Macro-F1

Training Behavior

Training/validation trends across epochs:

Final Results (Latest Run: run_2)

Metrics source:

• running_history/run_2/results/classification_report.txt

Headline results:

• Accuracy: 79.63%
• Macro F1: 0.5695
• Weighted F1: 0.8295
• Cohen's Kappa: 0.5063

Per-class report:

Class	Precision	Recall	F1-score	Support
N	0.96	0.80	0.87	18118
S	0.25	0.57	0.34	556
V	0.33	0.77	0.47	1448
F	0.18	0.69	0.28	162
Q	0.91	0.86	0.89	1608

Total test samples: 21892

Confusion Matrix

Precision/Recall/F1 by Class

One-vs-Rest ROC Curves

Error Analysis

Examples of wrongly predicted beats (useful for understanding failure modes in minority classes):

Saved Artifacts (run_2)

• Classification report: running_history/run_2/results/classification_report.txt
• Model checkpoint: running_history/run_2/results/snn_ecg_model.pth