Multi-Output Deep Learning for Early Risk Stratification in Myocardial Infarction Using Admission-Only Data
Abstract
Accurate, early risk stratification after myocardial infarction (MI) is critical for individualized clinical management but remains under-explored at the point of hospital admission. Established prognostic tools such as GRACE, TIMI, and ACTION remain clinically interpretable but are constrained by fixed feature sets and linear functional forms that cannot capture high-dimensional, nonlinear feature interactions. Existing machine learning approaches typically address a single outcome in isolation and rely on post-admission laboratory values collected over 24 to 72 hours, limiting real-time applicability at the moment a patient actually arrives. By relying exclusively on admission-time data, the proposed framework enables risk assessment at the earliest stage of care, allowing clinicians to identify high-risk patients promptly and potentially improve outcomes through earlier intervention. We formulate a hierarchical three-output framework predicting (O1) binary in-hospital mortality across the full cohort (N=1,700), (O2) cause of death among the 271 patients who died, and (O3) eleven binary complications, using only the 86 admission-time features retained after filtering variables with >9.5% missingness. A Feature Tokenizer Transformer (FTT) projects each feature into a shared D=192 token embedding, prepends a learnable [CLS] token, and processes the sequence with L=1 transformer block and H=2 attention heads; task contributions are balanced through an uncertainty-weighted multitask loss that learns per-task variances rather than fixing weights by hand (Figure 1, left). We compare this shared-encoder design against a hybrid that disconnects mortality into a separately trained Random Forest while the FTT jointly models cause of death and complications (Figure 1, right), alongside a Parallel Pipeline of independently tuned classical models used as a benchmark. Under stratified 5-fold cross-validation, the Parallel Pipeline achieves O1 AUC 0.896 and O1 F1 0.500 (LightGBM); for cause-of-death prediction it achieves O2 OvR AUC 0.688 and Macro-F1 0.298 (Random Forest); and for complications it achieves O3 macro-AUC 0.703 ± 0.019 (Random Forest), each the best of six benchmarked classical models on a single 80/20 split. The Unified FT-Transformer achieves O1 AUC 0.872 ± 0.030 with F1 0.596 and recall 0.623, the highest death-detection sensitivity of any architecture tested, alongside O2 OvR AUC 0.698 ± 0.051 and O3 micro-AUC 0.759 ± 0.020. The RF+FTT Hybrid trades recall for precision, reaching O1 precision 0.691 with only 47 false positives across 1,429 survivors, the strongest masked cause-of-death discrimination (O2 OvR AUC 0.744 ± 0.083), and O3 micro-AUC 0.756 ± 0.018. Makhmudov et al. [1], who trained on the full 72-hour post-admission feature window, report an O3 micro-AUC of 0.785; our admission-only models close this gap to within 0.026 to 0.029 micro-AUC units despite a strictly harder feature set and stricter 5-fold rather than single-split evaluation. Mortality risk proves largely independent of the complication pathway: the Random Forest matches the jointly-trained FTT on O1 AUC despite having no access to complication context, while cause of death and complications share a coherent feature-discriminability pattern (Spearman ρ ≈ 0.87 between their top-20 importance rankings, even though the two sub-models share no weights). [1] Makhmudov, F., Ravshanov, N., Akhmedov, D., Pekos, O., Turimov, D., and Cho, Y.-I. A multitask deep learning model for predicting myocardial infarction complications. Bioengineering (Basel), 12(5):520, 2025.