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.pngand_analysis/encoder_ratio_heatmap.pngwere 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.