Transformer-based heart language model with electrocardiogram annotations

Abstract

This paper explores the potential of transformer-based foundation models to detect Atrial Fibrillation (AFIB) in electrocardiogram (ECG) processing, an arrhythmia specified as an irregular heart rhythm without patterns. We construct a language with tokens from heartbeat locations to detect irregular heart rhythms by applying a transformers-based neural network architecture previously used only for building natural language models. Our experiments include 41, 128, 256, and 512 tokens, representing parts of ECG recordings after tokenization. The method consists of training the foundation model with annotated benchmark databases, then finetuning on a much smaller dataset and evaluating different ECG datasets from those used in the finetuning. The best-performing model achieved an F1 score of 93.33 % to detect AFIB in an ECG segment composed of 41 heartbeats by evaluating different training and testing ECG benchmark datasets. The results showed that a foundation model trained on a large data corpus could be finetuned using a much smaller annotated dataset to detect and classify arrhythmia in ECGs. This work paves the way for the transformation of foundation models into invaluable cardiologist assistants soon, opening the possibility of training foundation models with even more data to achieve even better performance scores.

Fig. 1

ECG illustration of (a) NSR, a regular rhythm with patterns, and (b) AFIB rhythm, an irregular rhythm without patterns.

Fig. 2

Durations of AFIB and non-AFIB rhythm episodes in seconds.

Fig. 3

Characteristic features of a heartbeat in ECG.

Fig. 4

ECG strip with annotated RR intervals (in ms), dRR values (in ms), and encoded characters (presented bold) using the mapping from Table 2. Text encoding of the given ECG interval is .

Fig. 5

Flowchart that systematically outlines the entire workflow from taking raw ECG annotations to finetuned AFIB classifier for sequences of 41 heartbeats.

Fig. 6

Duration-based approach to calculate TP, TN, FP, and FN based on time instead of the count of included heartbeats.

Fig. 7

Loss function versus epochs while training the foundation models with different context window sizes (41, 128, 256, and 512 tokens).

Fig. 8

Performance versus training epochs while finetuning the models with a context window of 41, 128, 256, and 512 tokens.

Fig. 9

Performance trends for models with a context window size of 41, 128, 256, and 512 tokens, evaluating TDB as a collection of ECG databases excluding MITDB used for training.

Fig. 10

F1 and F0 scores evaluating the models with a context window size of 41, 128, 256, and 512 tokens, trained on MITDB and evaluated on AFDB, LTAFDB, MITDB, and TDB (consisting of all databases excluding MITDB).

Fig. 11

Performance score Differences between beat and duration-based evaluations.

Fig. 12

An example of an FP error (MITDB record 200).

Fig. 13

An example of a false negative error (MITDB record 210).

Fig. 14

Error rates of our model for all analyzed ECG benchmark databases, including the FDB and TDB.