HealthMamba: An Uncertainty-aware Spatiotemporal Graph State Space Model for Effective and Reliable Healthcare Facility Visit Prediction

TL;DR

HealthMamba tackles the need for fine-grained, spatially aware healthcare facility visit forecasting with reliable uncertainty estimates. It introduces a three-part architecture: a Unified Spatiotemporal Context Encoder (STCE) for multi-source fusion, a GraphMamba backbone with adaptive graph learning for hierarchical spatiotemporal modeling, and a multi-faceted uncertainty quantification (UQ) module with node-based, distribution-based, and parameter-based components plus post-hoc calibration. Empirical results on four large-scale U.S. state datasets show about $6.0\%$ improvements in MAE and $3.5\%$ improvements in interval-related metrics over strong baselines, with robust performance across facility types and under abnormal events like pandemics and hurricanes. The approach offers a practical, efficient, and reliable framework for resource planning and emergency response, enabling trustworthy decision-making in public health settings.

Abstract

Healthcare facility visit prediction is essential for optimizing healthcare resource allocation and informing public health policy. Despite advanced machine learning methods being employed for better prediction performance, existing works usually formulate this task as a time-series forecasting problem without considering the intrinsic spatial dependencies of different types of healthcare facilities, and they also fail to provide reliable predictions under abnormal situations such as public emergencies. To advance existing research, we propose HealthMamba, an uncertainty-aware spatiotemporal framework for accurate and reliable healthcare facility visit prediction. HealthMamba comprises three key components: (i) a Unified Spatiotemporal Context Encoder that fuses heterogeneous static and dynamic information, (ii) a novel Graph State Space Model called GraphMamba for hierarchical spatiotemporal modeling, and (iii) a comprehensive uncertainty quantification module integrating three uncertainty quantification mechanisms for reliable prediction. We evaluate HealthMamba on four large-scale real-world datasets from California, New York, Texas, and Florida. Results show HealthMamba achieves around 6.0% improvement in prediction accuracy and 3.5% improvement in uncertainty quantification over state-of-the-art baselines.

Figures (5)

• Figure 1: Framework of HealthMamba, which consists of three major components: (i) Unified Spatiotemporal Context Encoder for heterogeneous input fusion, (ii) GraphMamba Backbone for hierarchical spatiotemporal modeling with adaptive graph learning, and (iii) comprehensive uncertainty quantification (UQ) module with three types of UQ mechanisms and post-hoc quantile calibration for reliable prediction.
• Figure 2: Prediction results on the California dataset aggregated in 24 hours, with best baselines selected for visualization.
• Figure 3: Prediction results under extreme events.
• Figure 4: Statistics of the Florida dataset. Red diagonal lines indicate rural counties.
• Figure 5: UQ analyses on the California dataset.