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

Risk prediction aims to estimate the chance $h(t)$ that a patient will have a certain event in an infinitesimal time interval $D\in[t,t+\Delta t)$: $h(t)=\lim_{\Delta t\rightarrow 0}\frac{p(D\in[t,t+\Delta t)\mid D\geq t)}{\Delta t}$, given that the event has not occurred before. A common approach is the Cox proportional hazards formulation (CoxPH) [17]. The CoxPH model is a statistical model for predicting the time-to-event outcomes based on the assumption that the hazard function $h(t)$ is proportional to a set of features or variables associated with the subject being studied. Mathematically, it can be defined as: $h(t|x)=h_{0}(t)\exp(g_{\theta}(x))$, where $h(t|x)$ models the hazard rate at which events occur at a time $t$ taking the subject information $x$ into account, $h_{0}(t)$ is the baseline hazard function shared by all observations, depending only on the time $t$, and $\exp(g_{\theta}(x))$ is the risk score for observation $x$. Conventional CoxPH models assume that the input observation is a set of covariates (e.g., sex, age, medical history, smoking) and constrain the function $g(\cdot)$ to the linear form: $g_{\theta}(x)=\theta^{\intercal}x=\theta_{1}\times x_{1}+\theta_{2}\times x_{2}+\cdots+\theta_{p}\times x_{p}$ with $\theta$ being the weights to the set of $p$ covariates in $x:{x_{1},x_{2},...,x_{p}}$ [17]. In the existing body of literature pertaining to HF risk prediction, the predominant methodologies employ a predefined set of variables to construct a linear model, subsequently isolating variables that exhibit a high correlation for the purpose of risk prediction. This approach includes parameters such as prolonged QRS duration [6, 7, 8, 9, 10], various conduction disorders (specifically left and right bundle-branch blocks) [11, 12], elongated QT intervals [7], abnormalities in QRS/T angles and T wave patterns [18, 19, 20, 21, 7], and ST-segment depression in the V5 precordial lead [7, 22].

More recently, deep learning approaches such as DeepSurv [23] consider replacing the linear function $g(\cdot)$ with a neural network $f$, which is a deep architecture parameterized by the weights of the network $\theta$. This relaxes the risk model to a non-linear form; also, the input can be high-dimensional without needing to be linearly independent, as the neural network is capable of extracting hierarchical features for risk score estimation in a non-linear fashion. Such an approach has been found to have superior performance against traditional approaches [23] on different risk prediction or survival prediction tasks. For example, researchers have successfully applied neural networks to automatically discover latent features from high-dimensional data for risk/survival prediction, such as whole-slide pathology images [24, 25], 4D cardiac shape plus motion information [26], etc. In this work, we adhere to a similar philosophy and concentrate on creating neural networks designed to autonomously extract features from intricate 12-lead ECG waves, bypassing the reliance on a predefined set of ECG parameters.

Large language models (LLMs) are a class of artificial intelligence (AI) algorithms to understand human language. They can answer questions, provide summaries or translations, and create stories or poems [27]. Recent studies have found that LLMs can be effective in guiding representation learning on image data [28], enabling the knowledge transfer to several downstream vision tasks such as image classification, object detection, and segmentation. In the medical domain, such kind of multi-modal pretraining has been exploited for better understanding of imaging data, such as chest X-rays and magnetic resonance images, to benefit the downstream disease classification and medical image segmentation tasks [29, 30, 31, 32, 33]. Apart from image data, most recently, concurrent works have been made on exploring the connection between natural language and signals (ECG, EEG) for better disease classification [34, 35, 36]. To the best of our knowledge, there is no existing risk prediction work that explores the benefit of combining ECG reports with ECG waves for better representation learning. Compared to disease classification, the task of risk prediction is more challenging due to the lack of event records for training, which amplifies the overfitting issue associated with deep neural networks. Our work provides a promising transfer learning approach to alleviate the need for a large number of event data, by utilizing LLMs and additional public large ECG-report datasets to conduct multi-modal pretraining.

Assume we have a dataset of $N$ triplets $\{\mathbf{x}_{i},\delta_{i},t_{i}\}_{i=1}^{N}$ to record the HF events in a population. Here, $\mathbf{x}_{i}\in\mathbb{R}^{12\times T}$ is a 12-lead ECG signal (I, II, III, AVL, AVR, AVF, V1-6) with a recording length of $T$; $\delta_{i}$ indicates whether there is known date of HF; $t_{i}$ is the number of month to the censoring time if there is no reported HF event during the follow‐up period ($\delta_{i}=0$, right censored) or the number of months until the patient was diagnosed with HF during the follow-up ($\delta_{i}=1$, uncensored). The objective is to have a risk prediction model $f$ parameterized by $\theta$ so that it can predict a patient’s risk of HF given the patient’s ECG data. To this end, we design an ECG dual attention network (ECG-DAN), as shown in Fig. 1(a), where the input is a 12-lead ECG median wave recording $\mathbf{x}$ and the output of the network is a single node $r$, which estimates the risk score $r=f(\mathbf{x};\theta)$.

A signature of ECG-DAN is a dual attention module, which is designed to extract morphological and spatial changes and relationships across different leads as well as the temporal dynamics inside each lead for a more comprehensive understanding of the heart’s electrical activity. We first process each lead signal via a group of 1D convolution layers for noise filtering and feature extraction, which gives $C_{in}$-channel features for each lead at each time point. We then employ a set of $K$ residual blocks with $K$ $2\times$ down-sampling layers along the temporal dimension to extract features at different scales. The output for each lead consists of feature maps with $C_{out}$ channels and a reduced time dimension of $\frac{T}{2^{K}}$: $\mathbf{h}_{i}\in\mathbb{R}^{C_{out}\times\frac{T}{2^{K}}}$. Those features from all 12 leads $\mathbf{h}:\{\mathbf{h}_{1},\mathbf{h}_{2},...,\mathbf{h}_{12}\}$ are then chained and sent to the Lead Attention (LA) block, facilitating the learning of cross-lead interactions to enhance feature aggregation globally. Concurrently, a 12-lead Temporal Attention (TA) component is employed to capture crucial temporal patterns within each lead $\mathbf{h}_{i}$ along the time dimension. Here, for each lead $\mathbf{h}_{i}$, we apply an individual temporal attention module TA${}_{i}$ across the time domain separately, given the fact that different leads correspond to different directions of cardiac activation within three-dimensional space. The outcomes of the two attention modules are added and then summed with a concat-pooling module, producing concatenated flattened features from max-pooling and average pooling operations following common practice from previous ECG analysis work [13]:

With this latent feature $\mathbf{z}_{ecg}$, we then send it to a multi-layer perceptron network (MLP), which consists of three linear layers, reducing the feature by half and then projecting the feature into a 3-dimensional feature space, finally, regressing it to a risk score $r$.

Next, we will explain the two key modules for improved feature learning and model explainability: the lead attention module and the temporal attention, in more detail.

Deep learning-based risk prediction often struggles with small datasets, particularly when events are rare, as in our case where HF events are below $5\%$. To overcome this and prevent overfitting, we adopt a two-stage training approach, see Fig. 2. Initially, we train our network on a large, diverse public ECG dataset. Of note, this dataset does not contain any HF records for risk prediction. After that, we initialize the model with the pretrained parameters and then conduct fine-tuning, focusing on risk prediction for particular populations with documented HF from the UK Biobank study [41]. Subsequently, we fine-tune the pretrained model on the UK Biobank data, specifically targeting HF risk prediction. The pretraining incorporates human-verified ECG reports to align features with clinical knowledge, aiming to improve the model’s ability to discern pathological ECG patterns relevant to HF risk.

In this study, we focus on two populations: patients with hypertension (HYP) and MI, which are highly related to disease progression to HF. Subjects are selected from the UK Biobank dataset (UKB), which is a large-scale biomedical database containing genetic, demographic, and disease information and is regularly updated with comprehensive follow-up studies, from approx. 500,000 subjects. The UKB dataset consists of a large portion of healthy subjects as well as those with a range of cardiovascular and other diseases. To assess the future risk of HF, our analysis is confined to individuals who have electrocardiogram (ECG) together with imaging data available and have not been diagnosed with HF either prior to or at the time of the ECG evaluation.

For ECG signals, we used 12-lead median waveforms (covering a single beat) with a frequency of 500 Hz as input. Each lead is preprocessed with $z$-score normalization and then zero-padded to a length of $1024$. We employ 5 residual blocks ($K=5$, $C_{in}=4$, $C_{out}=16$) to obtain multi-scale features. The dimension of latent ECG feature $\mathbf{z}_{ecg}$ is 512, and the dimension of projected features $\mathbf{e}_{ecg}$, $\mathbf{e}_{text}$ is 128. The default convolutional kernel size is 5. Further network details are provided in the Appendix. We used a batch size of 128 ($n=128$) for model updates at pre-training, with random lead masking for data augmentation [44]. AdamW optimizer [45] with stratified batches was used for training. For fine-tuning in post-MI risk prediction (UKB-MI), we maintained a batch size of 128. Given the low HF incidence ($<2\%$) in UKB-HYP, we increased the batch size to 1024 to include enough uncensored subjects for calculating $\mathcal{L}_{risk}$. Batch-wise dropout was used to stabilize fine-tuning [24]as well as avoid overfitting. The loss function’s weighting parameter ($\alpha$) is set at 0.5. Pre-training and fine-tuning were conducted over 300 and 100 epochs, respectively, to ensure convergence.

We report the concordance index (C-index) [46] as the primary evaluation metric. This metric measures the accuracy of the ranking of predicted time based on the radio of concordant pairs:$\text{ C-index }=\frac{\text{ \# concordant pairs }}{\text{ \# concordant pairs }+\text{ \# discordant pairs }}$. A concordant pair refers to when the predicted time ranking of two subjects aligns with their actual ranking, while a discordant pair means the opposite. A pair $(i,j)$ is concordant if $t_{i}<t_{j}$ and risk scores $r_{i}>r_{j}$, or vice versa. A value of 1 denotes perfect prediction, while a value of 0.5 indicates a prediction quality equivalent to random guessing.

Robust evaluation with multiple two-fold stratified cross-valiation. Due to the scarcity of uncensored subjects (with HF events), using conventional deep learning data splits (e.g, five-fold cross validation or 7:1:2 for training, validation and testing) would result in too few HF events for reliable C-index evaluation. Hence, we opted for a 1:1 training-to-testing ratio with stratified cross-validation, aligned with patient HF status during follow-up. This approach constituted a two-fold stratified cross-validation, which also helps to maintain the same proportion of each class as in the original dataset in each split. To further avoid over-fitting, $20\%$ of the training data was randomly allocated as a validation set for hyper-parameter search (e.g. choice of the learning rate) and model selection. Specifically, we split the dataset into two folds, one for training and validation, the rest for testing. This process is then repeated once more, with the roles of the two parts reversed. To ensure that the result is not biased by the splitting, we perform the above procedure five times, each with a different split across two datasets (UKB-HYP and UKB-MI). The final model performance is reported as the average (with standard deviation) of these 10 trials.

We first compared our proposed pretraining strategy against various pretraining tasks:

• •

  No pretraining (represented by a hyphen), serving as a baseline, with networks trained for 400 epochs (equivalent to the sum of pretraining and finetuning epochs);

• •

  Pretraining on signal reconstruction only;

• •

  Pretraining on signal reconstruction combined with multi-label disease classification (NORM, MI, STTC, CD, hypertrophy), using a cross-entropy loss for classification;

• •

  Pretraining on signal reconstruction and ECG-Report alignment using raw text reports;

• •

  Our proposed method, pretraining on signal reconstruction and ECG-Report alignment with structured report and confidence information, as detailed in Sec. III-C1.

Results are shown in Table II. It is clear that the models pretrained using the proposed strategy (signal reconstruction + ECG2Text alignment (structured w/ confidence)) consistently obtain the highest accuracy among the two tasks, obtaining the average C-index: 0.6349 (0.0156) on UKB-HYP and 0.5805 (0.0580) on UKB-MI.

We further compare model performance using different encoder architectures, including 1) the encoder used in a variational autoencoder (VAE)-like network architecture, which has been found effective for ECG signal feature representation learning in previous works [47, 48, 49, 50]; 2) XResNet1D [13], which is the top performing network architecture on a wide range of different ECG analysis tasks such as ECG disease classification, age regression, and form/rhythm prediction on the public PTB-XL benchmark dataset [39] and ICBEB2018 dataset [51]. We report their performance trained w/o and w/ the proposed language model informed pre-training strategy with structured SCP report in Table III. The proposed ECG attention network has the fewest parameters. Yet, models with this type of network architecture and the proposed pre-training strategy obtain the highest C-index scores on both two tasks.

We further compare our method to the traditional method with well-established ECG parameters. Specifically, we collect a set of ECG parameters from the ECG signals, which have been identified in previous studies for incident HF prediction and mortality estimation [6, 8, 9, 11, 12, 18, 19, 20, 21, 7, 22]. All of these parameters were automatically extracted by the ECG device (using the GE CardioSoft V6 ⁷⁷7https://biobank.ndph.ox.ac.uk/ukb/ukb/docs/CardiosoftFormatECG.pdf) with supporting evidence from previous work: 1) Ventricular rate; 2) Left-axis deviation; 3) Right-axis deviation; 4) Prolonged P-wave duration ($>120~{}\textit{ms}$) [22]; 5) Prolonged PR interval ($>200~{}\textit{ms}$) [22]; 6) Prolonged QRS duration ($>100~{}\textit{ms}$) [22]; 7) Prolonged QT interval ($\geq 460~{}\text{(women)}/450~{}\text{(men)}~{}\textit{ms}$ using the Framingham formula) [52]; 8) Delayed intrinsicoid deflection (DID time) (the maximum value in leads V5 and V6 $>50~{}\textit{ms}$) [53]; 9) Abnormal P-wave axis (values outside the range of $0^{\circ}$ and $75^{\circ}$) [54]; 10) Left ventricular hypertrophy [55]; 11) Abnormal QRS-T angle ($>77^{\circ}~{}\text{(women)}/88^{\circ}~{}\text{(men)}$) [22]; 12) Low QRS voltage [56]; 13) ST-T abnormality [56]; 14) Right bundle‐branch block [56]; 15) Left bundle‐branch block [56].

Following [22], we fit these ECG parameters into the conventional CoxPH linear regression model [17] and use it as a strong competing model for HF risk prediction. Table IV shows the performances using our approach and the traditional ECG parameter-based approach, and Fig. 3 shows Kaplan-Meier plots which depict the survival probability estimates over time, stratified by risk groups defined by each model’s predictions. We further plot the features learned in the last hidden layer ($\in\mathbb{R}^{3}$) of ECGs from the risk prediction branch in the proposed network and representative ECG waves with the lowest/highest risk score in Fig. 4.