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

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.

• The paper introduces MATA-Former, a transformer that integrates unified semantic embeddings with a dynamic, query-dependent temporal attention mechanism for ICU risk prediction.
• It leverages the SIICU dataset with over 506,000 clinical events and employs Plateau-Gaussian Soft Labeling to transform discrete diagnoses into continuous, multi-horizon regression targets.
• Experimental results demonstrate superior precision and interpretability over traditional models, offering enhanced early warning capabilities for rare critical events.

Semantic-Temporal Alignment for ICU Risk Prediction: The MATA-Former and SIICU Framework

Introduction

This paper presents an advanced framework for ICU risk forecasting using clinical time series, addressing the semantic-temporal gap inherent in heterogeneous medical records. The core innovation is the Medical-semantics Aware Time-ALiBi Transformer (MATA-Former), which leverages unified semantic embeddings and a dynamic, query-dependent temporal attention mechanism. Complementing this, the SIICU dataset is introduced, offering fine-grained, expert-verified annotations across over 506,000 clinical events. The work establishes a paradigm shift from binary classification toward continuous multi-horizon regression using Plateau-Gaussian Soft Labeling (PSL), achieving superior predictive fidelity and interpretability.

SIICU Dataset: Construction and Statistical Characteristics

The SIICU dataset is constructed from the HIS of a tertiary ICU, featuring rigorous privacy protocols and expert-in-the-loop annotation. Clinical trajectories are serialized by timestamp and categorized into 12 event classes, with a collaborative human-AI annotation protocol yielding precise event-level risk mapping.

Figure 1: Event count and statistical heterogeneity in SIICU trajectories, highlighting sequence length variability and risk sparsity.

Statistical analysis reveals heavy-tailed sequence length distributions and highly irregular sampling intervals. Class imbalance and long-tail risk distribution underscore the need for robust modeling able to prioritize rare, high-priority risks.

Unified Semantic Embedding and Problem Formalization

To harmonize unstructured text and structured vitals, the authors serialize all events via templates anchored by clinical context, embedding them using Qwen3-embedding-8B. The resulting $4096$-dimensional vectors preserve semantic specificity, crucial for discriminating context-dependent signals.

Prediction is framed as continuous regression over risk trajectories, with PSL transforming discrete diagnoses into horizon-specific Gaussian targets. This enables supervision on full clinical paths, capturing risk accumulation and remission over multiple temporal scales.

MATA-Former Architecture and Attention Dynamics

The MATA-Former augments standard MHSA with a dynamic Laplacian bias, parameterized by query-specific geometric parameters ($\mu$, $\alpha$). Temporal distances are log-transformed for scale normalization, while attention is causally masked to avoid look-ahead bias.

Figure 2: Overview of the SIICU pipeline and MATA-Former architecture, integrating semantic embeddings and dynamic temporal features via Laplacian-biased attention.

Dynamic parameter initialization covers a continuum of temporal anchors, allowing attention heads to model dependencies from immediate vitals to long-term history. Optimization employs unweighted MSE, validated as robust against zero-inflated label distributions.

Experimental Results and Performance Analysis

Evaluations on SIICU and MIMIC-IV benchmarks demonstrate decisive gains in precision-oriented metrics:

• On SIICU: Sample AUPRC $0.428$ for MATA-Former vs $0.372$ (Transformer) and $0.353$ (TALE-EHR); P@1 of $0.496$ and P@5 of $0.456$.
• On MIMIC-IV: Highest AUPRC ($0.741$) and AUROC ($0.933$) in IMV prediction, and robust performance for Mortality (AUPRC $\mu$0).

MATA-Former consistently maintains superior signal-to-noise isolation, reducing false alarms and offering early warning capabilities unattainable by recurrent or non-time-aware baselines.

Ablation Studies: Temporal Encoding and Loss Objectives

Ablations indicate that traditional additive temporal encoding severely disrupts the semantic manifold, while MATA's attention bias preserves integrity and achieves optimal signal separation. Unweighted MSE outperforms Focal Loss, maintaining a clear separation between noise and true signals across label sparsity extremes.

Figure 3: Temporal encoding strategies and ablation analysis—MATA attention bias optimally preserves semantic signal and model precision.

Interpretability and Parameter Evolution

Layer-wise analysis of attention parameters ($\mu$1, $\mu$2) reveals emergent multi-scale perception: shallow layers span broad temporal horizons, while deeper layers discretize into anchor points corresponding to clinically pivotal events. The magnitude of $\mu$3 relaxes in deeper layers, facilitating integration of chronic and acute information.

Figure 4: Dynamic parameter evolution and aggregated attention alignment, highlighting structural adaptation of attention mechanisms for clinical causality.

Multi-head aggregation functions as a mixture of Laplacians, enabling the model to correlate non-contiguous, sparse clinical events spanning diverse time windows.

Figure 5: Visualization of multi-head attention patterns in Layer 0, illustrating causal masking and diversity of effective receptive fields.

Case studies demonstrate alignment with clinical logic, where the model anticipates DIC risk by sustaining attention on surgery and postoperative biomarkers, validated by subsequent imaging confirmation.

Figure 6: Case study on interpretability and early prediction, evidencing sustained attention on critical clinical events and successful pre-warning capability.

Robustness and Label Sensitivity

Sensitivity analyses confirm that model performance is invariant to label binarization thresholds, indicating robust capture of continuous disease trajectories and avoidance of threshold overfitting.

Figure 7: Sensitivity analysis of label binarization threshold, demonstrating stability of ranking metrics across risk severity definitions.

Limitations

Several constraints are acknowledged:

• Single-center dataset may limit generalization under cross-center distribution shifts.
• Occasional model focus on non-causal noise.
• Instability of semantic manifold, requiring careful vector calibration.
• Lack of integrated high-frequency physiological waveforms or imaging, presenting an information bottleneck.

Conclusion

The MATA-Former framework establishes semantic-temporal alignment as a foundational principle for ICU risk prediction, leveraging continuous regression via PSL to capture fine-grained clinical dynamics. Evaluation confirms substantial gains in ranking and precision metrics, especially for rare risk classes. Interpretability is enhanced via explicit parameterization of temporal horizons. Future work should address modal completeness and cross-institutional robustness to further advance the deployment of interpretable, reliable clinical decision support in critical care.