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

Dahai Yu; Lin Jiang; Rongchao Xu; Guang Wang

cs.LGFeb 5, 2026

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

Authors: Dahai Yu, Lin Jiang, Rongchao Xu, Guang Wang

Organizations: Department of Computer Science, Florida State University

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 & tables

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.

Panel of Figure 2 — Figure 2: Prediction results on the California dataset aggregated in 24 hours, with best baselines selected for visualization.

Category	Model	MAE $\downarrow$	RMSE $\downarrow$	MPIW $\downarrow$	IS $\downarrow$	COV
Hospitals	Mamba	12.384	42.156	38.720	95.832	✓
	UrbanGPT	12.058	45.623	52.418	101.254	✓
	HealthMamba	11.628	39.524	35.105	91.250	✓
Ambulatory	Mamba	8.152	31.478	29.640	68.425	✓
	UrbanGPT	8.536	34.127	42.815	73.680	✓
	HealthMamba	7.683	29.352	27.218	65.127	✓

Table 2: Fine-grained prediction on four healthcare facility types.

Variants	MAE $\downarrow$	RMSE $\downarrow$	MPIW $\downarrow$	IS $\downarrow$	COV
w/o STCE	13.610	48.535	39.807	108.110	✗
w/o G-Mamba	15.782	49.092	42.625	117.382	✗
w/o Node-based	9.682	34.925	32.540	80.115	✗
w/o Distribution-based	10.154	36.218	34.112	85.330	✗
w/o Parameter-based	10.573	37.665	35.890	90.224	✗
w/o UQ (Total)	11.240	39.150	37.520	-	-

Table 3: Ablation study on the California dataset.

Panel of Figure 3 — Figure 3: Prediction results under extreme events.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 4 — Figure 4: UQ analyses on the California dataset.

Spatiotemporal Forecasting

Spatiotemporal forecasting predicts how quantities change across locations and time, from traffic and weather to disease and environmental systems.

See all 423 papers

EnergyMamba: An Uncertainty-Aware Graph-Enhanced Selective State Space Model for Energy Consumption Prediction
May 30, 2026Dahai Yu, Rongchao Xu, Lin Jiang +1Selective SSMs Time Series Forecasting
Florida State University · Tallahassee, Florida, USA
Spatio-temporal stochastic graph-based learning for infectious disease forecasting
May 28, 2026Luz Stefani Sotomayor Valenzuela, Susanna Cramb, Darren WraithEpidemic Forecasting Probabilistic Forecasting
School of Public Health and Social Work, Queensland University of Technology, Brisbane, Australia · QUT Centre for Data Science, Queensland University of Technology, Brisbane, Australia
Beyond Time Series: Spatial Reasoning for Epidemic Forecasting via Multimodal Learning
Jun 20, 2026Diana Guadalupe Gomez, Chenwei Wu, Zhiyi Wang +2Epidemic Forecasting Spatiotemporal Reasoning
University of Michigan Ann Arbor, MI, USA

May 30, 2026J/K move · Enter open · S save

EnergyMamba: An Uncertainty-Aware Graph-Enhanced Selective State Space Model for Energy Consumption Prediction

Dahai Yu, Rongchao Xu, Lin Jiang, Guang Wang

Florida State University · Tallahassee, Florida, USA

Selective SSMs Time Series Forecasting Probabilistic Forecasting Conformal Prediction Spatiotemporal Forecasting Electricity Load Forecasting+1

Open paper AskPDF arXiv

Energy consumption prediction is essential for efficient grid management, demand-side optimization, and sustainable energy planning. Although advanced machine learning methods have been employed for better prediction performance, existing works have two key limitations: (1) they usually formulate this task as a purely time-series prediction problem without explicitly modeling the spatial dependencies among different regions, and (2) they fail to provide reliable predictions with uncertainty estimates under abnormal situations such as extreme weather events. To advance existing research, we propose EnergyMamba, an uncertainty-aware spatiotemporal learning framework for accurate and reliable energy consumption prediction, which comprises two key components: (i) a novel Graph-Enhanced Selective State Space Model (GE-Mamba) that injects spatial context learned from the grid topology into the temporal dynamics, enabling coupled spatiotemporal modeling, and (ii) an Adaptive Sequential Conformalized Quantile Regression (AS-CQR) module, which includes locally adaptive normalization and an online feedback mechanism to dynamically calibrate prediction intervals under potential distribution shifts. We evaluate EnergyMamba on four large-scale real-world datasets from Florida, New York, and California. Results show EnergyMamba achieves around 5% improvement in prediction accuracy and 6% improvement in uncertainty quantification over 15 state-of-the-art baselines.

Abstract

Spatiotemporal Forecasting

Explore similar work

EnergyMamba: An Uncertainty-Aware Graph-Enhanced Selective State Space Model for Energy Consumption Prediction

Spatio-temporal stochastic graph-based learning for infectious disease forecasting

Beyond Time Series: Spatial Reasoning for Epidemic Forecasting via Multimodal Learning