ECG Generation with Flow Matching

GitHub - https://github.com/KarthikSundar2002/Flow-EEG

One Core Problems in using AI in Health is Data Availability. One possible way to mitigate this is Synthetic Data Generation.

As of now, I have trained a DiT based model with Flow Matching Framework. it is trained using the MIT BIH Dataset. It is tiny dataset and I do plan on training this with PTB XL Dataset.

Some of the Decent Results Yet

Pasted image 20251124213205.png
Pasted image 20251124213338.png

After Training in PTB-XL Dataset

I increased the capacity of the DiT model. I also added class-based conditioning and PTB-XL Dataset have 5 classes.
0 - 'CD' - Conduction Disturbance
1 -'HYP - Hypertrophy
2 - 'MI' - Myocardial Infarction
3 - 'NORM' - Normal ECG
4 - 'STTC' - ST/T Change

I trained it on a data from a single lead.

These are the results
ptbxl_generated.png

One of the key things I observed while doing this project was the effect of an Cosine Annealing Learning Rate Schedule.

I was first training with an ReduceLRonPlateau Learning Rate Scheduler. I had a loss curve that plateaued
Pasted image 20251125205906.png

However, then I took that checkpoint, then switched to Cosine Annealing LR Scheduler, I was able to get good results

Pasted image 20251125210033.png

Note the above training starts from the checkpoint from the plateaued training.