Multidimensional Demand Forecasting using State Space Models

World Conference on Transport Research (WCTR), 2026

Trained checkpoints and archived results of the corrected 816-run production benchmark. The benchmark asks whether forecasting each state × commodity × flow trade lane — and letting a model look across lanes — beats forecasting one aggregate series, and how recurrent, Transformer, tree and state-space (S4 / S4ND / Mamba) models compare under one shared pipeline.

Code · Results workbook · All results · Run provenance

Experiments

Experiment Directory Runs Design
Exp 1 — one aggregate series outputs/agg/ 453 18 trainable arms × 25 seeds + 3 closed-form baselines (seed 26)
Exp 2 — 918 per-lane series outputs/full/ 171 57 arms × seeds 26, 105, 115
Exp 3 — window × lag sweep outputs/sweep/ 192 LSTM, Transformer, S4, S4-bidir × windows {6,12,24,36} × lags {6,12,24,36} × 3 seeds

Exp 2 arms cover per-lane models (one-hot / embedding IDs), all-lanes models that read every lane at each step (cross-attention, FiLM, and axial 2d/3d/4d variants of LSTM, GRU, Transformer, GPT, Mamba, S4 and S4-bidir), Mamba-ND, upstream S4 / S4ND (s4_s4* directories), XGBoost, LightGBM, random forest, and moving-average, random-walk and seasonal-naive baselines.

Data and protocol

Property Specification
Source U.S. Census monthly vessel trade, 8 Mississippi-corridor states × HS2 commodities × import/export
Targets Next-month trade value (US$) and shipping weight (kg); pinned measure legacy_mixed
Lanes 918 of 1,542 state × commodity × flow combinations (≥ 95% non-zero value in training months)
Split train ≤ 2023-12, validation 2024, test 2025-01 … 2025-05 (rolling one-step)
Inputs 36-month window; value, weight, 12 value lags, 12 weight lags, month sine/cosine
Scaling per-series MinMax fitted on training months only
Selection early stopping on validation normalized MSE; test results are descriptive

The 2025 test months were inspected during development, and release vintages are not verified, so results are a retrospective backtest rather than a real-time forecast (see Evaluation_Scope in the workbook).

Repository layout

Path Contents
outputs/agg/{arm}_seed{S}/ Exp 1 runs
outputs/full/{arm}_seed{S}/ Exp 2 native runs
outputs/full/s4_s4*_seed{S}/hydra/checkpoints/ Exp 2 upstream S4 / S4ND (val/loss.ckpt is the validation-selected Lightning checkpoint)
outputs/sweep/{model}_embeddings_w{W}_l{L}_seed{S}/ Exp 3 runs
outputs/{agg,full,sweep}_run_extras.tar.gz every other per-run file: training and inference logs, logs/, plots/, artifacts/ (cohort and scaling records), Hydra configs
outputs/pipeline_runs.tar.gz controller manifests and logs for every launch of the production pipeline
results/ workbook, graph archive, post-inference report, SHAP analyses
provenance/ runbook and launch scripts; qualification.tar.gz, incidents.tar.gz and launch-logs.tar.gz for the production run

Each run directory holds the files people download directly: best.pth (validation-selected; every reported number comes from it), last.pth (resume state), checkpoint_manifest.json, metrics.json and predictions.csv. Tree arms store best_xgb_value.joblib / best_xgb_weight.joblib and best_xgb_metadata.json under the same prefix for XGBoost, LightGBM and random forest. The Hub caps a repository at 20,000 files, so the remaining per-run files are packed per experiment in outputs/*_run_extras.tar.gz with their original paths; extracting an archive at the repository root restores the full run directories. The per-run copies of the normalized dataset (artifacts/*_normalized.csv, 42 GB in total) are omitted except for one set per experiment; the pipeline regenerates them from source.

from huggingface_hub import hf_hub_download
import torch

path = hf_hub_download("Celsia/WCTR2026", "outputs/full/s4_cross_attention_3d_seed26/best.pth")
state = torch.load(path, map_location="cpu", weights_only=False)

Model definitions and the evaluation path live in the code repository (src/models/, scripts/inference.py, scripts/s4_eval.py); checkpoints are loaded with that code rather than a generic from_pretrained call.

Results

  • results/results.xlsx — 18 sheets: per-seed and seed-mean metrics for all three experiments, the Exp 2 seed ensemble, Exp 3 summaries, run details, model tensors, metric definitions and evaluation scope.
  • results/post_inference_report/ — value vs tonnage metrics, per-lane error distributions by lane size, and the aggregate model vs per-lane models on the same 918-lane total.
  • results/shap_exp1/ — SHAP for all 453 Exp 1 runs (every run passed prediction-reproduction and completeness checks).
  • results/shap_exp2_partial/ — SHAP for 78 Exp 2 runs on 24 sampled lanes (all baselines, per-lane neural and tree runs, and the first all-lanes arms at seed 26); the run stopped when the GPU host shut down.
  • results/graphs.tar.gz — figures from the code repository's graph builder. Only _analysis/leaderboard.png and _analysis/encoder_ratio_heatmap.png were regenerated from this run; the other figures predate it.

Descriptive headline on the five 2025 test months: a 12-month moving average leads Exp 1; in Exp 2 the S4 and S4-bidir all-lanes arms lead (best: S4 cross-attention 3d, normalized MSE 0.0487, value R² 0.61, weight R² 0.45), with every learned arm close together and all well ahead of random-walk and seasonal-naive baselines; Exp 3 window × lag settings move each model's error by only 1–1.4%.

Provenance

Production pipeline: scripts/run_all_experiments.sh on one NVIDIA H20, launched 2026-09-23, completed 2026-09-28 (controller run 20260928T055951Z_pid373495, exit 0, deployed commit 810cf3c). Three gate bugs that rejected correctly trained runs were fixed during the run and are documented under provenance/incidents/; the corresponding code changes are in PR #28.

This revision replaces the repository's earlier checkpoints (August 2026, pre-correction protocol), which remain available at revision 70e3479ef5eeb5c20fdaa708a4b930cc5513408f.

License

This model repository is distributed under the MIT license. Source trade data is published by the U.S. Census Bureau.

Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support