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

The paper proposes HealthMamba, an uncertainty-aware spatiotemporal graph state space model that integrates a unified context encoder, a novel GraphMamba architecture, and comprehensive uncertainty quantification to significantly improve the accuracy and reliability of healthcare facility visit predictions across four large-scale real-world datasets.

The Gist

Imagine you are the mayor of a large state, and you need to know exactly how many people will visit the hospital, the urgent care clinic, or the nursing home tomorrow. If you guess wrong, you might run out of doctors (causing chaos) or have too many sitting idle (wasting money).

For a long time, computers tried to solve this by looking at time. They said, "Last Tuesday was busy, so next Tuesday will probably be busy too." But this is like trying to predict traffic by only looking at a clock, ignoring that a car accident on a specific road changes everything.

The paper introduces HealthMamba, a new "smart crystal ball" that doesn't just look at time; it looks at space, context, and uncertainty.

Here is how it works, broken down into simple concepts:

1. The Problem: The "Blindfolded" Forecasters

Previous methods had three big flaws:

• They treated everyone the same: They didn't distinguish between a small dentist's office and a massive trauma hospital.
• They ignored geography: They didn't realize that if a flu hits one town, the neighboring town might get sick too. They treated every county as an island.
• They were overconfident: When a hurricane hit or a pandemic started, these models kept giving the same "normal" answers, acting like nothing was wrong. They couldn't say, "I'm not sure, but it might be crazy."

2. The Solution: HealthMamba

HealthMamba is like a super-smart traffic controller for healthcare. It has three main tools in its toolkit:

A. The "Context Chef" (Unified Spatiotemporal Context Encoder)

Imagine you are trying to guess how hungry a crowd will be. You wouldn't just look at the time of day; you'd also look at the weather, the age of the crowd, and if there's a festival nearby.

• What it does: HealthMamba mixes all this "ingredients" together. It looks at static data (like how many old people live in a county) and dynamic data (like a sudden heatwave or a flu outbreak) and blends them into a single, rich picture before making a prediction.

B. The "Neighborhood Watch" (GraphMamba)

This is the heart of the system. Imagine a neighborhood where everyone talks to their neighbors. If one house has a party, the neighbors hear about it.

• What it does: HealthMamba builds a digital map (a graph) connecting all the counties. It uses a special engine called Mamba (which is great at remembering long stories) to watch how diseases or visits spread from one county to another.
• The Magic: It doesn't just use a fixed map. It learns on the fly! If a hurricane hits, the "roads" between counties change (people can't travel), and HealthMamba updates its map instantly to reflect that new reality.

C. The "Honesty Meter" (Uncertainty Quantification)

This is the most important part. Most AI models are like a confident but wrong weatherman who says, "It will be sunny," even when a storm is coming.

• What it does: HealthMamba has three different ways to check its own confidence:
  1. Local Check: "Does this specific town usually have weird numbers?"
  2. Distribution Check: "Does the pattern of numbers look normal, or is it chaotic?"
  3. Memory Check: "If I run this simulation 100 times with slight changes, do I get the same answer?"
• The Result: Instead of just saying "500 people will visit," it says, "We are 90% sure it will be between 450 and 550." If a hurricane hits, the range gets wider, honestly admitting, "We don't know exactly, but it could be huge."

3. The Real-World Test

The researchers tested HealthMamba on data from four huge US states (California, New York, Texas, Florida).

• The Score: It was 6% more accurate than the best existing models.
• The Safety Net: It was 3.5% better at giving honest "confidence ranges."
• The Emergency Test: When they tested it on real disasters (like the COVID-19 lockdowns or Hurricane Hanna), the old models failed completely. HealthMamba, however, adapted. It saw the sudden drop in visits during lockdowns and the sudden spike during the hurricane, and it adjusted its predictions accordingly.

The Bottom Line

Think of HealthMamba as the difference between a weatherman who only looks at yesterday's calendar and a meteorologist who looks at the radar, the wind, the humidity, and admits when a storm is coming.

By understanding that healthcare visits depend on where you are, who you are, and what is happening in the world, HealthMamba helps hospitals and governments prepare for the future without being caught off guard. It doesn't just predict the future; it tells you how sure it is about that prediction.

Technical Summary

1. Problem Definition

The paper addresses the challenge of predicting healthcare facility visits across different facility types (e.g., hospitals, ambulatory clinics, nursing facilities, social assistance) at a regional level (county).

Key Limitations of Existing Works:

• Lack of Granularity: Most prior studies focus on aggregated visit counts, ignoring fine-grained differences between facility types which have distinct visit patterns.
• Neglect of Spatial Dependencies: Existing methods often treat the problem as a pure time-series forecasting task, failing to model the spatial correlations between different regions (e.g., how a surge in one county affects neighbors).
• Unreliability in Emergencies: State-of-the-art models often fail to provide reliable predictions or uncertainty estimates during abnormal events (e.g., pandemics, hurricanes), leading to overconfident and potentially dangerous decision-making.

Goal: To develop an uncertainty-aware spatiotemporal framework that predicts type-specific visits while explicitly modeling spatial dependencies and providing reliable uncertainty quantification (UQ) for robust decision-making.

---

2. Methodology: HealthMamba

HealthMamba is a novel framework comprising three core components:

A. Unified Spatiotemporal Context Encoder (STCE)

This module fuses heterogeneous data sources into compact node-time representations.

• Inputs:
  • Dynamic Visits ($V$): Historical visit counts across $C$ facility types.
  • Static Attributes ($D$): Demographics (e.g., age composition).
  • Dynamic External Factors ($E$): Weather, accessibility, etc.
  • Spatial Structure ($A$): A weighted adjacency matrix based on geographic distance.
• Mechanism:
  • Embeds all inputs into a common hidden space.
  • Applies Graph Convolution to capture spatial dependencies at each time step.
  • Uses a Light-weight Temporal Mixer (combining depthwise 1D convolution and channel mixing) to capture temporal dynamics.
  • Outputs a unified representation $R$ for the backbone.

B. GraphMamba Backbone (G-Mamba)

This is the core prediction engine, innovatively integrating adaptive graph learning into a UNet-style Mamba architecture.

• Architecture: A hierarchical encoder-decoder structure with skip connections.
  • Encoder: Processes features through multiple scales using G-Mamba blocks, followed by downsampling.
  • Decoder: Performs symmetric upsampling and refinement.
• G-Mamba Block:
  • Adaptive Graph Learning: Instead of relying solely on a static distance matrix, the model learns a data-driven, scale-specific adjacency matrix using an attention mechanism. It can blend this learned graph with a prior static graph.
  • Spatial Mixing: Performs graph convolution using the learned adjacency.
  • Temporal/Channel Mixing: Utilizes a State Space Model (SSM) (the core of Mamba) to model long-range temporal dependencies efficiently, followed by channel mixing.
• Advantage: Combines the linear-time efficiency of Mamba with the ability to model complex, evolving spatial relationships.

C. Comprehensive Uncertainty Quantification (UQ) Module

To ensure reliability, HealthMamba integrates three complementary uncertainty mechanisms and a calibration step:

1. Node-based UQ: Learns prediction intervals (quantiles) directly for each node to capture local heteroscedasticity using Pinball Loss.
2. Distribution-based UQ: Fits a parametric distribution (e.g., Heteroscedastic Gaussian) to capture aleatoric uncertainty (data noise) via Negative Log-Likelihood (NLL).
3. Parameter-based UQ: Uses Monte Carlo (MC) Dropout during inference to estimate epistemic uncertainty (model uncertainty) by aggregating stochastic forward passes.
4. Post-hoc Quantile Calibration: A final step on a held-out calibration set to adjust prediction intervals, ensuring they meet the target coverage rate (e.g., 90%).

Training Objective: A joint loss function combining quantile loss, NLL, parameter consistency loss, and a calibration loss to enforce standardized residuals.

---

3. Key Contributions

• Conceptual: Shifts the paradigm from aggregated time-series forecasting to type-specific uncertainty-aware spatiotemporal forecasting, explicitly modeling spatial dependencies.
• Technical: Proposes HealthMamba, the first framework to combine:
  • A unified encoder for heterogeneous health data.
  • A GraphMamba backbone (UNet + Mamba + Adaptive Graph Learning).
  • A tripartite UQ module (Node, Distribution, Parameter) with post-hoc calibration.
• Empirical: Extensive evaluation on four large-scale real-world datasets (CA, NY, TX, FL) demonstrating superior performance over 13 state-of-the-art baselines (including GNNs, Transformers, and LLMs).

---

4. Experimental Results

The model was evaluated on datasets from California, New York, Texas, and Florida covering 40 months of data (Jan 2019 – Apr 2022).

• Prediction Accuracy: HealthMamba achieved a ~6.0% improvement in Mean Absolute Error (MAE) compared to the best baseline (Mamba).
• Uncertainty Quantification: It showed a ~3.5% improvement in Interval Score (IS) and consistently achieved the target 90% coverage, whereas many baselines failed to meet coverage constraints.
• Robustness in Emergencies:
  • COVID-19 (CA): Successfully predicted the sudden drop in visits during lockdowns.
  • Hurricane Hanna (TX): Accurately captured the surge in emergency visits.
  • Baselines struggled to adapt to these abrupt distribution shifts.
• Efficiency:
  • Training Time: ~65 seconds/epoch (comparable to lightweight models, significantly faster than LLM-based methods like UrbanGPT which took ~2000s).
  • Memory: ~1.5 GB GPU usage (lower than many baselines).
  • Parameters: 892K (moderate size, far more efficient than LLMs with millions/billions of parameters).

Ablation Studies: Confirmed that removing any core component (STCE, G-Mamba, or any UQ module) significantly degraded performance, particularly the ability to maintain valid prediction intervals.

---

5. Significance

• Public Health Policy: Provides a reliable tool for resource allocation (e.g., staffing, bed management) by distinguishing between different facility types and predicting demand under normal and crisis conditions.
• Trustworthy AI: The integrated uncertainty quantification prevents "overconfident" predictions during emergencies, allowing policymakers to understand the risk and reliability of the forecast.
• Methodological Advancement: Demonstrates the efficacy of combining State Space Models (Mamba) with Graph Neural Networks for spatiotemporal tasks, offering a scalable and efficient alternative to heavy Transformer or LLM-based approaches.