UQGNN

Uncertainty Quantification of Graph Neural Networks for Multivariate Spatiotemporal Prediction

Spatiotemporal prediction usually returns a single number per region and time step, which says nothing about how much that number can be trusted. UQGNN predicts a full multivariate Gaussian for every region and forecast step, jointly over the M interacting urban phenomena (e.g. taxi / TNP / bike demand):

p(y_n | X) = N( mu_n , Sigma_n ),    mu_n in R^M,   Sigma_n in R^{MxM}  (positive definite)

so accuracy and uncertainty come out of the same forward pass, and the off-diagonal terms of Sigma_n capture how the phenomena co-vary at that location.

Method

Component	What it does
ISTE — Interaction-aware SpatioTemporal Embedding	A spatial branch (MDGCN, multivariate diffusion graph convolution over the forward/reverse random-walk supports A_q, A_h) and a temporal branch (ITCN, interaction-aware temporal convolution with a learnable inter-variable embedding) each produce an embedding of shape (B, N, M, e); the two are fused by a Hadamard product E = E_s ⊙ E_t.
MPP — Multivariate Probabilistic Prediction	Two heads map the fused embedding to the mean mu and to a lower-triangular vector that is assembled into a positive-definite covariance Sigma (Algorithm 1: symmetric fill → eigendecomposition → eigenvalue clamping at min_vec → reconstruction).
Objective	Multivariate-Gaussian negative log-likelihood (MGAU in base/metrics.py).

forward(X) takes (B, seq_len, N, M) and returns (mu, Sigma); the engine consumes the second element as the per-prediction covariance. Because the model emits its own distribution rather than a quantile triple, it sets cqr_compatible = False and the --cqr flag is rejected with a clear message.

Reported metrics: MGAU, MAE, MAPE, RMSE, CRPS, KL — all computed per forecast horizon in the original (inverse-transformed) data space.

Repository layout

UQGNN/
  src/flow/uqgnn/
    uqgnn_model.py      MDGCN + ITCN + covariance assembly (the UQGNN model)
    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          MAE / RMSE / MAPE / CRPS / KL / MGAU / 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 (sym / transition / Chebyshev / Laplacian)
    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

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.npy       Training sample indices
    idx_val.npy         Validation sample indices
    idx_test.npy        Test sample indices
    idx_all.npy         All sample indices

Generate it from a raw array of shape (T, N, M):

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

Supported input layouts: NDT (T×N×D), NTD (N×T×D), NT (N×T, D=1 appended). Then register the dataset in utils/registry.yaml:

chicago_15min:
  data: chicago_15min
  adj: chicago_15min/chicago.npy

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

Usage

# Train
python src/flow/uqgnn/main.py --dataset chicago_15min --years 2018

# Test from a checkpoint
python src/flow/uqgnn/main.py --dataset chicago_15min --years 2018 \
    --mode test --model_path /path/to/UQGNN_<timestamp>.pt

# Test and export prediction archives
python src/flow/uqgnn/main.py --dataset chicago_15min --years 2018 --mode test --export

# Group results under result/<proj>/
python src/flow/uqgnn/main.py --dataset chicago_15min --proj MyExperiment

Slurm:

sbatch jobs/train.sh
DATASETS="chicago_15min nyc_manhattan_15min" YEARS=2018 sbatch jobs/train.sh

Compare runs:

python utils/res.py --path result/MyExperiment
python utils/res.py --path result/MyExperiment --select RMSE
python utils/res.py --log result/MyExperiment/UQGNN/chicago_15min/<timestamp>.log

Results land in result/<proj>/UQGNN/<dataset>/<timestamp>.log next to UQGNN_<timestamp>.pt.

Arguments

Model (UQGNN)

Argument	Default	Description
--hidden_dim_s	64	Hidden width of the spatial (MDGCN) branch
--hidden_dim_t	64	Hidden width of the temporal (ITCN) branch
--emb_dim	32	Interaction-aware embedding dimension e
--kernel_size	3	Temporal convolution kernel size
--temporal_layers	2	Number of ITCN layers
--min_vec	1e-6	Eigenvalue floor used when clamping Sigma to positive definite

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 (Adam)
--wdecay	5e-4	Weight decay
--dropout	0.5	Dropout
--clip_grad_norm	5	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	Number of variables M, 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 (Intel Xeon, 64 GB RAM) with an NVIDIA A100 GPU. 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.

Baselines

Deterministic baselines follow STGCN, DCRNN, GWNET, StemGNN, DSTAGNN, AGCRN, and SUMformer.

Probabilistic baselines follow TimeGrad, STZINB, DeepSTUQ, CF-GNN, and DiffSTG.

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

Citation

@inproceedings{yu2025uqgnn,
  title     = {UQGNN: Uncertainty Quantification of Graph Neural Networks for Multivariate Spatiotemporal Prediction},
  author    = {Yu, Dahai and Zhuang, Dingyi and Jiang, Lin and Xu, Rongchao and Ye, Xinyue and Bu, Yuheng and Wang, Shenhao and Wang, Guang},
  booktitle = {Proceedings of the 33rd ACM International Conference on Advances in Geographic Information Systems},
  pages     = {52--65},
  year      = {2025},
  doi       = {10.1145/3748636.3762709}
}

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
• HealthMamba (IJCAI 2026) — graph state space model with three-mechanism uncertainty quantification

License

Released under the MIT License.