MATA-Former & SIICU: Semantic Aware Temporal Alignment for High-Fidelity ICU Risk Prediction

Zhichong Zheng; Xiaohang Nie; Xueqi Wang; Yuanjin Zhao; Haitao Zhang; Yichao Tang

cs.LGApr 2, 2026

MATA-Former & SIICU: Semantic Aware Temporal Alignment for High-Fidelity ICU Risk Prediction

Authors: Zhichong Zheng, Xiaohang Nie, Xueqi Wang, Yuanjin Zhao, Haitao Zhang, Yichao Tang

Organizations: Tongji University · Shanghai Innovation Institute · Harbin Institute of Technology · Department of Critical Care Medicine, Shanghai East Hospital, Tongji University School of Medicine.

Abstract

Forecasting evolving clinical risks relies on intrinsic pathological dependencies rather than mere chronological proximity, yet current methods struggle with coarse binary supervision and physical timestamps. To align predictive modeling with clinical logic, we propose the Medical-semantics Aware Time-ALiBi Transformer (MATA-Former), utilizing event semantics to dynamically parameterize attention weights to prioritize causal validity over time lags. Furthermore, we introduce Plateau-Gaussian Soft Labeling (PSL), reformulating binary classification into continuous multi-horizon regression for full-trajectory risk modeling. Evaluated on SIICU -- a newly constructed dataset featuring over 506k events with rigorous expert-verified, fine-grained annotations -- and the MIMIC-IV dataset, our framework demonstrates superior efficacy and robust generalization in capturing risks from text-intensive, irregular clinical time series.

Figures & tables

Panel of Figure 1 — Figure 1: Statistics of SIICU. (a) Heavy-tailed distribution of sequence lengths. (b) Irregular time intervals. (c) Density of positive sample ratios. (d) Long-tail distribution of risks. These distributions highlight the dataset’s high heterogeneity and sparsity.

Figure 2: Overview of the SIICU & MATA-Former framework. (A) The SIICU pipeline transforms heterogeneous clinical records into chronological timelines, utilizing expert-verified collaborative annotation for ground truth generation. (B) The MATA-Former Architecture integrates semantic embeddings with log-transformed temporal features. The core MATA mechanism modulates self-attention via dynamic Laplacian biases, optimized against PSLs for multi-horizon risk prediction.

Model	AUPRC (Sample)	AUPRC (Micro)	AUROC	Brier Score	P@1	P@5
Doctor AI	$0.205\pm 0.010$	$0.187\pm 0.001$	$0.960\pm 0.001$	$0.064\pm 0.001$	$0.331\pm 0.003$	$0.286\pm 0.006$
LSTM	$0.189\pm 0.017$	$0.153\pm 0.027$	$0.936\pm 0.017$	$0.062\pm 0.006$	$0.288\pm 0.033$	$0.264\pm 0.025$
Transformer	$0.372\pm 0.009$	$0.364\pm 0.003$	$0.929\pm 0.003$	$0.006\pm 0.002$	$0.432\pm 0.015$	$0.378\pm 0.009$
HiTANet	$0.336\pm 0.014$	$0.274\pm 0.010$	$0.956\pm 0.002$	$0.006\pm 0.001$	$0.401\pm 0.019$	$0.365\pm 0.014$
EHRMamba	$0.300\pm 0.007$	$0.223\pm 0.008$	$0.912\pm 0.014$	$0.007\pm 0.001$	$0.391\pm 0.009$	$0.367\pm 0.007$
TALE-EHR	$0.353\pm 0.016$	$0.295\pm 0.006$	$\mathbf{0.960\pm 0.002}$	$0.006\pm 0.001$	$0.375\pm 0.004$	$0.333\pm 0.010$

Table 1: Comparative performance on the SIICU dataset. Values denote mean

\pm

standard deviation across four independent runs. Given the extreme label sparsity, evaluation prioritizes AUPRC over AUROC to mitigate the impact of class imbalance. MATA-Former achieving decisive gains in precision-oriented metrics (P@1, P@5) and demonstrating superior sensitivity to rare, high-priority clinical risks.

Model	Sepsis		IMV		Mortality
Model	AUPRC	AUROC	AUPRC	AUROC	AUPRC	AUROC
HiTANet	$0.563\pm 0.018$	$0.861\pm 0.009$	$0.557\pm 0.031$	$0.804\pm 0.023$	$0.395\pm 0.102$	$0.864\pm 0.036$
EHRMamba	$\mathbf{0.753\pm 0.006}$	$\mathbf{0.948\pm 0.020}$	$0.679\pm 0.021$	$0.830\pm 0.021$	$0.385\pm 0.036$	$0.818\pm 0.040$
TALE-EHR	$0.657\pm 0.065$	$0.940\pm 0.011$	$0.652\pm 0.027$	$0.920\pm 0.025$	$0.501\pm 0.111$	$\mathbf{0.902\pm 0.005}$
MATA-Former	$0.709\pm 0.033$	$0.929\pm 0.016$	$\mathbf{0.741\pm 0.037}$	$\mathbf{0.933\pm 0.037}$	$\mathbf{0.523\pm 0.091}$	$0.899\pm 0.003$

Table 2: Performance comparison on the MIMIC-IV v3.1 benchmark across three critical clinical prediction tasks. We report mean ± std over four independent runs. MATA-Former maintains robust stability and generalizability across different datasets

Panel of Figure 3 — Figure 3: Ablation and mechanism analysis. (a) Additive bias superiorly preserves semantic integrity compared to temporal encoding. (b)-(c) MSE maximizes AUPRC by maintaining a distinct signal-noise separation valley. (d) Decoupling $\mu$ and $\alpha$ achieves optimality in capturing heterogeneous clinical dynamics.

Panel of Figure 4 — Figure 4: Visualization of learnable temporal dynamics. (a) Dynamic Parameter Evolution: The distributional shift of Laplacian parameters ( $\mu,\alpha$ ) from Layer 0 to 5 reveals a transition from discrete temporal filtering to high-level semantic abstraction. (b) Aggregated Attention Alignment: To accommodate extended clinical trajectories, the heatmap utilizes $30\times 30$ block-wise average pooling. The resulting vertical strata demonstrate consistent anchoring on pivotal prognostic determinants—specifically abnormal labs and interventional history—validating the structural alignment between the Laplacian prior and clinical causality.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Peak ( $\mu$ )	Slope ( $\alpha$ )	Log Radius ( $X$ )	Physical Range $[T_{min},T_{max}]$	Clinical Semantic Interpretation
1.0	2	2.5	0 sec – 32.2 min	Captures minute-level fluctuations
5.0	2	2.5	11.2 min – 30.1 hours	Captures diurnal trends
9.5	2	2.5	18.2 hours – 113.1 days	Captures seasonal patterns

Appendix

Table 3: Examples of Physical Receptive Fields under different parameter combinations (

\Gamma=5,\tau=60

). These values demonstrate how the interaction of

\mu

and

\alpha

maps to clinically meaningful time windows.

Panel of Figure 5 — Figure 5: Visualization of bias parameters. (Top) No-Slope setting; (Middle) No-Peak setting; (Bottom) Full model configuration. The columns represent $\mu$ Distribution, $\alpha$ Distribution, and Joint Distribution respectively.

Figure 9

Figure 7: Visualization of Multi-Head Attention Patterns in Layer 0. The figure displays the attention weight matrices for all 32 heads. The heatmaps illustrate the strictly causal structure (lower triangular mask) and the diversity of attention distributions, where varying degrees of concentration indicate different effective receptive fields learned by the MATA mechanism.

Key event	Event type	Event content
key1	History	Initial medical record: …Patient, Female, 71 years old… Admission Diagnosis: 1. Valvular Heart Disease, Severe Mitral Regurgitation, Cardiac Function Class III (NYHA); 2. Coronary Artery Disease, Post-CAG+PCI……Grade 3/6 systolic murmur at the apex… …History of Coronary Heart Disease for $>9$ years, oral Aspirin… Clopidogrel… Underwent first coronary stent implantation $>9$ years ago at [Hospital Name Omitted]… underwent coronary stent implantation in [Year]… underwent CAG+PCI surgery on [Year]… Left Anterior Descending branch… severe calcification… in-stent intimal hyperplasia… Right Coronary Artery… implanted Firebird2 $3.5\times 18$ mm stent… …Echocardiography: Severe Mitral Regurgitation… Left atrial and ventricular enlargement… Aortic valve calcification… Left ventricular diastolic function decreased… [Previous] Echocardiography indicated: Massive eccentric mitral regurgitation… Left heart enlargement…
key2	Surgery	Mitral Valve Bio-prosthesis Replacement (Preserving Subvalvular Apparatus)
key3	Clinical Order	Pathogen DNA Detection (KPC)… Pathogen DNA Detection (NDM)… Pathogen DNA Detection (IMP)… Pathogen DNA Detection (VIM)… Pathogen DNA Detection (OXA-48)… Blood Cell Analysis (Three-part differential or above) Determination of C-Reactive Protein (CRP)… Determination of Thrombin Time (TT)… Determination of Plasma Prothrombin Time (PT) Determination of Activated Partial Thromboplastin Time (APTT) Determination of Plasma Fibrinogen Determination of Plasma Antithrombin Activity (ATA) Determination of Fibrin(ogen) Degradation Products (FDP) Determination of Plasma D-Dimer Determination of Thrombin-Antithrombin Complex (TAT) Detection of Plasma Thrombomodulin Antigen (TMAg) Determination of Plasma Plasmin- $\alpha$ 2-plasmin Inhibitor Complex (PAP) Detection of Plasma Tissue Plasminogen Activator Inhibitor Antigen (t-PAI)… Detection of Procalcitonin (Chemiluminescence) Determination of Interleukin-6 (IL-6)…
key4	Lab Test	…Red Blood Cell Count: 2.4 $\times$ $10^{12}$ /L (Reference: 3.8–5.1), Low… Hemoglobin: 71 g/L (Reference: 115–150), Low… Procalcitonin: 0.683 ng/ml (Reference: … $>2$ High risk sepsis), Elevated… Interleukin-6: 261 pg/ml (Reference: 7.0), Elevated… International Normalized Ratio (cobas): 1.16 (Reference: 0.91–1.15), Elevated… Lactate Dehydrogenase: 415 U/L (Reference: 120–246), Elevated… Antithrombin (cobas): 55.5 (Reference: 79.0–113.2), Low… Prothrombin Time (cobas): 11.4 Seconds (Reference: 7.67–10.50), Elevated… Activated Partial Thromboplastin Time (cobas): 56.1 Seconds (Reference: 26.8–42.3), Elevated… Fibrin(ogen) Degradation Products (werfen): 40.6 mg/L (Reference: 0–5), Elevated… D-Dimer (werfen): 7.46 mg/L FEU (Reference: 0–0.5), Elevated… Plasmin- $\alpha$ 2-Plasmin Inhibitor Complex: 6.123 $\mu$ g/mL (Reference: $<0.85$ ), Elevated… Platelet Count: 109 $\times$ $10^{9}$ /L (Reference: 125–350), Low… Thrombin-Antithrombin III Complex: 11.3 ng/mL (Reference: $<4.07$ ), Elevated… High-Sensitivity C-Reactive Protein: 72.86 mg/L (Reference: 0–6), Elevated… Serum Amyloid A: $>350.00$ mg/L (Reference: 0–10), Elevated…
key5	Lab Test	…Neutrophil Count: 10.14 $\times$ $10^{9}$ /L (Reference: 1.8–6.3), Elevated… Procalcitonin: 0.594 ng/ml (Reference: 0–0.046 Normal 0.047–0.50 Indicates low risk sepsis or septic shock $>2$ High risk sepsis), Elevated… Interleukin-6: 66.7 pg/ml (Reference: 7.0), Elevated… Hemoglobin: 72 g/L (Reference: 115–150), Low… Fibrinogen cobas: 4.4 g/L (Reference: 2.08–3.85), Elevated… Activated Partial Thromboplastin Time cobas: 62.5 Seconds (Reference: 26.8–42.3), Elevated… Platelet Count: 98 $\times$ $10^{9}$ /L (Reference: 125–350), Low… Antithrombin cobas: 52.3 (Reference: 79.0–113.2), Low… Prothrombin Time cobas: 16.6 Seconds (Reference: 7.67–10.50), Elevated… International Normalized Ratio cobas: 1.66 (Reference: 0.91–1.15), Elevated… D-Dimer cobas: 1.19 mg/L FEU (Reference: 0–0.510), Elevated… White Blood Cell Count: 11.52 $\times$ $10^{9}$ /L (Reference: 3.5–9.5), Elevated… Thrombin-Antithrombin III Complex: 5.81 ng/mL (Reference: $<4.07$ ), Elevated…
key6	Radiology Report	Abdominal Aortography (Abdominal CTA) (64-slice Contrast-Enhanced CT) …The abdominal aorta and its main visceral branches are visualized… Bilateral renal arteries, iliac arteries, and the celiac trunk show normal course and morphology… The lumens of various vessels show uneven caliber; vessel walls are thickened with visible plaque-like calcifications… A local filling defect is observed in the superior mesenteric artery… Diagnostic Conclusion: Atherosclerotic changes of the abdominal aorta… Local filling defect in the superior mesenteric artery, considering thrombosis… Please correlate with clinical findings…

Appendix

Table 6: Detailed Events and Medical History Records.

Figure 8: Case Study on Interpretability and Early Prediction (DIC). The top panel tracks the predicted risk score (blue) against ground truth labels (orange), highlighting a successful pre-warning window (red box). The bottom panel displays the attention weight distribution, where distinct vertical high-attention bands indicate the model’s mechanism of “locking” onto critical clinical events (Key1–Key6) and sustaining this attention over long durations. This verifies the model’s capability to capture non-contiguous, sparse dependencies in irregular clinical trajectories.

Criterion	Sepsis Prediction	In-Hospital Mortality	IMV Prediction
Target Definition	Sepsis-3 (SOFA shift)	Survival Status (0/1)	First Intubation
Inclusion Criteria	Age $\geq$ 18, History $\geq$ 30h	Age $\geq$ 18, LOS $\geq$ 24h	Age $\geq$ 18, Event $>36$ h
Obs. Window	24h ( $T_{onset}-30$ h to $-6$ h)	First 24h of Admission	24h ( $T_{onset}-30$ h to $-6$ h)
Prediction Gap	6 hours	Remainder of Stay	6 hours
Control Sampling	Pseudo-onset( $T_{pseudo}$ )	All non-survivors	Pseudo-onset( $T_{pseudo}$ )
positive samples	1822	7557	1999

Appendix

Table 7: Cohort selection criteria and experimental design for the three downstream clinical prediction tasks derived from the MIMIC-IV database.

Figure 9: Sensitivity Analysis of Label Binarization Threshold (

\tau

). The figure illustrates the model’s performance stability across a wide range of risk severity definitions (

\tau\in[0.3,0.9]

). While Precision@1 naturally decays as the threshold broadens, the ranking metrics (Global and Sample-wise AUPRC) exhibit negligible fluctuation (

\Delta<0.013

). This threshold invariance confirms that MATA-Former robustly captures the continuous progression of risk rather than overfitting to a specific binary cutoff.

Explore similar work

TA-RNN-Medical-Hybrid: A Time-Aware and Interpretable Framework for Mortality Risk Prediction
Mar 9, 2026Zahra Jafari, Azadeh Zamanifar, Amirfarhad FarhadiMortality Prediction Longitudinal EHR Modeling
Department of Computer Engineering, SRC.,Islamic Azad university, Tehran, Iran · School of Computer Engineering, Iran university of Science and technology, Tehran, Iran
INTERVenE: Temporal-Abstraction-Interval Based Transformers for Short-Horizon Medical Event Prediction
Aug 30, 2026Shahar Oded, Yuval ShaharEvent Time Prediction Interpretable ML
The Stein Faculty of Computer and Information Science, Ben-Gurion University of the Negev, Beer-Sheva, Israel
OC-Distill: Ontology-aware Contrastive Learning with Cross-Modal Distillation for ICU Risk Prediction
Apr 18, 2026Zhongyuan Liang, Junhyung Jo, Hyang-Jung Lee +2Contrastive Learning Cross-Modal Knowledge Distillation
Computational Precision Health UC Berkeley, UCSF · Samsung Advanced Institute of Technology (SAIT) · Computational Precision Health UC Berkeley, UCSF[cs.LG]

Mar 9, 2026J/K move · Enter open · S save

TA-RNN-Medical-Hybrid: A Time-Aware and Interpretable Framework for Mortality Risk Prediction

Zahra Jafari, Azadeh Zamanifar, Amirfarhad Farhadi

Department of Computer Engineering, SRC.,Islamic Azad university, Tehran, Iran · School of Computer Engineering, Iran university of Science and technology, Tehran, Iran

Mortality Prediction Longitudinal EHR Modeling Recurrent Neural Networks Irregular Time-Series Modeling Electronic Health Records Neural Network Interpretability

Open paper AskPDF arXiv

Accurate and interpretable mortality risk prediction in intensive care units (ICUs) remains a critical challenge due to the irregular temporal structure of electronic health records (EHRs), the complexity of longitudinal disease trajectories, and the lack of clinically grounded explanations in many data-driven models. To address these challenges, we propose TA-RNN-Medical-Hybrid, a time-aware and knowledge-enriched deep learning framework that jointly models longitudinal clinical sequences and irregular temporal dynamics through explicit continuous-time encoding, along with standardized medical concept representations. The proposed framework extends time-aware recurrent modeling by integrating explicit continuous-time embeddings that operate independently of visit indexing, SNOMED-based disease representations, and a hierarchical dual-level attention mechanism that captures both visit-level temporal importance and feature/concept-level clinical relevance. This design enables accurate mortality risk estimation while providing transparent and clinically meaningful explanations aligned with established medical knowledge. We evaluate the proposed approach on the MIMIC-III critical care dataset and compare it against strong time-aware and sequential baselines. Experimental results demonstrate that TA-RNN-Medical-Hybrid consistently improves predictive performance in terms of AUC, accuracy, and recall-oriented F

_2

-score. Moreover, qualitative analysis shows that the model effectively decomposes mortality risk across time and clinical concepts, yielding interpretable insights into disease severity, chronicity, and temporal progression. Overall, the proposed framework bridges the gap between predictive accuracy and clinical interpretability, offering a scalable and transparent solution for high-stakes ICU decision support systems.

Abstract

Explore similar work

TA-RNN-Medical-Hybrid: A Time-Aware and Interpretable Framework for Mortality Risk Prediction

INTERVenE: Temporal-Abstraction-Interval Based Transformers for Short-Horizon Medical Event Prediction

OC-Distill: Ontology-aware Contrastive Learning with Cross-Modal Distillation for ICU Risk Prediction