Bridging Graph and State-Space Modeling for Intensive Care Unit Length of Stay Prediction

TL;DR

This paper tackles ICU length-of-stay prediction by bridging temporal dynamics and inter-patient relationships. It introduces S^2G-Net, a dual-path architecture that combines a state-space temporal encoder (Mamba) with a multi-view graph encoder (GraphGPS-based) and a static feature branch, fused into a unified representation for regression in the log domain. The authors construct heterogeneous patient graphs from diagnosis-based and semantic features, augmented with MST or GDC, and demonstrate superior predictive performance, efficiency, and calibration on the MIMIC-IV dataset, supported by ablations and interpretability analyses. The work advances clinical prediction by leveraging multi-modal priors and principled graph construction to yield scalable, interpretable ICU LOS estimates with real-world implications for resource management.

Abstract

Predicting a patient's length of stay (LOS) in the intensive care unit (ICU) is a critical task for hospital resource management, yet remains challenging due to the heterogeneous and irregularly sampled nature of electronic health records (EHRs). In this work, we propose S$^2$G-Net, a novel neural architecture that unifies state-space sequence modeling with multi-view Graph Neural Networks (GNNs) for ICU LOS prediction. The temporal path employs Mamba state-space models (SSMs) to capture patient trajectories, while the graph path leverages an optimized GraphGPS backbone, designed to integrate heterogeneous patient similarity graphs derived from diagnostic, administrative, and semantic features. Experiments on the large-scale MIMIC-IV cohort dataset show that S$^2$G-Net consistently outperforms sequence models (BiLSTM, Mamba, Transformer), graph models (classic GNNs, GraphGPS), and hybrid approaches across all primary metrics. Extensive ablation studies and interpretability analyses highlight the complementary contributions of each component of our architecture and underscore the importance of principled graph construction. These results demonstrate that S$^2$G-Net provides an effective and scalable solution for ICU LOS prediction with multi-modal clinical data. The code can be found at https://github.com/ShuqiZi1/S2G-Net.

Figures (8)

• Figure 1: Multi-view graph construction and augmentation pipeline. Graphs based on ICD codes and BERT semantics are merged to form an initial patient graph. Connectivity and global structure are optionally enhanced via MST bridging or GDC. The final graph captures local, semantic, and long-range relations for downstream learning.
• Figure 2: S$^2$G-Net model architecture. It consists of three branches: (i) a temporal encoder based on SSMs, (ii) a multi-view graph encoder capturing patient similarities, and (iii) a flat feature encoder for static attributes. The final representation is formed via weighted fusion of these components.
• Figure 3: Multi-dimensional analysis of S$^2$G-Net and baselines. (a) Training cost vs. $R^2$; (b) Parameter count vs. $R^2$; (c) $R^2$ vs. MSE; (d) Overall ranking by $R^2$. XGBoost excluded for non-comparable computational metrics.
• Figure 4: Ablation study on temporal, feature, and graph structure for $R^2$ and MSE. Dashed lines show the 48h baseline. (a,d) 6/24/48h input windows. (b,e) Removing static features Physiological (Phys), Vital signs (Vital), and Ethnicity (Ethn). (c,f) Edge dropout at 30%, 50%, and 70%.
• Figure 5: Top 12 features by SHAP importance, combining diagnosis-related and static features.