Lab-Integrated Multi-Aspect Pretraining for Prognostic EHR Modeling
Keywords:
diagnosis prediction, heart failure predictionAbstract
The widespread adoption of electronic health records (EHRs) has created vast datasets containing valuable patient information, yet many machine learning models underutilize single-visit records due to the lack of future outcome labels. This research addresses the challenge of fully leveraging both single-visit and multi-visit patient data by proposing MPLite, a lightweight multi-aspect pretraining framework. The objective is to improve predictive healthcare models by using lab results as auxiliary input features to enhance medical concept representation. The proposed method employs a multi-layer perceptron (MLP) that pretrains on the relationship between lab results and diagnosis codes, allowing integration as a plug-and-play module into various downstream models without modifying their core architectures. Experiments were conducted on the MIMIC-III and MIMIC-IV datasets, focusing on diagnosis prediction and heart failure prediction tasks. Results showed consistent improvements across all tested models, with notable gains in weighted-F1 score, recall, and area under the curve (AUC), demonstrating that MPLite significantly strengthens the predictive performance of baseline models. The findings confirm that incorporating auxiliary features from lab results can effectively address data underutilization and enhance generalizability across diverse predictive tasks. In conclusion, MPLite provides a scalable and efficient solution for advancing predictive modeling in healthcare, offering practical contributions to early intervention strategies and personalized patient care while opening pathways for future research to extend the framework to other data modalities and clinical applications.
References
[2] A. Johnson, L. Bulgarelli, T. Pollard, S. Horng, L. A. Celi, and R. Mark, “MIMIC-IV,” PhysioNet, 2020, [Online]. Available: https://physionet.org/content/mimiciv/1.0/.
[3] E. Choi, M. T. Bahadori, A. Schuetz, W. F. Stewart, and J. Sun, “Doctor AI: Predicting clinical events via recurrent neural networks,” Proc. 1st Machine Learning in Healthcare Conf., vol. 56, pp. 301–318, 2016.
[4] K. Cho, B. van Merrienboer, C. Gulcehre, D. Bahdanau, F. Bougares, H. Schwenk, and Y. Bengio, “Learning phrase representations using RNN encoder-decoder for statistical machine translation,” Proc. EMNLP 2014, pp. 1724–1734, 2014.
[5] T. Bai, S. Zhang, B. L. Egleston, and S. Vucetic, “Interpretable representation learning for healthcare via capturing disease progression through time,” Proc. 24th ACM SIGKDD Int. Conf. Knowledge Discovery & Data Mining (KDD), pp. 43–51, 2018.
[6] E. Choi, M. T. Bahadori, J. Sun, J. Kulas, A. Schuetz, and W. F. Stewart, “RETAIN: An interpretable predictive model for healthcare using reverse time attention mechanism,” Advances in Neural Information Processing Systems (NeurIPS), vol. 29, pp. 3504–3512, 2016.
[7] J. Shang, T. Ma, C. Xiao, and J. Sun, “Pre-training of graph augmented transformers for medication recommendation,” Proc. IJCAI 2019, pp. 5953–5959, 2019, doi: 10.24963/ijcai.2019/827.
[8] L. Rasmy, Y. Xiang, Z. Xie, C. Tao, and D. Zhi, “Med-BERT: Pretrained contextualized embeddings on large-scale structured electronic health records for disease prediction,” npj Digital Medicine, vol. 4, no. 1, pp. 1–13, 2021, doi: 10.1038/s41746-021-00455-y.
[9] F. Ma, R. Chitta, J. Zhou, Q. You, T. Sun, and J. Gao, “Dipole: Diagnosis prediction in healthcare via attention-based bidirectional recurrent neural networks,” Proc. 23rd ACM SIGKDD Int. Conf. Knowledge Discovery & Data Mining (KDD), pp. 1903–1911, 2017, doi: 10.1145/3097983.3098088.
[10] E. Choi, Z. Xu, Y. Li, M. Dusenberry, G. Flores, E. Xue, and A. M. Dai, “Learning the graphical structure of electronic health records with graph convolutional transformer,” Proc. AAAI Conf. Artificial Intelligence, vol. 34, no. 4, pp. 606–613, 2020.
[11] C. Lu, C. K. Reddy, P. Chakraborty, S. Kleinberg, and Y. Ning, “Collaborative graph learning with auxiliary text for temporal event prediction in healthcare,” Proc. IJCAI 2021, pp. 3529–3535, 2021, doi: 10.24963/ijcai.2021/485.
[12] E. Choi, C. Xiao, W. F. Stewart, and J. Sun, “MiME: Multilevel medical embedding of electronic health records for predictive healthcare,” Advances in Neural Information Processing Systems (NeurIPS), vol. 31, pp. 4552–4562, 2018.
[13] P. Nguyen et al., “Deepr: A convolutional net for medical records,” IEEE J. Biomed. Health Informatics, vol. 21, no. 1, pp. 22–30, 2017, doi: 10.1109/JBHI.2016.2633961.
[14] L. Ma et al., “AdaCare: Explainable clinical health status representation learning via scale-adaptive feature extraction and recalibration,” Proc. AAAI Conf. Artificial Intelligence (AAAI), pp. 825–832, 2020.
[15] E. Choi et al., “GRAM: Graph-based attention model for healthcare representation learning,” Proc. 23rd ACM SIGKDD Int. Conf. Knowledge Discovery & Data Mining (KDD), pp. 787–795, 2017.
[16] T. Wu et al., “Leveraging graph-based hierarchical medical entity embedding for healthcare applications,” Scientific Reports, vol. 11, no. 1, p. 5858, 2021, doi: 10.1038/s41598-021-85100-5.
[17] J. Luo et al., “HiTANet: Hierarchical time-aware attention networks for risk prediction on electronic health records,” Proc. 26th ACM SIGKDD Int. Conf. Knowledge Discovery & Data Mining (KDD), pp. 647–656, 2020, doi: 10.1145/3394486.3403094.
[18] C. Lu et al., “Self-supervised graph learning with hyperbolic embedding for temporal health event prediction,” IEEE Trans. Cybern., vol. 53, no. 4, pp. 2124–2136, 2023, doi: 10.1109/TCYB.2022.3146230.
[19] S. Park et al., “Graph-text multimodal pre-training for medical representation learning,” Proc. ACM Conf. Health, Inference, and Learning (CHIL), pp. 261–281, 2022, doi: 10.1145/3536169.3537800.
[20] W. Zhang, B. Ingale, H. Shabir, T. Li, T. Shi, and P. Wang, “Event detection explorer: An interactive tool for event detection exploration,” Proc. 28th Int. Conf. Intelligent User Interfaces (IUI), pp. 171–174, 2023.
[21] W. Zhang, K. Zeng, X. Yang, T. Shi, and P. Wang, “Text-to-ESQ: A two-stage controllable approach for efficient retrieval of vaccine adverse events from NoSQL database,” Proc. 14th ACM Int. Conf. Bioinformatics, Computational Biology, and Health Informatics, pp. 1–10, 2023.
[22] Tian, J., Cui, R., Song, H. et al. Prediction of acute kidney injury in patients with liver cirrhosis using machine learning models: evidence from the MIMIC-III and MIMIC-IV. Int Urol Nephrol 56, 237–247 (2024). doi: s11255-023-03646-6.
[23] C. Lu et al., “Collaborative graph learning with auxiliary text for temporal event prediction in healthcare,” Proc. IJCAI 2021, pp. 3529–3535, 2021, doi: 10.24963/ijcai.2021/485.
[24] S. Park et al., “Graph-text multimodal pre-training for medical representation learning,” Proc. ACM Conf. Health, Inference, and Learning (CHIL), pp. 261–281, 2022, doi: 10.1145/3536169.3537800.
