Publications

Deep Learning for Stroke Mortality Prediction in eICU: A Dual-Tower Transformer Framework

Zhengrong Jia*, Kwong-Cheong Wong* ★ Best Industrial Paper Award ★ Best Presentation Award

2026 6th International Conference on Computer Communication and Artificial Intelligence (CCAI), Nanjing, China, May 22–24, 2026, pp. 76–81. DOI: 10.1109/CCAI69603.2026.11641950

IEEE Xplore Code Data Pipeline Slides

Abstract

We propose a novel Dual-Tower Transformer (DT-Transformer) for stroke mortality prediction on the multicenter eICU Collaborative Research Database. The decoupled architecture encodes categorical features with learnable embeddings and numerical vitals with a Transformer encoder built on multi-head self-attention, then fuses the two representations for final prediction. The model achieves an AUPRC of 0.6171 — a 14.4% relative improvement over the strongest neural baseline. An Adaptive Runtime Safeguard keeps out-of-range categorical indices from reaching the embedding lookup.

Post-publication Notes

The paper's descriptions differ from the archived implementation in a few places: which features the categorical tower embeds, the feature counts, the scope of the runtime safeguard, and Figure 4, which was included in error and shows illustrative values rather than attention weights from the trained model. The reported metrics are those of the original pipeline and are unchanged. Details: model repository corrections · preprocessing audit.

Baseline Comparison

Model AUROC AUPRC
DT-Transformer (ours) 0.8848 ± 0.0034 0.6171 ± 0.0058
XGBoost 0.8908 0.6467
Random Forest 0.8806 0.6236
Neural Network 0.8582 ± 0.0018 0.5394 ± 0.0054
Standard Transformer 0.8457 ± 0.0129 0.5279 ± 0.0195
Standard MLP 0.8534 ± 0.0058 0.5170 ± 0.0081

Neural model metrics are aggregated over 5 random seeds on a fixed 70/10/20 split, stratified by outcome and keyed on the ICU stay (mean ± std). XGBoost and Random Forest use single-run evaluation.

Architecture & Methodology

  • Dual-Tower Design — Separate encoding pathways for categorical (demographics) and numerical (vitals) features
  • Self-Attention (Numerical Tower) — A two-layer Transformer encoder models interactions among the 91 numerical inputs; the categorical tower uses learnable embeddings
  • Late Fusion — Tower outputs concatenated before classification head
  • Adaptive Runtime Safeguard — Range checks and index clamping for the categorical inputs, so that out-of-range values cannot reach the embedding lookup

Stack: Python · PyTorch · Scikit-learn · XGBoost · Pandas · NumPy

Dataset & Reproducibility

eICU Collaborative Research Database — Multicenter critical care database (200,859 ICU stays across 208 hospitals). Stroke cohort extracted with ICD-9 codes 430, 431, 433, 434, 436, and 437, restricted to the first 72 hours of each ICU stay. Features: 94 inputs, comprising demographics, the hospital identifier, admission height and weight, and aggregated vital signs and laboratory values. Access via PhysioNet credentialed access.

Open-source implementation is split across two repositories. Model training — Dual-Tower-Transformer-eICU-Stroke: PyTorch model, unified training engine covering seven model types, multi-seed experiment runner, Optuna hyperparameter search, and SHAP/LIME explainability. Data pipeline — Data-Preprocessing-for-eICU: raw eICU extraction, ICD-9 stroke cohort selection, first-72-hour aggregation, and stratified train/validation/test splits (by ICU stay in the published version, by patient on the corrected fix branch).

Back to Publications