TimeKAN Ensemble: Stacking with Time-Series Foundation Models

Code accompanying the paper "" (NeurIPS 2026 submission).

We study whether ensembling TimeKAN with five time-series foundation models (TTM, Timer, TimesFM, UniTS, Moirai) improves long-horizon forecasting on standard LTSF benchmarks (ETTh1/2, ETTm1/2, Weather, Electricity), and analyze when and why the resulting hybrid succeeds or fails.

Repository structure

.
├── benchmark_ltsf_full.py        # Train TimeKAN + run FM inference (canonical hparams)
├── benchmark_ensemble_ctx512.py  # Main ensembling: TimeKAN + FM stacking (OLS / NNLS)
├── build_predictions_npz.py      # Bridges benchmark caches → npz format for the above
├── run_nnls_decorrselect.py      # NNLS oracle, DecorrSelect, Diebold-Mariano tests
├── threshold_sensitivity.py      # Decorrelation percentile sweep (Pearson r, DecorrScore)
├── analyze_sample_errors.py      # Per-sample error coupling analysis
├── make_figures.py               # Generate the 4 paper figures from cached CSVs/JSON
├── envs/                         # Per-env package hint files
│   ├── timekan.txt               # TimeKAN training + orchestration
│   ├── ttm.txt                   # TTM
│   ├── timesfm-moirai.txt        # TimesFM + Moirai
│   ├── units.txt                 # UniTS
│   └── timer.txt                 # Timer
│
├── UpdatedFMs/                   # FM inference wrappers (one per FM)
│   ├── ttm_infer.py
│   ├── timer_infer.py
│   ├── timesfm_infer.py
│   ├── units_infer.py
│   ├── moirai_infer.py
│   └── UniTS/                    # git submodule → mims-harvard/UniTS
│
├── TimeKAN/                      # git submodule → huangst21/TimeKAN
│
└── paper_results/ensemble_ctx512/
    ├── ensemble_unconstrained.csv         # All MSE/MAE results (288 rows)
    ├── ols_weights_unconstrained.csv      # Per-(dataset, horizon, FM) OLS weights
    ├── pivot_mse_unconstrained.csv        # MSE pivot for side-by-side comparison
    ├── pivot_mae_unconstrained.csv
    ├── decorrselect_results.json          # NNLS / DecorrSelect / DM stats
    ├── paper_tables/                      # Final tables used in the paper
    ├── sample_error_analysis/             # Decorrelation analysis outputs
    └── figures/                           # Output of make_figures.py (PNG + PDF)

Setup

1. Clone with submodules

git clone --recurse-submodules https://github.com/fnafees/TimeKanEnsemble.git
cd timekan-ensemble-paper

If you already cloned without --recurse-submodules:

git submodule update --init --recursive

2. Conda environments

Five conda envs are required because the FMs' upstream packages have mutually incompatible dependency pins (notably transformers versions and torch versions):

Create each env, then pip install -r envs/<env>.txt as a starting point and resolve any GPU/CUDA-version conflicts locally. The hint files are pip freeze snapshots from the dev machine, not strict lockfiles.

Then export the python interpreter paths so the orchestrating scripts can launch FM subprocesses in the right env:

export TIMEKAN_PY=/path/to/conda/envs/timekan/bin/python
export FRONTIER_PY=/path/to/conda/envs/ttm/bin/python              # TTM only
export ENS_PY=/path/to/conda/envs/timesfm-moirai/bin/python        # TimesFM + Moirai
export UNITS_PY=/path/to/conda/envs/units/bin/python               # UniTS
export TIMER_PY=/path/to/conda/envs/timer/bin/python               # Timer

If unset, scripts fall back to sys.executable. The orchestrating scripts themselves (benchmark_ltsf_full.py, benchmark_ensemble_ctx512.py, make_figures.py, run_nnls_decorrselect.py, build_predictions_npz.py, threshold_sensitivity.py, analyze_sample_errors.py) should always be invoked with $TIMEKAN_PY. The env vars above tell those orchestrators which interpreter to use when they spawn FM-inference subprocesses.

The historical env-var names (FRONTIER_PY, ENS_PY, etc.) are kept for code-stability; the FM-to-env mapping is hardcoded as FM_ENVS in benchmark_ensemble_ctx512.py if you need to verify it.

3. UniTS pretrained checkpoint

UniTS requires the public units_x128_pretrain_checkpoint.pth. Download it from the UniTS HuggingFace page and place it at:

UpdatedFMs/units_x128_pretrain_checkpoint.pth

4. Datasets

Standard LTSF datasets used in the paper:

• ETTh1, ETTh2, ETTm1, ETTm2 — https://github.com/zhouhaoyi/ETDataset (place CSVs under TimeKAN/dataset/ETT-small/)
• Weather, Electricity — Time-Series-Library Google Drive: https://drive.google.com/drive/folders/13Cg1KYOlzM5C7K8gK8NfC-F3EYxkM3D2 (the weather/weather.csv and electricity/electricity.csv files; place under TimeKAN/dataset/weather/ and TimeKAN/dataset/electricity/ respectively)

Both benchmark_ltsf_full.py and benchmark_ensemble_ctx512.py read from TimeKAN/dataset/ via the TimeKAN data loader.

5. Determinism

SEED=2021 is set at the top of benchmark_ltsf_full.py and benchmark_ensemble_ctx512.py. On the same hardware, two runs of benchmark_ltsf_full.py --datasets ETTh1 --pred-lens 96 produce identical losses to all printed digits (verified on Quadro RTX 8000, CUDA 12.x, torch 2.x).

torch.use_deterministic_algorithms(True) is not set — some FM ops (e.g., scatter operations in TimesFM) are not deterministic-safe. Different GPU architectures may produce drift in the 4th decimal of MSE.

The cached CSVs and JSON in paper_results/ensemble_ctx512/ were produced from one specific full sweep on Quadro RTX 8000. A fresh re-train will produce slightly different absolute numbers (typically 1–3% variation) due to non-deterministic CUDA ops; e.g., the cached TimeKAN ETTh1 H=96 MSE is 0.3746, while a fresh re-train on the same hardware produces 0.3669. Reviewers should expect qualitative match (model rankings, OLS-vs-NNLS gap, sign of decorrelation correlations), not bitwise equality.

Heads up: Re-running benchmark_ensemble_ctx512.py appends/updates rows in the cached paper_results/ensemble_ctx512/ensemble_unconstrained.csv. If you want to preserve the canonical CSV as a reference, copy it aside first (e.g., cp ensemble_unconstrained.csv ensemble_unconstrained.canonical.csv) before re-running.

6. Expected runtime (Quadro RTX 8000)

• TimeKAN training: ~30s per (dataset, horizon). Full sweep (24 combos) ≈ 15 min.
• TTM / Timer / TimesFM inference: ~5–10 min per combo each.
• UniTS inference: ~30 min per combo.
• Moirai inference: ~20 min per combo.
• Full FM sweep across 5 FMs × 24 combos: several hours to ~1 day depending on GPU.
• Ensembling once caches exist: <1 min.
• run_nnls_decorrselect.py + threshold_sensitivity.py + analyze_sample_errors.py: <5 min combined.
• make_figures.py: <30 s.

Reproducing results

All commands assume the env vars from §2 are exported. As a quick check:

echo "$TIMEKAN_PY" "$FRONTIER_PY" "$ENS_PY" "$UNITS_PY" "$TIMER_PY"
# Each should print an absolute path to a `python` binary; none should be empty.

Step 1: Train TimeKAN (per dataset / horizon)

$TIMEKAN_PY benchmark_ltsf_full.py \
  --datasets ETTh1,ETTh2,ETTm1,ETTm2,Weather,Electricity \
  --pred-lens 96,192,336,720

This trains TimeKAN with the canonical hyperparameters used in the paper and saves checkpoints to TimeKAN/checkpoints/{setting}/checkpoint.pth. Note that the bundled TimeKAN scripts in TimeKAN/scripts/<dataset>/ use slightly different hyperparameters (des='Exp') and are not used in this paper.

Step 2: Run FM inference + ensembling

$TIMEKAN_PY benchmark_ensemble_ctx512.py \
  --fms TTM,Timer,TimesFM,UniTS,Moirai \
  --datasets ETTh1,ETTh2,ETTm1,ETTm2,Weather,Electricity \
  --pred-lens 96,192,336,720

This appends/updates rows in paper_results/ensemble_ctx512/ensemble_unconstrained.csv (see the heads-up note in §5 about preserving the canonical CSV).

Step 3: NNLS / DecorrSelect / DM tests

$TIMEKAN_PY build_predictions_npz.py        # Build per-config .npz from caches
$TIMEKAN_PY run_nnls_decorrselect.py \
  --pred_dir ./predictions \
  --threshold_sweep \
  --output_json decorrselect_results.json

Step 4: Decorrelation threshold sensitivity

$TIMEKAN_PY threshold_sensitivity.py

Step 5: Generate figures

$TIMEKAN_PY make_figures.py

Outputs the 4 figures (PNG + PDF) to paper_results/ensemble_ctx512/figures/.

Quick smoke test (one combo, ~10 min)

A fast end-to-end check that exercises TimeKAN training → FM inference subprocess → OLS stacking → figures. Requires $TIMEKAN_PY and $FRONTIER_PY (TTM env) exported per §2; other FM env vars are not needed since this only runs TTM.

# (a) Train TimeKAN on ETTh1 H=96
$TIMEKAN_PY benchmark_ltsf_full.py --models TimeKAN --datasets ETTh1 --pred-lens 96
# Expected: final mse: ≈ 0.367 (deterministic on the same hardware with SEED=2021)

# (b) Run TTM-only ensemble for the same combo
$TIMEKAN_PY benchmark_ensemble_ctx512.py --fms TTM --datasets ETTh1 --pred-lens 96

# (c) Structural verification — Hybrid-TTM should beat both TimeKAN and raw TTM
$TIMEKAN_PY - <<'PY'
import pandas as pd
df = pd.read_csv('paper_results/ensemble_ctx512/ensemble_unconstrained.csv')
sub = df[(df.Dataset=='ETTh1') & (df.Horizon==96)].drop_duplicates('Model', keep='last')
tk  = sub[sub.Model=='TimeKAN'].MSE.iloc[0]
ht  = sub[sub.Model=='Hybrid-TTM'].MSE.iloc[0]
ttm = sub[sub.Model=='TTM'].MSE.iloc[0]
print(f'TimeKAN    MSE = {tk:.4f}')
print(f'TTM        MSE = {ttm:.4f}')
print(f'Hybrid-TTM MSE = {ht:.4f}')
assert ht < tk  and ht < ttm, 'Ensemble should beat both components'
print('OK  ensemble pipeline reproduces qualitative result')
PY

# (d) Generate figures
$TIMEKAN_PY make_figures.py

A successful run will print OK ensemble pipeline reproduces qualitative result in step (c) and emit 8 files (4 figures × {png, pdf}) in step (d). Note: a fresh re-train produces slightly different absolute MSEs than the cached CSVs due to non-deterministic CUDA ops; we verify qualitative match (ensembling improves over both components), not bitwise equality.

Citation