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

1 Introduction

Healthcare facilities encompass a wide range of locations providing healthcare services, from small clinics and physicians' offices to urgent care centers and large hospitals. They are indispensable to population well-being for preventing and treating disease, enabling continuous care, and forming the backbone of resilient public health systems. Consequently, predicting population visits to healthcare facilities is critical for resource allocation and policy planning [Marcusson et al., 2020; Xu et al., 2026a; Xu et al., 2026b].

∗Prof. Guang Wang is the corresponding author.

Recent studies on have made progress by using statistical methods [Pun et al., 2019; Alghamdi et al., 2019], deep [Piccialli et al., 2020; Zhong et al., 2026],

and hybrid approaches [Jiang et al., 2025a; Xu et al., 2025; Xiao and Liu, 2025]. However, there are three key limitations of existing work. First, prior studies predominantly focus on aggregated visit prediction, overlooking fine-grained healthcare facility types such as ambulatory clinics, hospitals, nursing facilities, and social assistance providers. Second, most existing work formulates healthcare facility visit prediction as a purely time-series forecasting problem, neglecting spatial dependencies that could provide critical information for accurate prediction. Third, most state-of-the-art approaches fail to deliver reliable predictions under public health emergencies and extreme weather events, such as the COVID-19 pandemic and hurricanes.

Hence, in this paper, we aim to develop an uncertaintyaware spatiotemporal framework for type-specific healthcare facility visit prediction that explicitly models spatial dependencies while providing reliable uncertainty estimates. However, there are two key challenges to achieving this. The first challenge lies in the heterogeneity and sensitivity of healthrelated mobility. Our data analysis indicates that healthcare facility visits exhibit a strong correlation with population structure (e.g., age composition), spatial accessibility, and service availability. These characteristics lead to highly uneven visit patterns across regions and facility types, making it difficult for existing models to capture fine-grained spatiotemporal heterogeneity. Second, reliable uncertainty quantification under abnormal situations remains challenging. During public emergencies such as pandemics, healthcare facility visit behaviors and demand distributions shift dramatically. Without explicit uncertainty modeling, predictions become overconfident and unreliable, potentially leading to misguided decision-making.

To address the above challenges and advance existing research, we propose HealthMamba, which comprises three key components: (i) a Unified SpatioTemporal Context Encoder (STCE), which fuses heterogeneous static and time-varying contextual information and historical visit data into compact node-time representations, (ii) a novel GraphMamba Backbone (G-Mamba), which innovatively integrates adaptive graph learning into a UNet-style Mamba architecture for hierarchical spatiotemporal modeling, and (iii)

a comprehensive uncertainty quantification (UQ) module integrating node-based, distribution-based, and parameterbased UQ mechanisms, followed by a post-hoc quantile calibration to achieve higher prediction reliability.

The key contributions of this paper are as follows:

• Conceptually, unlike existing works that focus on aggregated healthcare facility visit prediction and treat this task as a time-series forecasting problem, we target typespecific healthcare facility visit prediction and formulate it as an uncertainty-aware spatiotemporal forecasting task with explicit spatial modeling.
• Technically, we propose HealthMamba, an uncertaintyaware spatiotemporal framework comprising three novel components: (i) a Unified Spatiotemporal Context Encoder for heterogeneous information fusion, (ii) a graph state space model called GraphMamba for hierarchical spatiotemporal modeling, and (iii) a comprehensive UQ module that integrates three different UQ mechanisms and a quantile calibration for reliable prediction.
• Empirically, we evaluate HealthMamba on four realworld datasets from California, New York, Texas, and Florida by comparing it with 13 state-of-the-art baselines across five metrics. Extensive results demonstrate that HealthMamba achieves 6.0% higher prediction accuracy and 3.5% better uncertainty quantification than the best baseline. Our implementation is available at https://github.com/UFOdestiny/HealthMamba. A longer version of this paper can be found at https: //arxiv.org/pdf/2602.05286.

2 ProblemFormulation

Based on data analysis using real-world data, we found that there is a disparity in healthcare facility distribution and accessibility among different regions. For example, there are fewer healthcare facilities in rural areas, and rural residents need to travel longer distances to visit healthcare facilities. In addition, the visit patterns of different types of healthcare facilities are also different. Hence, we intuitively build graphs to model spatial dependencies to predict multi-type healthcare facility visits.

2.1 Spatiotemporal Graph Construction

We define a spatiotemporal graph G = (V, E,A), where V = {v1, . . . , vN} denotes N nodes (e.g., counties), E is the edge set, and A ∈ RN×N is the weighted adjacency matrix. At each time step t, node i has a feature vector vi,t ∈ RC representing visit counts across C healthcare facility categories. We construct A using a Gaussian kernel with sparsity thresholding:

where dij is the centroid distance between nodes vi and vj, σ > 0 controls spatial decay, and ϵ ∈ [0, 1] is a sparsity threshold that removes weak connections.

2.2 Healthcare Facility Visit Prediction

Let V ∈ RN×Tin×C denote the input visit tensor including features of all N nodes over Tin historical time steps across C healthcare facility types. In addition to historical visits, we also incorporate static attributes D ∈ RN×ddem (e.g., demographics) and time-varying external factors E ∈ RN×Tin×dext (e.g., weather, accessibility) as auxiliary inputs, where ddem and dext denote the dimensions of demographic attributes and external factors, respectively. The goal of healthcare facility visit prediction is to predict future visits for the next Tout steps given historical observations and auxiliary information.

Deterministic Prediction

Conventional deterministic prediction learns a mapping f from inputs to point forecasts:

Probabilistic Prediction

In this work, we focus on probabilistic prediction, which quantifies uncertainty using prediction intervals. Given a target miscoverage rate α (e.g., α = 0.1 for 90% coverage), we predict lower, upper, and median quantiles:

where ˆL = ˆ qα/2(V), ˆU = ˆ q1-α/2(V), and ˆM= ˆ q0.5(V) represent the lower bound, upper bound, and median predictions, respectively, all in RN×Tout×C . The median ˆMalso serves as the point prediction, analogous to ˆY in Eq. (2).

3 Methodology

In this paper, we propose HealthMamba, a spatiotemporal framework for effective and reliable healthcare facility visit prediction. An overview of HealthMamba is shown in Figure 1, which comprises three key components: (i) a Unified SpatioTemporal Context Encoder (STCE), which fuses heterogeneous contextual information and historical visit data into compact node-time representations, (ii) a novel GraphMamba Backbone (G-Mamba), which innovatively integrates adaptive graph learning into a UNet-style Mamba architecture for hierarchical spatiotemporal modeling, and (iii) a comprehensive uncertainty quantification module integrating node-based, distribution-based, and parameter-based uncertainty quantification, followed by a post-hoc quantile calibration to achieve higher prediction reliability.

3.1 Unified Spatiotemporal Context Encoder

Given the input visit tensor V, static attributes D, dynamic external factors E, and the adjacency matrix A as defined in Section 2, STCE produces a unified node-time representation R ∈ RN×Tin×dmodel that fuses these heterogeneous inputs while considering spatial structure and temporal dynamics:

We utilize dhid and dmodel to denote the hidden-layer dimension and model output embedding dimension, respectively.

Specifically, we first map inputs to a common hidden space via feature embedding layers:

where φ∗ denotes a nonlinear function (e.g., SiLU/GELU) with dropout. The initial fused embedding at each (i, t) is:

For spatial encoding, we apply graph convolution with activation function ReLU on node features X∶,t at each time t:

For temporal encoding, we mix information across the Tin steps for each node i using a light-weight temporal mixer combining depthwise 1D convolution and channel mixing:

where DepthConv1D mixes information across the temporal axis and MLPchan performs channel mixing. LayerNorm and residual connections stabilize training. The final projection is:

The encoder output R ∈ RN×Tin×dmodel serves as the input to the GraphMamba backbone.

3.2 GraphMamba Backbone

In this part, we innovatively integrate an adaptive graph learning module into a UNet-style Mamba architecture for hierarchical spatiotemporal modeling. The input feature at scale s is X(s) ∈ RN×Tin×ds , with X(0) = R. The encoder applies a sequence of G-Mamba blocks followed by downsampling; the decoder performs symmetric upsampling followed by additional G-Mamba blocks. At the coarsest scale, we initialize the decoder by setting Z(S) = Y(S) . For other scales, we use the following equations.

Here, DownConv denotes stride convolution, UpConv denotes transposed convolution, and Fuse represents skipconnection fusion. Each scale maintains its own G-Mamba block without shared weights, allowing scale-specific modeling capacity.

G-Mamba Block

Each GraphMamba Block (G-Mamba) is designed to capture spatial dependencies, temporal dynamics, and channel interactions, incorporating an Adaptive Graph Learning mechanism that learns scale-specific, data-driven graph structures.

Adaptive Graph Learning. We first derive a node-level embedding by aggregating temporal information:

Node affinities are learned using an attention-style mechanism:

The resulting adjacency is symmetrized and normalized:

where ˜D is the degree matrix of A(sym) . If prior adjacency A0 is available, we blend the two sources:

Graph-enhanced Spatial Mixing. Using the learned adjacency A∗ , the block performs graph convolution at each time step:

A residual connection stabilizes learning:

Temporal and Channel Mixing. Temporal dependencies are modeled via a State-Space Module (SSM), followed by channel mixing:

Projection and Residual Refinement. The block output is produced by projecting back to the original dimension with a residual connection:

Multi-scale Fusion and Skip Connections. During decoding, the upsampled representation at scale s is fused with the encoder's corresponding feature Y(s) via concatenation and projection:

The fused tensor is passed through another G-Mamba block, enabling joint refinement of encoder and decoder information.

3.3 Uncertainty-aware Prediction

We design a modular uncertainty quantification component atop the GraphMamba backbone that produces calibrated uncertainty estimates. This component comprises three complementary mechanisms, node-based, distribution-based, and parameter-based uncertainty quantification, which together capture local heteroscedasticity, predictive distributional shape, and epistemic uncertainty.

Node-based Uncertainty

Node-based uncertainty produces per-node prediction intervals expressing local uncertainty at the region level. For a prediction target yi,t+τ at node i and horizon τ, we learn lower and upper quantile heads ℓα(i, t, τ ) and uα(i, t, τ ) for a nominal miscoverage level α (e.g., α = 0.1 for a 90% interval).

Each quantile head is implemented as a projection from the backbone's final hierarchical representation Z. The pinball loss for a set Q of quantile levels is:

where ρq(z) = max(qz, (q - 1)z) is the pinball loss, ˆ qi,t+τ denotes the predicted q-quantile, and S is the training set. Node-based intervals are robust to misspecification of the full predictive distribution, as they directly target coverage without assuming any parametric form.

Distribution-based Uncertainty

Distribution-based uncertainty fits a parametric predictive distribution (e.g., Gaussian or Laplace) per node-time. The most common instantiation is a heteroscedastic Gaussian predictor outputting mean μi,t+τ and variance σ2 i,t+τ :

This head is trained by minimizing the negative log likelihood (NLL):

Distribution-based heads capture aleatoric uncertainty (data noise and heteroscedasticity) and allow analytic computation of prediction intervals, e.g., μ ± z1-α/2σ for Gaussian.

Parameter-based Uncertainty

Parameter-based uncertainty estimates epistemic uncertainty arising from model parameters. We adopt the commonlyused MC Dropout, which performs M stochastic forward passes at inference with dropout enabled. Given M stochastic predictions {μ(m) i,t+τ }Mm=1 and {σ2,(m) i,t+τ }Mm=1 , the predictive mean and total variance decompose as:

This decomposition yields interval estimates that combine both sources of aleatoric and epistemic uncertainty. The parameter-based loss minimizes the variance of ensemble predictions to encourage consistency, i.e., the model should produce similar predictions across different dropout masks when it is confident about a given input:

A lower value of Lparam indicates that the stochastic predictions are tightly clustered around their mean, reflecting lower epistemic uncertainty.

Integrated Training Objective

We jointly train all the above uncertainty heads using:

where Lquant, Lnll, and Lparam correspond to node-based, distribution-based, and parameter-based uncertainty losses, respectively. The calibration loss Lcalib encourages standardized residuals to match a standard normal distribution. Over a training minibatch B, we compute

which enforces the standardized residuals to have zero mean and unit variance, thereby improving probabilistic calibration. Here, the standardized residual is defined as

where ε > 0 is a small constant for numerical stability.

Post-hoc Quantile Calibration

To improve the empirical coverage of the trained quantile intervals, we apply a post-hoc quantile calibration on a held-out calibration set. Let (ℓi, ui) be the predicted interval and yi be the ground truth, we compute the empirical coverage gap:

where C is the calibration set and α is the target miscoverage level. We then compute the adjustment margin:

The calibrated interval for a new prediction becomes

which adjusts the quantile regression intervals to achieve the target coverage level.

4 Evaluation

In this section, we conduct a comprehensive experimental evaluation of our proposed HealthMamba. Specifically, we aim to address the following five research questions:

• RQ 1: How does HealthMamba perform compared to state-of-the-art baselines?
• RQ 2: Is HealthMamba effective for visit prediction of different types of healthcare facilities?
• RQ 3: Are all components in HealthMamba effective?
• RQ 4: How does HealthMamba perform under abnormal scenarios such as public emergencies (COVID-19)?

4.1 Evaluation Setup

Datasets. We utilize healthcare facility visit data from four major U.S. states (California, New York, Texas, and Florida), provided by Dewey. The dataset captures daily visit counts to

• (a) UrbanGPT vs. HealthMamba (b) Mamba vs. HealthMamba

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

four categories (i.e., Ambulatory Health Care Services, Hospitals, Nursing and Residential Care Facilities, and Social Assistance) of healthcare facilities at the county level from January 2019 to April 2022, spanning 40 months.

Baselines. We compare HealthMamba with 5 categories of 13 state-of-the-art baselines: (1) GNN-based: DCRNN [Li et al., 2017], STGCN [Yu et al., 2018], AGCRN [Bai et al., 2020], DGCRN [Li et al., 2023], UQGNN [Yu et al., 2025]; (2) Attention-based: DSTAGNN [Lan et al., 2022], ASTGCN [Zhu et al., 2021]; (3) Transformer-based: GluonTS [Alexandrov et al., 2020], PatchTST [Nie, 2022]; (4) LLM-based: ST-LLM [Liu et al., 2024], UrbanGPT [Li et al., 2024b]; (5) Mamba-based: Mamba [Gu and Dao, 2024], U-Mamba [Ma et al., 2024].

Metrics. We utilize two widely used metrics, including Mean Absolute Error (MAE) and Root Mean Squared Error (RMSE), to evaluate the performance of deterministic prediction, and three other commonly used metrics, including Mean Prediction Interval Width (MPIW), Interval Score (IS), and Coverage (COV), to evaluate the performance of uncertainty quantification. It is worth noting that our model outputs prediction intervals, and the deterministic metrics are computed based on the mean of the upper and lower bounds.

4.2 Overall Performance Comparison (RQ 1)

An overall comparison of our HealthMamba and other baselines is presented in Table 1. We found that our HealthMamba reduces MAE by approximately 6.0% compared to the best baseline Mamba based on the average of all four datasets. In addition, HealthMamba also demonstrates superior performance on prediction reliability, with a 3.5% improvement in IS and reaching the target coverage. As shown in Figure 2, the green shadow of HealthMamba is more compact while still covering most observations, indicating more reliable and precise estimates compared to baselines. Our experiments show that HealthMamba consistently achieves the best performance for both short-term (next-day) and medium-term (7-day-ahead) prediction.

4.3 Type-specific Visit Prediction (RQ 2)

We further show the prediction results on different healthcare facility categories in Table 2. We found that HealthMamba consistently outperforms Mamba and UrbanGPT across all categories. Notably, Mamba fails to achieve target coverage in the Nursing and Social categories, while UrbanGPT exhibits substantially wider prediction intervals despite achieving coverage. In contrast, HealthMamba maintains valid and compact confidence intervals across all categories, demonstrating robustness to facility-type heterogeneity.

4.4 Ablation Study (RQ 3)

To evaluate the contribution of each key component in HealthMamba, we conduct an ablation study by comparing the full model against six variants: (1) w/o STCE, which removes the contextual information; (2) w/o G-Mamba, which replaces the GraphMamba with conventional Mamba; (3) w/o Node-based UQ module, (4) w/o Distribution-based UQ module, and (5) w/o Parameter-based UQ module; and (6) w/o UQ, which uses only quantile regression.

Table 3 validates our design choices. Architecturally, removing GraphMamba or STCE significantly degrades MAE (up to 15.782), highlighting the crucial role of spatiotemporal and contextual modeling. Furthermore, omitting any individual UQ module (Node-, Parameter-, or Distribution-based) increases MAE and fails the coverage constraint (COV = ✗). Finally, completely removing UQ fails to provide valid prediction intervals entirely (missing IS and COV), demonstrating the necessity of our synergistic UQ framework.

4.5 Prediction under Abnormal Situations (RQ 4)

A key challenge in healthcare visit prediction is handling abnormal scenarios (e.g., emergencies, extreme weather) that cause sudden pattern deviations. We evaluate HealthMamba on two representative events: the COVID-19 lockdown in LA County (March 2020) causing a sudden drop in visits, and Hurricane Hanna in Texas (July 2020) triggering a surge. As shown in Figure 3, under these situations, HealthMamba successfully maintained accurate predictions, and ground truth values are also within the prediction intervals in most cases, whereas the best baseline struggled to adapt to these abrupt changes. These results highlight HealthMamba's superior adaptability to both demand surges and drops under abnormal scenarios.

5 RelatedWork

5.1 Healthcare Facility Visit Prediction

Healthcare facility visit prediction has attracted much interest from both academia and the government, as it is important for health resource allocation and public emergency responses. Traditional methods rely on statistical models such as ARIMA [Alghamdi et al., 2019; Orhan and Kurutkan, 2025] and regression-based methods [Pun et al., 2019; Avinash et al., 2025; Shen et al., 2025] for daily or weekly patient count prediction, but overlook temporal dynamics and spatial dependencies. Recent works have explored different approaches to address this problem, including GNN-based methods [Yu et al., 2018; Li et al., 2024a; Shen et al., 2026], attention-based methods [Lan et al., 2022], Transformerbased methods [Nie, 2022], and LLM-based methods [Li et al., 2024b; Li et al., 2026; Li et al., 2025]. However, most existing works focus on aggregated prediction without considering different types of healthcare facilities, such as hospitals and nursing facilities, which is more important for practical decision support. In addition, existing solutions usually consider this problem a time-series prediction problem, without capturing spatial dependencies, which can potentially be used to improve prediction performance.

5.2 Uncertainty-aware Spatiotemporal Prediction

Uncertainty quantification has attracted much interest from the spatiotemporal prediction community because it lays the foundation for reliable and safe decision-making [Jiang et al., 2025b; Yu et al., 2026; Cheng et al., 2025]. DeepSTUQ [Qian et al., 2023] employs dual sub-networks to separately quantify aleatoric and epistemic uncertainty at each node. DiffSTG [Wen et al., 2023] generalizes denoising diffusion probabilistic models to spatiotemporal graphs to capture intrinsic uncertainties. STZINB-GNN [Zhuang et al., 2022] models sparse spatiotemporal data through zero-inflated negative binomial distributions. CF-GNN [Huang et al., 2024] extends conformal prediction to graph-based models for guaranteed coverage, providing distribution-free prediction intervals. UQGNN [Yu et al., 2025] integrates graph neural networks with node-level uncertainty estimation for spatiotemporal prediction. Different from existing works [Yang et al., 2023; Yang et al., 2024], we design a comprehensive uncertainty quantification module that integrates distributionbased, parameter-based, and node-based uncertainty quantification mechanisms, combined with post-hoc quantile calibration to produce well-calibrated prediction intervals for uncertainty-aware spatiotemporal prediction.

6 Conclusion

In this paper, we propose HealthMamba, an uncertaintyaware spatiotemporal framework for accurate and reliable healthcare facility visit prediction. There are three key components in HealthMamba: (i) a Unified Spatiotemporal Context Encoder that integrates both heterogeneous static and dynamic contextual information, (ii) a novel graph state space model called GraphMamba for spatiotemporal modeling, and (iii) a comprehensive UQ module integrating three different types of UQ mechanisms and a post-hoc quantile calibration for higher prediction reliability. Extensive experiments on four real-world datasets from California, New York, Texas, and Florida demonstrate that HealthMamba outperforms 13 state-of-the-art baselines, improving prediction accuracy by 6.0% and uncertainty quantification by 3.5%. Our uncertainty-aware design also enables robust predictions under public emergencies and extreme weather events, supporting trustworthy decision-making under abnormal scenarios.

Acknowledgments

We sincerely thank all anonymous reviewers for their insightful comments and valuable suggestions. This work is partially supported by the National Science Foundation under Grant 2411152 and the FSU Institute for Successful Longevity (ISL) Planning Grant.

Ethics Statement

All authors comply with the IJCAI Code of Ethics and take full responsibility for this research. We use only de-identified data without human subjects or sensitive information, posing no foreseeable harm. There are no conflicts of interest or external sponsorships.