TL;DR. Evolving-graph continual forecasters degrade sharply when the sensor network grows by orders of magnitude over its lifetime (e.g. +9433% on PEMS05). A2TTA keeps any frozen backbone and attaches a tiny per-node FiLM calibrator adapted at test time from delayed ground-truth labels; at each step it additionally spins up a discardable local clone specialised on context-weighted recent labels. It recovers accuracy at a fraction of the cost of retraining and is backbone-agnostic — validated on both the Online-AN and STAEformer backbones.

---

📖 Method

A2TTA decouples what the model knows (a frozen backbone) from how it adapts on the fly (a lightweight calibrator), so adaptation cost stays constant even as the graph explodes. The wrapper is backbone-agnostic — it consumes only the backbone's prediction, the input window, and node ids, so any pretrained forecaster plugs in unchanged (we report both Online-AN / TrafficStream and STAEformer backbones).

1. Frozen backbone. A per-year checkpoint (--backbone_method, default Online-AN; STAEformer also supported) is loaded and frozen — no backbone gradients at test time. It emits a raw-scale H-step forecast y_base.

2. FiLM calibrator (src/model/a2tta.py, --calibrator_arch film). A small per-node MLP consuming y_base, the input window x_in, four temporal statistics (last / mean / std / OLS-slope) and a learnable per-node embedding, emitting a per-horizon affine ŷ = γ ⊙ y_base + β with γ = 1 + 0.5·tanh(·) (init 1) and β = std·β_norm (init 0). The head is zero-init, so the calibrator initially reproduces the frozen backbone. The node table grows automatically as new sensors appear. (Legacy archs residual / affine / adapter / … remain selectable.)

3. Delayed-label online adaptation (src/trainer/a2tta_trainer.py). Test windows are processed in true chronological order; a window's ground truth is only revealed after its horizon H has physically elapsed — enforced by a pending → candidate_pool queue. Under the default strict protocol, the logical evaluation batch contains one window and the persistent calibrator updates every 64 windows using only released labels (supervised L1 + optional consistency + proximal-to-init).

4. Discardable local clone (src/trainer/ctx_local.py, --method tta_ctx_local — our method). Before predicting each window, the stable global calibrator is cloned; the clone takes a few steps on pool samples re-weighted by relevance to the current window (time-of-day / day-of-week phase · input+base-prediction cosine similarity · recency, softmax-normalised with an ESS guard), predicts that window, and is then discarded — so the global calibrator is never biased by any single context. On free delayed-label TTA this is the only mechanism we found that consistently beats adapting on all labels (all other sample-selection schemes are ablated below).

5. Identical metrics. Predictions are scored with the same cal_metric used by every baseline, so A2TTA numbers are directly comparable.

Ablation-D (--method × --calibrator_arch)

Each row of Ablation-D is one a2tta_main.py config (numbers in tables/ablation.md):

The selection-study modes (tta_random / tta_recent / tta_error, and active a2tta_lite) also remain available on the --method flag.

---

📦 Repository layout

a2tta/
├── main.py                 # entry for all backbones / continual baselines
├── a2tta_main.py           # entry for A2TTA (+ its ablation variants)
├── stkec_main.py           # entry for STKEC
├── src/
│   ├── model/              # model.py (all backbones+baselines), a2tta.py (FiLM calibrator), ...
│   ├── trainer/            # a2tta_trainer.py (online loop), ctx_local.py (local clone), ...
│   └── dataer/             # SpatioTemporalDataset.py
├── utils/                  # data_convert, initialize, metric, common_tools
├── conf/                   # per-dataset JSON configs (a2tta_olan_*, a2tta_stae_*, baselines)
├── scripts/                # runners for A2TTA + baselines + analysis helpers
├── tables/                 # result tables — main / ablation-D / new-sensor (md + tex + PDF)
├── notebook/               # figure notebooks (HP sensitivity, new-sensor, per-horizon,
│                           #   dataset maps) + render scripts + method spec
├── data/                   # dataset skeleton + processing notebooks (see data/README.md)
└── environment.yaml

Before publishing, run the privacy check. Site-specific cluster, queue, account, user, path, email, and institution markers should be supplied at runtime rather than committed:

bash scripts/check_public_release.sh
# Add private words that are specific to your environment:
PRIVATE_MARKERS_REGEX='cluster-name|queue-name' bash scripts/check_public_release.sh

---

🚀 Getting started

1. Environment

conda env create -f environment.yaml
conda activate stg

Core dependencies: python, pytorch, torch-geometric, networkx, scipy, numpy, tqdm. A single CUDA GPU is enough; CPU works for debugging (--gpuid -1).

2. Data

The processed tensors are large and are released separately — see data/README.md

☁️ Cloud-disk download link: https://pan.baidu.com/s/1llz16kYY33TrWlKENNHC5A?pwd=xxtf code: xxtf ☁️ Raw data Link: https://pan.baidu.com/s/1BPuxL96npWlfRXDv38duww?pwd=xxtf code: xxtf

Place tensors under data/<dataset>/{RawData,FastData,graph}/ and the stable-ID sidecars under data/sensor_metadata/<dataset>/<year>/sensor_metadata.csv; the configs already point to the tensor directories. Datasets: XXL expanding-sensor benchmarks pems03 … pems12, with dataset-specific year spans ending in 2025 and sensor counts that grow by up to two orders of magnitude.

---

🏃 Running A2TTA

Prerequisite. A2TTA adapts on top of a frozen per-year backbone. By default it loads the Online-AN checkpoints. Make sure you have run the Online-AN stage for the dataset first (it is step 5 of every scripts/pemsXX_run.sh; for PEMS05 you can also run python main.py --conf conf/PEMS05/oneline_st_an_pems05.json --gpuid 0 --seed 51).

Single run — full method (Ours), Online-AN backbone (PEMS05, one seed):

python a2tta_main.py \
    --conf conf/PEMS05/a2tta_olan_pems05.json \
    --method tta_ctx_local --calibrator_arch film \
    --dataset PEMS05 --backbone_method TrafficStream --freeze_backbone 1 \
    --backbone_ckpt_logname oneline_st_an_pems05 --backbone_seed 51 \
    --eval_batch_size 1 --adapt_every_windows 64 --label_delay_steps 12 \
    --input_normalization train --sensor_id_mode metadata \
    --sensor_metadata_root data/sensor_metadata \
    --gpuid 0 --seed 51

Same, on the STAEformer backbone (backbone_seed = seed − 9 for the STAE checkpoints):

python a2tta_main.py \
    --conf conf/PEMS05/a2tta_stae_pems05.json \
    --method tta_ctx_local --calibrator_arch film \
    --dataset PEMS05 --backbone_method STAEFORMER --freeze_backbone 1 \
    --backbone_ckpt_logname retrain_staeformer_pems05 --backbone_seed 42 \
    --eval_batch_size 1 --adapt_every_windows 64 --label_delay_steps 12 \
    --input_normalization train --sensor_id_mode metadata \
    --sensor_metadata_root data/sensor_metadata \
    --gpuid 0 --seed 51

Ablation-D — sweep the five rows by changing only --method / --calibrator_arch:

for method in backbone calibrator tta_all tta_ctx_local; do
  python a2tta_main.py --conf conf/PEMS05/a2tta_olan_pems05.json --dataset PEMS05 \
      --backbone_method TrafficStream --backbone_ckpt_logname oneline_st_an_pems05 \
      --backbone_seed 51 --method "$method" --calibrator_arch film \
      --eval_batch_size 1 --adapt_every_windows 64 --label_delay_steps 12 \
      --input_normalization train --sensor_id_mode metadata \
      --sensor_metadata_root data/sensor_metadata --gpuid 0 --seed 51
done
python a2tta_main.py --conf conf/PEMS05/a2tta_olan_pems05.json --dataset PEMS05 \
    --backbone_method TrafficStream --backbone_ckpt_logname oneline_st_an_pems05 \
    --backbone_seed 51 --method tta_ctx_local --calibrator_arch affine \
    --eval_batch_size 1 --adapt_every_windows 64 --label_delay_steps 12 \
    --input_normalization train --sensor_id_mode metadata \
    --sensor_metadata_root data/sensor_metadata --gpuid 0 --seed 51

Core component variants on PEMS05 (four variants × five seeds):

bash scripts/a2tta_lite_pems05_run.sh
# single seed / GPU:        GPU=0 SEEDS="51" bash scripts/a2tta_lite_pems05_run.sh
# only the main method:     METHODS="tta_ctx_local" bash scripts/a2tta_lite_pems05_run.sh
# quick sanity (1yr,4 bat): FAST_DEV_RUN=1 bash scripts/a2tta_lite_pems05_run.sh

All datasets (a2tta_lite_all_datasets_6gpu.sh dispatches jobs across the GPUs you give it):

GPUS="0 1" DATASETS="PEMS03 PEMS04" bash scripts/a2tta_lite_all_datasets_6gpu.sh

Outputs:

• per-year metric logs → log/<DATASET>/a2tta_*-<seed>/*.log
• aggregated results CSV → run_logs/a2tta_lite_*_results.csv (year × method × seed × horizon)
• summarize a CSV into a table with python scripts/a2tta_summarize.py <csv>

Key knobs (env-overridable in the script, or CLI flags on a2tta_main.py): ADAPT_LR, ADAPT_STEPS, ADAPT_EVERY_WINDOWS, LABEL_DELAY_STEPS, BUDGET_FRAC, POOL_SIZE, LOCAL_STEPS, WARMUP_EPOCHS, LAMBDA_CONS, LAMBDA_REG, HIDDEN_DIM, NODE_EMB_DIM, and the active-score weights --w_err / --w_unc / --w_shift / --w_recency. Defaults (from the sensitivity study): adapt_lr 1e-3 · adapt_steps 3 · candidate_pool_size 512 · budget_frac 0.25 · warmup 3 · local_steps 3; only adapt_lr is materially sensitive (flat optimum 1e-3–3e-3).

---

📈 Results & figures

Result tables (Markdown + LaTeX + a zoomable vector PDF rendering) live in tables/:

Figures are reproducible notebooks in notebook/ (each embeds its rendered PNG so it displays without re-running; regeneration needs the released run_logs/ summaries at repo root — see notebook/README.md):

The method write-up used to design the overview figure is notebook/ours_method_spec.md.

---

📊 Running the baselines

All baselines reported in the main table are launched from scripts/. Each runner takes the same env overrides: GPU=<id>, SEEDS="...", DATASETS="...", METHODS="...", and NOHUP=1 to background with a timestamped log under run_logs/.

scripts/pemsXX_run.sh runs the full per-dataset pipeline end-to-end (Retrain → auto-link → Pretrain → Online-NN → Online-AN → TrafficStream → STKEC → EAC) and is the recommended starting point, because it also produces the Online-AN checkpoints A2TTA needs.

Examples

# Full pipeline on PEMS05 (produces naïve + continual baselines + Online-AN ckpts)
bash scripts/pems05_run.sh

# STRAP-paper backbones + PECPM + STRAP, just two datasets, on GPU 0
DATASETS="PEMS04 PEMS05" GPU=0 bash scripts/baselines_pems_run.sh

# Extra static backbones, only GWN + STID, all datasets
METHODS="gwn stid" bash scripts/extra_baselines_run.sh

# ST-TTC on a subset
DATASETS="pems05 pems06" GPU=0 bash scripts/sttc_run.sh

# Run a single method directly via main.py
python main.py --conf conf/PEMS05/eac.json --gpuid 0 --seed 51

Per-year metrics for every baseline are written to log/<DATASET>/<logname>-<seed>/<logname>.log.

---

🤖 Foundation-model baselines

The zero-shot (ZS) and per-year fine-tuned (FT) foundation-model columns of the main table (TimesFM-2.5, Chronos-2, Moirai-MoE, Moirai-2.0) are run from their own entry points, each in a separate conda env (their deps conflict with the STGNN stack). Setup and usage are documented in TIMESFM_README.md and MOIRAI_README.md.

# TimesFM-2.5 zero-shot on PEMS05
conda activate timesfm && python timesfm_main.py --conf conf/PEMS05/timesfm_pems05.json --gpuid 0
# Moirai-MoE zero-shot on PEMS05
conda activate moirai && python moirai_main.py --conf conf/PEMS05/moirai_moe_pems05.json --gpuid 0
# Chronos-2 zero-shot (univariate) on PEMS05
python chronos_baseline.py --dataset PEMS05 --gpuid 0

All FM predictions use the same test windows and cal_metric as the other methods. The FT variants use a 20% per-year few-shot training slice, whereas the conventional trained spatiotemporal baselines use the 60% training split. Foundation-model columns are therefore presented as references and are excluded from best/second highlighting.

Wrappers: src/model/timesfm_wrapper.py, src/model/moirai_wrapper.py; TimesFM uses a RevIN patch (timesfm_revin_patch.py) to align input normalisation.

---

🙏 Acknowledgements

This benchmark builds on the data and code of several prior works, which we gratefully acknowledge:

• TrafficStream (IJCAI'23) — paper · repo
• EAC (ICLR'25) — paper · code
• STKEC (TITS'23) — paper · repo
• ST-TTC (NeurIPS'25) — paper · repo
• STRAP (NeurIPS'25), paper · repo
• PECPM (KDD'23),
• and the conventional backbones GWN / STID / iTransformer / DLinear / ST-Norm / STAEformer.

Foundation-model baselines:

• TimesFM — paper · repo
• Chronos — paper · repo
• Moirai / Moirai-MoE (uni2ts) — paper · repo

License

Released under the Apache-2.0 License — see LICENSE.