HealthMamba

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

Healthcare facility visit prediction underpins resource allocation and public health policy, but existing work usually treats it as plain time-series forecasting — ignoring the spatial dependencies between different types of facilities — and gives no reliable answer under abnormal situations such as public emergencies. HealthMamba is an uncertainty-aware spatiotemporal framework with three components:

1. STCE — a Unified Spatiotemporal Context Encoder that fuses heterogeneous static and dynamic information.
2. G-Mamba — a Graph State Space Model (GraphMamba) for hierarchical spatiotemporal modelling.
3. Comprehensive uncertainty quantification — three complementary mechanisms combined with post-hoc quantile calibration.

On four large-scale real-world datasets from California, New York, Texas and Florida, HealthMamba improves prediction accuracy by ~6.0% and uncertainty quantification by ~3.5% over state-of-the-art baselines.

Method

Internal tensor convention is (B, N, T, D) — batch, node, time, hidden — so the graph convolution mixes across N at each step while the SSM scans along T.

STCE (STCE in src/flow/healthmamba/hm_model.py)

Embeds the visit features, mixes them spatially with a prior-graph convolution (D̃^{-1/2}(A+I)D̃^{-1/2}), then temporally with a depthwise-conv + channel-MLP mixer, producing R ∈ R^{B×N×T×d_model}.

The framework's dataloader exposes only the visit tensor V. The paper's optional static demographics D and dynamic externals E are not part of this pipeline, so STCE here fuses the visit stream alone.

G-Mamba backbone

A U-Net over the temporal axis. Each GMambaBlock is fully residual and performs, in order:

Stage	Detail
Adaptive graph learning	An attention-style, data-driven, symmetric normalized adjacency computed per forward pass from a learnable node embedding, blended with the prior graph.
Graph-convolutional spatial mixing	One-hop convolution over the blended adjacency.
SSM temporal mixing	A selective state-space (Mamba) scan per node along T, on mamba_ssm's fused selective_scan_fn.
Channel mixing	A pointwise MLP with RMSNorm.

unet_depth encoder stages downsample time, a bottleneck follows, and matching decoder stages upsample with skip connections; num_layers blocks per stage.

Uncertainty-aware heads and PHQC (src/flow/healthmamba/PHQC_engine.py)

HealthMamba has two modes, switched by --cqr:

• Point mode (--cqr no, default) — a single regression head returns (B, H, N, F), consumed by the shared BaseEngine.
• UQ mode (--cqr horizon / --cqr global) — the runner passes the true feature count F as cqr_channels, HealthMamba builds its own uncertainty heads (it does not use the runner's 3F output widening), and forward returns a dict {q_lo, q_mid, q_hi, mu, logvar, mc_var}, consumed by PHQC_Engine.

Three complementary mechanisms are trained jointly:

1. Node-based — ordered lower / median / upper quantile heads, trained with the pinball loss L_quant.
2. Distribution-based — a heteroscedastic Gaussian head (mu, sigma²), trained with the Gaussian NLL L_nll plus a calibration loss L_calib forcing the standardized residuals (y − mu)/sigma to zero mean and unit variance.
3. Parameter-based — MC dropout: at inference M stochastic passes decompose the predictive variance into aleatoric and epistemic parts, and the epistemic part widens the quantile interval. A consistency penalty L_param on the variance of the stochastic mean head is minimized during training.

L_total = L_quant + w_nll · L_nll + w_param · L_param + w_calib · L_calib

Post-hoc quantile calibration. After training, on a held-out calibration split the engine computes the one-sided adjustment margin

c = Quantile_{1-alpha} { max( l_i - y_i , y_i - u_i , 0 ) }

and widens every test interval to [l − c, u + c]. The max(·, 0) clamp makes c ≥ 0, so calibration only ever widens intervals to reach the target coverage. c is per-horizon (--phqc horizon) or scalar (--phqc global), computed at the conformal level ceil((n+1)(1−α))/n for the finite-sample coverage guarantee, and persisted in the checkpoint. Calibration and metrics run in the original (inverse-transformed) data space; the per-feature inverse transform is monotonic, so quantile ordering is preserved.

Ablation toggles --use_stce, --use_gmamba, --use_node, --use_dist, --use_param reproduce the paper's ablation variants.

Repository layout

HealthMamba/
  src/flow/healthmamba/
    hm_model.py         STCE + G-Mamba U-Net + three uncertainty heads
    PHQC_engine.py      Joint UQ objective + post-hoc quantile calibration
    main.py             Entry point: model args, adjacency setup, run_experiment()
  base/                 Shared framework
    runner.py           run_experiment(): the single experiment driver
    model.py            BaseModel contract
    engine.py           Training / validation / test loop, checkpointing, export
    CQR_engine.py       Conformalized Quantile Regression engine
    metrics.py          Point, distributional and interval metrics
    efficiency.py       Hardware info, memory, inference time, FLOPs
  utils/
    args.py             Common CLI arguments, path config, set_seed
    dataloader.py       Dataset / DataLoader, dataset registry lookup
    generate.py         Raw array -> his.npz / info.json / split indices
    registry.yaml       Dataset name -> data & adjacency paths
    graph_algo.py       Adjacency normalizations
    get_adj_mat.py      Build an adjacency matrix from geographic shapefiles
    log.py              Logger
    res.py              Result collection / comparison CLI
  jobs/train.sh         Slurm submission script

Installation

conda create -n st python=3.10 -y
conda activate st

# PyTorch (CUDA 12.8 build)
pip install torch --index-url https://download.pytorch.org/whl/cu128

pip install -r requirements.txt

mamba_ssm needs a CUDA toolchain matching your PyTorch build; install it after PyTorch.

Data

Point the framework at your data root (defaults to <repo>/datasets):

export POPST_DATA=/path/to/datasets      # where datasets live
export POPST_RESULT=/path/to/result      # where logs & checkpoints go

Each dataset folder has this layout, and every entry is produced by utils/generate.py:

<POPST_DATA>/<dataset>/
  <adj_name>.npy        Adjacency matrix (N x N)
  <years>/
    his.npz             Normalized data + scaler parameters
    info.json           Shape, scaler, split sizes, seq_length_x / seq_length_y
    meta.json           Scaler parameters and raw data shape
    idx_{train,val,test,all}.npy   Split sample indices

Generate it from a raw visit array of shape (T, N, F) and register it in utils/registry.yaml:

python utils/generate.py --data_path /path/to/raw.npy --dataset my_visits --years 2018 --fmt NDT

my_visits:
  data: my_visits
  adj: my_visits/adj.npy

N is read from info.json at runtime; seq_len / horizon / input_dim / output_dim are auto-filled from the same file unless given on the command line.

Usage

# Point prediction
python src/flow/healthmamba/main.py --dataset chicago_15min --years 2018

# Uncertainty-aware: joint UQ objective + post-hoc quantile calibration
python src/flow/healthmamba/main.py --dataset chicago_15min --years 2018 \
    --cqr horizon --quantile_alpha 0.1 --phqc horizon --phqc_mc 20

# A single shared calibration margin instead
python src/flow/healthmamba/main.py --dataset chicago_15min --cqr global --phqc global

# Ablation: drop the parameter-based (MC-dropout) mechanism
python src/flow/healthmamba/main.py --dataset chicago_15min --cqr horizon --use_param 0

# Test from a checkpoint (the calibrated margin is restored with it)
python src/flow/healthmamba/main.py --dataset chicago_15min --cqr horizon \
    --mode test --model_path /path/to/HealthMamba_CQR_<timestamp>.pt

# Export prediction archives
python src/flow/healthmamba/main.py --dataset chicago_15min --mode test --export

Slurm:

sbatch jobs/train.sh
EXTRA="--cqr horizon" DATASETS="chicago_15min" sbatch jobs/train.sh

Compare runs:

python utils/res.py --path result/MyExperiment
python utils/res.py --log result/MyExperiment/HealthMamba_CQR/chicago_15min/<timestamp>.log

Results land in result/<proj>/HealthMamba/<dataset>/<timestamp>.log (HealthMamba_CQR under --cqr) next to the matching .pt checkpoint.

Arguments

Model (STCE + G-Mamba)

Argument	Default	Description
--d_model	64	Hidden width
--d_hid	128	STCE feed-forward width
--num_layers	2	G-Mamba blocks per U-Net stage
--unet_depth	2	Encoder / decoder stages S
--d_state	16	SSM state dimension
--expand	2	Mamba inner expansion factor
--d_conv	4	Depthwise causal-conv width inside the SSM
--emb_dim	32	Adaptive-graph node embedding dimension
--dropout	0.3	Dropout (also the MC-dropout rate)

Ablation toggles (1 = on, 0 = off)

Argument	Default	Description
--use_stce	1	Unified Spatiotemporal Context Encoder
--use_gmamba	1	Adaptive-graph spatial mixing inside the G-Mamba blocks
--use_node	1	Node-based (quantile) uncertainty mechanism
--use_dist	1	Distribution-based (Gaussian) uncertainty mechanism
--use_param	1	Parameter-based (MC-dropout) uncertainty mechanism

Uncertainty & calibration (used only with --cqr)

Argument	Default	Description
--cqr	no	no (point model), horizon / global — switches in PHQC_Engine
--quantile_alpha	0.1	Target miscoverage α; intervals target 1 − α coverage
--phqc	horizon	Calibration-margin granularity: horizon (per forecast step) or global
--phqc_mc	20	MC-dropout passes M at inference
--phqc_w_nll	1.0	Weight w_nll of the Gaussian NLL term
--phqc_w_param	1.0	Weight w_param of the MC-consistency term
--phqc_w_calib	1.0	Weight w_calib of the residual-calibration term

Training

Argument	Default	Description
--bs	64	Batch size
--max_epochs	2000	Maximum epochs
--patience	30	Early-stopping patience on validation loss
--lrate	1e-3	Learning rate (AdamW)
--wdecay	5e-4	Weight decay
--clip_grad_norm	1.0	Gradient-norm clipping
--step_size	200	StepLR decay interval
--gamma	0.95	StepLR decay factor
--seed	2025	Random seed

Data & system

Argument	Default	Description
--dataset	chicago_15min	Dataset name (must exist in registry.yaml)
--years	2018	Data sub-folder
--seq_len / --horizon	auto	Input length / forecast steps, auto-filled from info.json
--input_dim / --output_dim	auto	Feature counts, auto-filled from info.json
--no_normalize	--	Disable MinMax normalization (on by default)
--device	cuda	Device
--mode	train	train or test
--model_path	--	Checkpoint to load in test mode
--export	off	Save prediction archives with the final evaluation
--proj	--	Sub-folder name for grouping results

Implementation details

Experiments were run on a Linux server with an NVIDIA A100 GPU (80 GB). The reference results in the paper use PyTorch 2.3.0 / CUDA 11.8; this release is pinned to PyTorch 2.8.0 / CUDA 12.8 with mamba_ssm 2.2.6.

Baselines

Baselines follow DCRNN, STGCN, AGCRN, DGCRN, UQGNN, DSTAGNN, ASTGCN, GluonTS, PatchTST, ST-LLM, UrbanGPT, Mamba, and U-Mamba.

Runnable implementations of all of them, on the same runner and metric pipeline used here, are available in POPST.

Citation

The paper is to appear at IJCAI 2026; please cite the arXiv version until the proceedings are published.

@article{yu2026healthmamba,
  title   = {HealthMamba: An Uncertainty-aware Spatiotemporal Graph State Space Model for Effective and Reliable Healthcare Facility Visit Prediction},
  author  = {Yu, Dahai and Jiang, Lin and Xu, Rongchao and Wang, Guang},
  journal = {arXiv preprint arXiv:2602.05286},
  year    = {2026},
  note    = {To appear in the Proceedings of the 35th International Joint Conference on Artificial Intelligence (IJCAI 2026)}
}

Related work

• POPST — the unified spatiotemporal benchmarking framework this release is extracted from (~30 flow models, ~17 OD models, shared conformal-prediction engines)
• EnergyMamba (KDD 2026) — graph-enhanced selective state space model with adaptive sequential CQR
• TrustEnergy (AAAI 2026) — memory-augmented spatiotemporal GNN with sequential CQR
• UQGNN (SIGSPATIAL 2025) — multivariate Gaussian spatiotemporal prediction

Acknowledgements

The selective scan builds on Mamba; the U-Net arrangement follows U-Mamba.

License

Released under the MIT License.