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

arXiv:2602.05286 · cs.LG, cs.AI · Submitted 2026-02-05 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "HealthMamba: An Uncertainty-aware Spatiotemporal Graph State Space Model for Effective and Reliable Healthcare Facility Visit Prediction".

Jane: Here's the gist of the paper. The gist: HealthMamba proposes an uncertainty-aware spatiotemporal framework that models healthcare facility visit prediction by explicitly incorporating spatial dependencies and providing reliable uncertainty estimates.

Tom: Problem Formulation Healthcare facility visit prediction is critical for optimizing healthcare resource allocation and informing public health policy 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 The first challenge lies in the heterogeneity and sensitivity of healthrelated mobility Second, reliable uncertainty quantification under abnormal situations remains challenging.

Jane: HealthMamba Framework HealthMamba comprises three key components for accurate and reliable healthcare facility visit prediction (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.

Tom: The Unified Spatiotemporal Context Encoder (STCE) produces a unified node-time representation R by fusing heterogeneous inputs such as visit tensor V, static attributes D, and time-varying external factors E with the adjacency matrix A This process involves mapping inputs to a common hidden space via feature embedding layers (Eq.

Title and authors: Tom: The paper summarizes the three main parts of HealthMamba: first, this Unified Spatiotemporal Context Encoder that fuses all those different inputs into one representation.

Jane: That means they take the visit counts, the facility attributes, and external factors like weather information and map them together before anything else starts modeling.

Lu: Then they have this GraphMamba Backbone which is a novel Graph State Space Model designed for hierarchical spatiotemporal modeling, which captures those spatial dependencies as well.

Meng: I’m interested in that hierarchy part. It suggests they aren't just looking at the immediate neighbors but can model dependencies across different scales of spatial interaction and time.

Lalam: And then they have this comprehensive Uncertainty Quantification module that uses three different ways to estimate uncertainty, followed by a post-hoc quantile calibration for better intervals.

Tom: So, the main takeaway is that they aren't just predicting a number; they are predicting a number with an explicit understanding of how uncertain we might be about it.

Jane: It’s about moving from just making a prediction to making a trustworthy one, especially when you need to know the range of possibilities.

The paper's summary: Tom: The authors point out several technical improvements that make this framework work better than what was out there before.

Lu: One major improvement is the introduction of adaptive graph learning within the GraphMamba Block, which derives a node-level embedding and learns node affinities using an attention-style mechanism.

Meng: Adaptive graph learning sounds like it’s trying to make the spatial structure itself dynamic based on what it learns from the data at each step.

Jane: And then they use that learned adjacency matrix to perform graph convolution at every time step, stabilizing it with a residual connection and blending it with a prior adjacency matrix.

Lalam: They also introduce three different uncertainty quantification mechanisms: node-based, distribution-based, and parameter-based uncertainty estimation.

Tom: That third mechanism uses MC Dropout to decompose the predictive mean and total variance into aleatoric and epistemic components, which is really important for understanding the source of the uncertainty.

Jane: And then they finish up with a post-hoc quantile calibration step that adjusts those prediction intervals to hit a target coverage level.

The paper's improvements: Tom: So, wrapping up the HealthMamba paper, they show empirical results on four real-world datasets from California, New York, Texas, and Florida against thirteen state-of-the-art baselines across five metrics.

Jane: They found that HealthMamba improves prediction accuracy by around six point zero percent and uncertainty quantification by about three point five percent compared to the best baseline Mamba based on the average of all four datasets <ref:2602.05286#pg3>.

Lu: The study confirms that removing either the GraphMamba or the Unified Spatiotemporal Context Encoder significantly degrades prediction accuracy, showing how necessary those components are for this approach.

Meng: From a deployment standpoint, they show it can perform well for both short-term, like next-day predictions, and medium-term predictions, which is useful for operational planning.

Lalam: It also shows graceful degradation under different horizons and that HealthMamba is the only method that consistently achieves the target ninety percent coverage across all datasets for both horizons.

Tom: So, this Uncertainty-aware Spatiotemporal Graph State Space Model for Effective and Reliable Healthcare Facility Visit Prediction gives us a much more robust way to handle healthcare demand forecasting in complex environments.

Jane: It supports making decisions under abnormal situations, like public health emergencies or extreme weather events, because we now have calibrated intervals that are actually trustworthy.

Conclusion: Tom: So, we've just walked through HealthMamba: An Uncertainty-aware Spatiotemporal Graph State Space Model for Effective and Reliable Healthcare Facility Visit Prediction, which is essentially a system that uses spatial graphs and uncertainty quantification to predict when people will visit healthcare facilities.

Jane: Right, so the core idea is taking all those different data points—the visit counts, the facility details, even external stuff like weather—and fusing them into one unified representation before modeling how they move in space and time.

Lu: What I found really interesting is that GraphMamba Backbone doesn't just look at neighbors in a simple way; it learns an adaptive graph structure by deriving node embeddings and learning node affinities, which means the spatial dependencies change based on what the model learns.

Meng: From a deployment standpoint, having that hierarchy in the backbone means you can scale it up effectively because it handles different levels of spatial interaction at different times.

Lalam: And to make those predictions useful for real doctors or policymakers, they added three ways to measure uncertainty: node-based, distribution-based, and parameter-based uncertainty.

Tom: That third one using MC Dropout is key; it lets the model tell us the difference between what we know because of the data, that’s aleatoric uncertainty, and what we don't know because our assumptions are shaky.

Jane: And they didn't just stop there; they added a post-hoc quantile calibration step to adjust those prediction intervals so they actually hit the coverage level doctors expect.

Lu: The ablation study really hammered home how important every piece is; removing the GraphMamba or the STCE just crushed the accuracy, dropping it by up to fifteen percent.

Meng: If I were deploying this on a real traffic stream, I’d worry about the computational cost of that adaptive graph learning, but they kept it pretty lean with six fifty seconds per epoch and only one thousand five hundred gigabytes of peak GPU memory.

Tom: So, looking at the results across California, New York, Texas and Florida—they showed a six percent improvement in accuracy and a three point five percent improvement in uncertainty quantification over the best existing models.

Jane: That means HealthMamba is not just making better guesses; it’s giving us much more reliable confidence when we use those predictions to allocate resources.

Lu: It really handles facility type heterogeneity well, meaning it can still provide valid and compact confidence intervals even when predicting different kinds of places like hospitals versus ambulatory clinics.

Lalam: And the best part for me is that it kept accurate predictions even during real-world abnormal scenarios, like the COVID-nineteen lockdown in LA County or Hurricane Hanna in Texas.

Tom: So, to wrap up, this HealthMamba: An Uncertainty-aware Spatiotemporal Graph State Space Model for Effective and Reliable Healthcare Facility Visit Prediction is a solid framework that tackles spatial dependencies and uncertainty quantification together.

Jane: It’s a big step toward making healthcare resource allocation more robust under unpredictable conditions.

Lu: Definitely something worth keeping an eye on if we want to move beyond simple time-series forecasting in this domain.

Meng: For production, it seems like a viable candidate for scenarios where reliability is more important than just the raw prediction number.

Lalam: I’m pretty confident in its uncertainty estimates; they seem well calibrated across the board.

Dahai Yu, Lin Jiang, Rongchao Xu, Guang Wang

Florida State University

cs.LG, cs.AI

Submitted: 2026-02-05

Updated: 2026-10-21

Comments: IJCAI 2026

Code: https://github.com/UFOdestiny/HealthMamba

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 84/100

The gist: The gist: HealthMamba proposes an uncertainty-aware spatiotemporal framework that models healthcare facility visit prediction by explicitly incorporating spatial dependencies and providing reliable

Key concepts

Unified Spatiotemporal Context Encoder (STCE): This component takes different types of information—like visit counts, facility attributes, and external factors—and combines them into a single representation. It maps these diverse inputs into a common hidden space so the model can effectively understand both the spatial layout and the time-varying context of healthcare visits.
GraphMamba Backbone (G-Mamba): This is an innovative structure that uses an adaptive graph learning module within a Mamba architecture. It allows the model to capture complex, hierarchical spatiotemporal dependencies by dynamically learning how different locations are related spatially and how trends evolve over time at various levels of detail.
Uncertainty Quantification (UQ) Module: This module uses three techniques—node-based, distribution-based, and parameter-based uncertainty—to estimate prediction reliability. It provides calibrated confidence intervals that account for both inherent randomness (aleatoric) and model uncertainty (epistemic), ensuring predictions are trustworthy even when data is sparse or abnormal.

Terminology

Uncertainty Quantification — 3× in this paper · explained in 9 episodes
Accuracy — 1× in this paper · explained in 1 episode

Summary

The gist: HealthMamba proposes an uncertainty-aware spatiotemporal framework that models healthcare facility visit prediction by explicitly incorporating spatial dependencies and providing reliable uncertainty estimates.

Problem Formulation

Healthcare facility visit prediction is critical for optimizing healthcare resource allocation and informing public health policy 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 The first challenge lies in the heterogeneity and sensitivity of healthrelated mobility Second, reliable uncertainty quantification under abnormal situations remains challenging.

HealthMamba Framework

HealthMamba comprises three key components for accurate and reliable healthcare facility visit prediction (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.

The Unified Spatiotemporal Context Encoder (STCE) produces a unified node-time representation R by fusing heterogeneous inputs such as visit tensor V, static attributes D, and time-varying external factors E with the adjacency matrix A This process involves mapping inputs to a common hidden space via feature embedding layers (Eq. 5-7) and then applying graph convolution for spatial encoding (Eq. 9) followed by a light-weight temporal mixer for temporal encoding (Eq. 10-12) to produce the output R.

GraphMamba Backbone

The GraphMamba Backbone (G-Mamba) innovatively integrates an adaptive graph learning module into a UNet-style Mamba architecture for hierarchical spatiotemporal modeling The G-Mamba Block captures spatial dependencies, temporal dynamics, and channel interactions through several mechanisms.

Key aspects of the GraphMamba Block include:

Adaptive Graph Learning: This involves deriving a node-level embedding (Eq. 18), learning node affinities using an attention-style mechanism (Eq. 19), calculating the adjacency matrix A(ag) and normalizing it to get Â (Eq. 20, 21), and blending it with a prior adjacency A0 to obtain the final adjacency A∗ (Eq. 22).
Graph-enhanced Spatial Mixing: Using the learned adjacency A∗, the block performs graph convolution at each time step (Eq. 23) and stabilizes learning with a residual connection (Eq. 24).
Temporal and Channel Mixing: Temporal dependencies are modeled via a State-Space Module (SSM) followed by channel mixing using MLPchan (Eq. 25, 26).

The overall backbone follows a UNet-style structure where the encoder applies G-Mamba blocks followed by downsampling, and the decoder performs symmetric upsampling with additional G-Mamba blocks.

Uncertainty Quantification (UQ)

The comprehensive uncertainty quantification (UQ) module integrates three complementary mechanisms to produce calibrated uncertainty estimates These mechanisms are:

Node-based Uncertainty: This produces per-node prediction intervals using pinball loss (Eq. 29) and is robust to misspecification of the full predictive distribution.
Distribution-based Uncertainty: This fits a parametric predictive distribution, such as a heteroscedastic Gaussian predictor outputting mean µi,t+τ and variance σ2i,t+τ (Eq. 30), minimizing the negative log likelihood (NLL) to capture aleatoric uncertainty.
Parameter-based Uncertainty: This estimates epistemic uncertainty using MC Dropout to decompose the predictive mean and total variance into aleatoric and epistemic components (Eq. 32, 33), minimizing the variance of ensemble predictions via loss Lparam (Eq. 34).

The training objective jointly trains these heads using Ltotal = Lquant + Lnll + Lparam + Lcalib. A post-hoc quantile calibration is applied to adjust the quantile regression intervals to achieve the target coverage level by computing an adjustment margin c (Eq. 39) and setting the calibrated interval as [lcal, ucal] = [l - c, u + c].

Empirical Evaluation

HealthMamba was evaluated on four large-scale realworld datasets from California, New York, Texas, and Florida by comparing it with 13 state-of-the-art baselines across five metrics Results show that HealthMamba achieves around 6.0% improvement in prediction accuracy and 3.5% improvement in uncertainty quantification over state-of-the-art baselines.

The evaluation results indicate:

** RQ 1: How does HealthMamba perform compared to state-of-the-art baselines? **

HealthMamba reduces MAE by approximately 6.0% compared to the best baseline Mamba based on the average of all four datasets. It also demonstrates superior performance on prediction reliability, with a 3.5% improvement in IS and reaching the target coverage.

** RQ 2: Is HealthMamba effective for visit prediction of different types of healthcare facilities? **

HealthMamba consistently outperforms Mamba and UrbanGPT across all categories. Notably, HealthMamba maintains valid and compact confidence intervals across all categories, demonstrating robustness to facility-type heterogeneity.

** RQ 4: How does HealthMamba perform under abnormal scenarios such as public emergencies (COVID-19)? **

Under the COVID-19 lockdown in LA County (March 2020) and Hurricane Hanna in Texas (July 2020), HealthMamba successfully maintained accurate predictions, and ground truth values are also within the prediction intervals in most cases.

The ablation study confirms the necessity of all components: removing GraphMamba or STCE significantly degrades MAE (up to 15.782), and omitting any individual UQ module increases MAE and fails the coverage constraint (COV =). HealthMamba consistently achieves the best performance for both short-term (next-day) and medium-term (7-day-ahead) prediction.

The model exhibits graceful degradation under different horizons, with HealthMamba being the only method that consistently achieves the target 90% coverage across all datasets for both horizons. The computational complexity analysis shows that HealthMamba requires 65 seconds per epoch and uses only 1,520 MB of peak GPU memory. Overall, HealthMamba outperforms 13 state-of-the-art baselines by improving prediction accuracy by 6.0% and uncertainty quantification by 3.5% The uncertainty-aware design also enables robust predictions under public emergencies and extreme weather events, supporting trustworthy decision-making under abnormal scenarios.

Improvements for AI systems

Unified Spatiotemporal Context Encoder (STCE) fuses heterogeneous static and dynamic information into compact node-time representations using feature embedding layers to map inputs to a common hidden space. This allows the model to effectively fuse visit counts, demographic attributes, and external factors before spatial encoding begins.
GraphMamba Backbone (G-Mamba) innovatively integrates adaptive graph learning into a UNet-style Mamba architecture for hierarchical spatiotemporal modeling. This structure enables scale-specific modeling capacity by maintaining separate G-Mamba blocks at each level, capturing fine-grained spatiotemporal dependencies.
Comprehensive uncertainty quantification module integrates node-based, distribution-based, and parameter-based uncertainty quantification mechanisms followed by a post hoc quantile calibration to achieve higher prediction reliability. This allows the system to provide calibrated prediction intervals that are robust to misspecification of the full predictive distribution.
The system can perform type specific healthcare facility visit prediction by targeting type specific healthcare facility visit prediction rather than aggregated predictions, explicitly modeling spatial dependencies for different facility types like ambulatory clinics and hospitals.
HealthMamba can provide reliable predictions under abnormal situations such as public health emergencies and extreme weather events by demonstrating superior adaptability to both demand surges and drops under abnormal scenarios.

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.