Large Language Model-informed ECG Dual Attention Network for Heart Failure Risk Prediction

Abstract

Heart failure (HF) poses a significant public health challenge, with a rising global mortality rate. Early detection and prevention of HF could significantly reduce its impact. We introduce a novel methodology for predicting HF risk using 12-lead electrocardiograms (ECGs). We present a novel, lightweight dual attention ECG network designed to capture complex ECG features essential for early HF risk prediction, despite the notable imbalance between low and high-risk groups. This network incorporates a cross-lead attention module and 12 lead-specific temporal attention modules, focusing on cross-lead interactions and each lead's local dynamics. To further alleviate model overfitting, we leverage a large language model (LLM) with a public ECG-Report dataset for pretraining on an ECG-Report alignment task. The network is then fine-tuned for HF risk prediction using two specific cohorts from the UK Biobank study, focusing on patients with hypertension (UKB-HYP) and those who have had a myocardial infarction (UKB-MI). The results reveal that LLM-informed pre-training substantially enhances HF risk prediction in these cohorts. The dual attention design not only improves interpretability but also predictive accuracy, outperforming existing competitive methods with C-index scores of 0.6349 for UKB-HYP and 0.5805 for UKB-MI. This demonstrates our method's potential in advancing HF risk assessment with clinical complex ECG data.

Keywords: Large language model; electrocardiogram; heart failure; interpretable artificial intelligence; multi-modal learning; risk prediction.

Fig. 1

(a) Overview of the ECG dual attention encoder-based risk prediction network (ECG-DAN). A 12-lead ECG recording x is sent to the ECG dual attention encoder, which is capable of simultaneously extracting both cross-lead relationships as well as temporal dynamic patterns within each lead, for better feature aggregation. Then, features from two routes are added and then sent to a max-average pooling layer, producing a flattened feature vector z_ecg. Finally, we employ a multi-layer perceptron (MLP) module to map from a high-dimensional feature space into a risk score (scalar) r for heart failure. (b) Overview of the core attention module used in lead attention and temporal attention modules. See Sec. IIIB for more details.

Fig. 2

Training overview. Our model is (a) first pretrained on the ECG-Report alignment task and the signal reconstruction task on a large-scale public dataset (PTB-XL [39, 40]), and then (b) finetuned on the heart failure risk prediction task with two specific cohorts from the UK Biobank where the future HF event data is available. Here, in PTB-XL dataset, each report has been abstracted to a set of SCP codes with SCP-ECG statement description and confidence score (annotated by human experts). We construct a structured report based on SCP-ECG protocol [39] and then send it to a frozen LLM to extract clinical knowledge for better representation learning guidance. As one ECG may have multiple SCP-code relavant statements, we extract text features separately and then use confidence-based reweighting to aggregate features for feature summation. See below texts for more details.

Fig. 3

Kaplan-Meier risk curves for a) conventional model using a composite of 15 predefined ECG parameters/measurements, and b) the proposed ECG dual attention risk prediction model with the language-informed pretraining. For both models, patients were divided into low- and high-risk groups with a cutoff value referenced from the top 98^th percentile (for UKB-HYP) or top 96^th percentile (UKB-MI) risk scores predicted by the model, reflecting the statistics of the datasets in Table I

Fig. 4

3D visualization of the last 3-dim hidden feature learned in the risk prediction subnetwork along with the visualization of input ECG waves with lowest predicted risk score (dark purple) and highest predicted risk score (light yellow) on the (a) UKB-HYP and (b) UKB-MI datasets.

Fig. 5

Visualization of lead attention patterns and differences between low-risk and high-risk groups across two, different populations: UKB-HYP and UKB-MI.

Fig. 6

Visualization of cross-lead (a,b) and 12-lead temporal attention maps (d,e) obtained from a high HF risk ECG with HYP (a,d) and a high HF risk ECGs with MI (b,e). (c) is a schematic standard ECG for illustrative purpose. Source: Wikimedia Commons.

Fig. 7

U-map visualization of latent code embeddings z^SCP from the large language model using different structured SCP statements. Different colors represent the categorization of statements with disease labels.