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

AI summary

Overview

Research area: Spatiotemporal machine learning applied to public health — specifically, forecasting daily visits to different types of healthcare facilities at the county level, with calibrated uncertainty estimates.

Technical level: Advanced. The paper assumes familiarity with graph neural networks, state space sequence models (Mamba), quantile regression, and uncertainty quantification (aleatoric vs. epistemic).

One-sentence scope: The paper proposes HealthMamba, a three-part framework combining heterogeneous context encoding, a graph-enhanced Mamba backbone, and a three-mechanism uncertainty quantification module, and evaluates it on four U.S. state-level datasets against 13 baselines.

Authors and venue metadata: Dahai Yu, Lin Jiang, Rongchao Xu, and Guang Wang (corresponding author), all in the Department of Computer Science, Florida State University. Posted under arXiv:2602.05286v3 [cs.LG], 21 May 2026, licensed CC BY 4.0. Code is available at https://github.com/UFOdestiny/HealthMamba, with a longer version at https://arxiv.org/pdf/2602.05286.

What This Paper Is About

Hospitals, clinics, nursing homes, and social assistance providers each see very different demand patterns, and getting those patterns right matters for allocating staff, beds, and funding. The authors argue that prior work has two blind spots: it predicts aggregated visits rather than visits broken out by facility type, and it treats the problem as pure time-series forecasting, ignoring the spatial dependencies between neighboring counties.

The paper's goal is a framework that (1) models spatial structure explicitly across counties, (2) predicts visits separately for four facility categories, and (3) produces prediction intervals that stay trustworthy even when behavior shifts abruptly, such as during a pandemic lockdown or a hurricane.

Key Contributions

1. Problem reformulation. Rather than aggregated, time-series-only forecasting, the paper frames type-specific healthcare facility visit prediction as an uncertainty-aware spatiotemporal forecasting task with explicit spatial modeling, targeting four facility categories defined by NAICS codes 621 (Ambulatory Health Care Services), 622 (Hospitals), 623 (Nursing and Residential Care Facilities), and 624 (Social Assistance).

2. HealthMamba framework with three components. (i) A Unified SpatioTemporal Context Encoder (STCE) that fuses historical visit counts, static demographics, and time-varying external factors (e.g., weather, accessibility) into node–time representations; (ii) a GraphMamba Backbone (G-Mamba) that integrates adaptive graph learning into a UNet-style Mamba architecture for hierarchical spatiotemporal modeling; (iii) a comprehensive uncertainty quantification module combining node-based, distribution-based, and parameter-based mechanisms plus post-hoc quantile calibration.

3. A graph state space model design. Each G-Mamba block learns a data-driven adjacency via attention over pooled node embeddings, symmetrizes and normalizes it, optionally blends it with a prior adjacency through a mixing weight lambda, and then performs graph convolution, state space temporal mixing, and channel mixing with residual connections. Scales do not share weights, allowing scale-specific capacity.

4. Empirical validation at scale. Evaluation on four real-world datasets (California, New York, Texas, Florida) against 13 state-of-the-art baselines across five metrics, reporting around 6.0% better prediction accuracy and 3.5% better uncertainty quantification than the best baseline.

Main Findings

• Overall accuracy gains: HealthMamba reduces MAE by approximately 6.0% compared to the best baseline, Mamba, averaged across the four datasets, and improves the Interval Score (IS) by 3.5% while reaching the target coverage. On the California dataset HealthMamba records MAE 9.150, RMSE 33.401, MPIW 30.970, IS 74.953, and achieves coverage; on New York 3.926 / 14.395 / 14.110 / 29.361 with coverage; on Texas 6.275 / 46.566 / 24.015 / 58.934 with coverage; and on Florida 54.951 / 156.535 / 255.230 / 420.884 with coverage.

• Coverage failures among baselines: Many baselines fail the target coverage of at least 90%, marked with a cross in Table 1 — for example, DCRNN, STGCN, AGCRN, DGCRN, DSTAGNN, ASTGCN, PatchTST, ST-LLM, Mamba, and U-Mamba on California. HealthMamba reaches target coverage on all four datasets.

• Type-specific robustness: Across Hospitals, Ambulatory, Nursing, and Social Assistance categories, HealthMamba outperforms both Mamba and UrbanGPT. Notably, Mamba fails to achieve target coverage in the Nursing and Social categories, while UrbanGPT achieves coverage but with substantially wider prediction intervals.

• Ablation on California: Removing STCE raises MAE to 13.610 and removing G-Mamba raises it to 15.782, both with failed coverage. Dropping any single UQ mechanism degrades performance and fails coverage: w/o Node-based reaches MAE 9.682, w/o Distribution-based 10.154, and w/o Parameter-based 10.573, versus 9.150 for the full model. Removing UQ entirely yields MAE 11.240 and no valid prediction intervals at all (IS and COV are not available).

• Performance under abnormal conditions: The model was tested on the COVID-19 lockdown in LA County in March 2020 (a sudden visit drop) and Hurricane Hanna in Texas in July 2020 (a surge). HealthMamba maintained accurate predictions and ground truth fell within the prediction intervals in most cases, whereas the best baseline struggled to adapt to the abrupt changes.

• Forecast horizons: The paper reports that HealthMamba consistently achieves the best performance for both short-term (next-day) and medium-term (7-day-ahead) prediction; the main comparison table uses a horizon of 3 days.

Methodology in Plain English

The approach starts by turning geography into a graph. Each county is a node, and an edge between two counties gets a weight based on the distance between their centroids, passed through a Gaussian kernel and cut off below a sparsity threshold so only meaningful connections survive. Each node carries a vector of visit counts across the C facility categories at each time step.

Three streams of information are then fused: the historical visit counts, static county-level attributes such as demographics, and time-varying external factors such as weather and accessibility. Each stream is embedded into a shared hidden space and summed. The encoder applies graph convolution to share information between neighboring counties and a lightweight temporal mixer built from depthwise 1D convolution plus channel mixing to share information across time, with layer normalization and residual connections throughout. The output is a compact node–time representation.

The backbone then processes this representation in a UNet-like arrangement: a sequence of G-Mamba blocks with downsampling on the encoder side, symmetric upsampling and skip connections on the decoder side. Inside each block, the model learns its own adjacency rather than relying only on geography — it pools each node's temporal information into a node embedding, computes pairwise attention scores, normalizes them into a learned adjacency, symmetrizes and degree-normalizes it, and can blend it with the distance-based prior graph. Graph convolution, a state space module for temporal dynamics, and channel mixing follow, each with residual connections.

For reliability, three uncertainty heads are trained jointly. The node-based head learns lower and upper quantiles directly with a pinball loss, making no assumption about the shape of the predictive distribution. The distribution-based head outputs a mean and variance per node-time and is trained with negative log-likelihood, capturing noise and heteroscedasticity. The parameter-based head uses MC Dropout with M stochastic forward passes at inference, decomposing total variance into an aleatoric term (the average predicted variance across passes) and an epistemic term (the spread of the predicted means). A calibration loss pushes standardized residuals toward zero mean and unit variance. Finally, a post-hoc step on a held-out calibration set measures the empirical coverage gap and widens or shifts intervals by a computed margin to hit the target coverage.

Why This Matters

Impact on research. The paper pushes spatiotemporal forecasting toward a joint treatment of prediction and reliability, and demonstrates that a Mamba-based state space backbone adapted with adaptive graph learning is competitive with — and here better than — GNN, attention, Transformer, and LLM-based baselines on this task. It also argues that facility-type granularity, not just aggregate demand, is the right unit of analysis for health resource planning.

Real-world applications:

• Public health agencies allocating staff, beds, and outreach funding across hospital, ambulatory, nursing, and social assistance facilities at the county level.
• Emergency preparedness, where planners need not just a point forecast but a defensible interval during pandemics or extreme weather such as Hurricane Hanna.
• Rural health access planning, which the paper motivates with the observation that rural areas have fewer facilities and residents travel longer distances, captured in the Florida dataset statistics on facility counts and travel miles per visit.
• Insurance and provider-network planning, where type-specific demand shifts inform where capacity or coverage gaps will emerge.

Industry relevance. Any organization that must pre-position resources against uncertain demand — hospital systems, state health departments, insurers, and logistics providers supporting care delivery — benefits from forecasts that come with calibrated intervals rather than bare point estimates. The structured three-mechanism uncertainty design and the post-hoc calibration step are also transferable to other spatiotemporal forecasting settings beyond healthcare.

Future Directions

• Closing the gap between reliability and sharpness. UrbanGPT achieves coverage but with much wider intervals; future work could tighten intervals while preserving coverage, particularly in the Nursing and Social Assistance categories where baselines fail coverage outright.
• Extending beyond the study window and geography. The data spans January 2019 to April 2022 across four U.S. states; whether the framework transfers to other countries, data regimes, or more recent periods is not addressed in the content provided.
• Scaling the graph and richer event modeling. The paper tests two anomalous events (the March 2020 LA County COVID-19 lockdown and July 2020 Hurricane Hanna in Texas); a systematic treatment of many event types, including their lead time and severity, is a natural extension.
• Reducing inference cost. The parameter-based UQ mechanism requires M stochastic forward passes at inference, which trades compute for epistemic uncertainty; cheaper alternatives or adaptive M are open questions.

Target Audience

This paper is most valuable to graduate students and researchers working on spatiotemporal forecasting, graph neural networks, state space models, or uncertainty quantification, as well as applied machine learning practitioners in health analytics and urban computing. Public health planners and policy analysts interested in how reliable forecasts can inform resource allocation will find the framing and evaluation useful, though the methodology sections assume comfort with mathematical notation for graphs, quantile losses, and variance decompositions. Readers seeking a practitioner-level introduction should start with the introduction, problem formulation, and the RQ1 and RQ4 results.

Note on completeness: the provided paper content is truncated within Appendix A (the description of Figure 4, "Statistics of the Florida dataset," is cut off), and certain setup details such as hyperparameter values, training configuration, data split sizes, and the fifth research question referenced in the abstract's evaluation section are not reported in the available content. The paper's introduction states five research questions but the evaluation section enumerates four (RQ1–RQ4).

Authors’ 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.