China-AIHeart: A Transformer-Based Cardiovascular Disease Risk Prediction Model

The China-AIHeart model, developed by a Chinese research team, adopts the Transformer architecture combined with survival analysis techniques to predict the risk of three events (overall cardiovascular disease (CVD), myocardial infarction (MI), and stroke (STROKE)) in Chinese adults. The project provides a complete toolchain from training to deployment, including open-source code and an online prediction platform, to support clinical decision-making and public health screening.

Cardiovascular disease (CVD) is the leading cause of death globally. Due to population aging and lifestyle changes, the incidence and mortality of CVD in China continue to rise. Traditional models like the Framingham Risk Score perform well in Western populations but have prediction biases when applied to the Chinese population. Deep learning technologies (especially the Transformer architecture) have great potential in medical prediction, and the AIHeart project is a localized attempt to apply this technology to CVD risk prediction for the Chinese population.

AIHeart offers two versions: a simplified version (15 predictors, suitable for primary care screening) and a full version (22 predictors, suitable for high-precision needs of specialist hospitals). The core model uses the Transformer multi-head self-attention mechanism, introducing learnable prompt embeddings and a Top-K selection mechanism to capture feature interactions. Training uses an early stopping strategy, Adam optimizer (learning rate 1e-4), and weight decay regularization to prevent overfitting.

Inputs include demographic characteristics, physiological indicators, lifestyle factors, and medication history, with range limits for continuous variables. The county-level area deprivation index is specifically included as a proxy variable for socioeconomic status. Data preprocessing automatically handles standardization and missing value processing, supporting CSV, XLS, and XLSX format inputs.

The project provides a full-process toolchain:

• Training: Train on custom datasets via train.py and output complete checkpoints;
• Inference: infer.py (PyTorch native, suitable for research) and predict.py (ONNX Runtime, suitable for production);
• Export: onnx.py converts the model to ONNX format and verifies consistency;
• Online service: A web platform is provided (https://www.rjh.com.cn/2018RJPortal/4c/Clinicalpredictionmodels/ChinaAIHeart2/index.shtml) for use without programming.

Clinical level: Assist doctors in objective risk stratification and identify high-risk individuals missed by traditional scores; Public health level: Can be used for large-scale population screening and early intervention to reduce disease burden; Open-source features promote academic validation and improvement, supporting fine-tuning for specific subgroups.

Limitations: Training data is limited to the Chinese population, and cross-ethnic/regional applications require revalidation; The model's interpretability is weaker than traditional logistic regression. Future directions: Enhance interpretability (e.g., attention visualization); Integrate genomic/multi-omics data; Develop CVD subtype-specific models; Explore real-world long-term effects and cost-effectiveness.

China-AIHeart is an important progress in the localized application of medical AI, combining the Transformer architecture with the epidemiological characteristics of the Chinese population and providing a complete solution from research to clinical deployment. For researchers in cardiovascular preventive medicine, clinical epidemiology, and medical AI, it is both a practical toolkit and an example of cutting-edge technology translation.