Probe-based Imputation via Local Subspace Optimization for Time-Series Forecasting.
PILOT trains an imputer for a frozen, inference-only time-series foundation model
(TSFM) without backpropagating through the forecaster. At each optimizer step it
builds a local subspace from the optimizer's candidate descent direction and
recently applied directions, measures forecast-gradient coordinates in that
subspace with 2r symmetric finite-difference queries over the current batch
plus a without-replacement probe sample, forms a finite-population estimator
of the full-training gradient mean, and applies a safe-release rule: the raw
update is released only when it is robustly non-harming for a confidence
ellipsoid of the unobserved full-training mean; otherwise a minimum-change
robust correction is applied. A shared reconstruction warm-up trains the
imputer alone first; the B1/B2 branches then continue from the same captured
checkpoint (identical weights, optimizer, and RNG state), so the two arms
differ only in whether the forecast mechanism intervenes.
This repository contains the code for the controlled mechanism validation (RQ0), the main imputer–forecaster–dataset matrix (RQ2–RQ3), the component ablation, and the parameter-sensitivity scans.
Only the following are needed to reproduce the paper's experiments:
run_expman.py Entry point (expman batch runner; --staged protocol)
configs/
experiment.yaml Main matrix: 2 stages x 5 datasets x 5 imputers x 4 predictors (200 runs)
ablation.yaml Component ablation: 6 branches on one representative cell (6 runs)
ablation_base.yaml Shared protocol for ablation/sensitivity (cell + pilot + staged protocol)
sensitivity.yaml OFAT sensitivity, 19 one-factor B2 arms + B1 reference (20 runs)
sensitivity_refined.yaml Additional one-factor points (epsilon, probe count)
sensitivity_reference.yaml Alignment point for the refined scan
dataset/ imputer/ predictor/ Named defaults resolved by expman (source, item, model hparams)
src/
pipeline.py expman stages: PrepareData, Warmup (shared-checkpoint early stop), ContinueTraining, Evaluate
pipeline.py expman pipeline (data prep, config, run tags)
experiments.py B1--B5, B7 branch definitions and training drivers
models/jointModel.py The PILOT mechanism: subspace, finite-difference probing,
finite-population ellipsoid, safe-release gate, robust correction
data_loader.py Windowing, shuffled 7:1:2 split, block-missing generation
evaluator.py Test/validation metrics (forecast MSE on observed future positions)
scripts/
run_pilot_controlled.py Paired single-step controlled experiment (population seeds, probe draws)
runs/, results/, logs/, cache/ and tmp/ are runtime artifacts and are
git-ignored. scripts/check_scoping.py is a development utility, not part of
the reproduction path.
Requirements: Linux, Python 3.12 (managed by uv), one CUDA GPU (the paper's
runs used a single 24 GB consumer card), and Hugging Face access for the
datasets (Real-TSF/TIME) and the frozen forecasters (TimesFM 3, Chronos-2;
Toto2 ships with the toto-models package; Sundial is a local module under
src/models/predictors/Sundial).
uv sync # root environment (includes the TimesFM 3 PyTorch runtime)
For users in mainland China, point Hugging Face at a mirror before the first run (the runner also sets this by default):
export HF_ENDPOINT=https://hf-mirror.com
run_expman.py pins BLAS/torch thread pools and idle-pool waiting policy for
its worker processes automatically; no manual tuning is needed.
All three expman matrices below run through the staged pipeline
(src/pipeline.py, stages in execution order):
missing_seed fixes the realization across seeds).patience: 5 epochs without a new minimum validation forecast
loss; captures the best checkpoint including optimizer and all RNG
states.continuation_epochs: 1 epoch.
B2 engages the forecast mechanism (subspace rank 3, ε=1e-3 FD probes over
the current batch + 4 probe batches, safe_release gate with Hotelling
ellipsoid radius, cone_delta 0.3, preserve_norm so corrected steps keep
the optimizer's displacement norm).Branches (see src/experiments.py: variants): B1 reconstruction only;
B2 full PILOT; B3 random subspace basis; B4 point projection
(no uncertainty margin); B5 mean gate (no safe-release margin);
B7 direct application of the estimated forecast-gradient direction
(norm-preserving). The paired design means every B1/B2 difference at the
endpoint is attributable to the mechanism alone.
Every batch is resumable: re-running the same command, or
python run_expman.py --resume --output-dir runs/<name>, skips completed runs.
Progress lines print to stdout; per-run metrics accumulate in
runs/<name>/metrics.sqlite3 (table metrics, JSON metric_path such as
["test", "pred_mse"], stage [1] = continuation training, [2] = evaluation;
run metadata in runs/<name>/experiments/<id>/config.pkl).
python run_expman.py --config configs/experiment.yaml --output-dir runs/main
2 stages × 5 datasets (EWELD_Load_15T, Coastal_T_S_15T, Port_Activity_D,
Supply_Chain_Customer_D_Amazon, current_velocity_10T) × 5 imputers
(SAITS, DeepMVI, ModernTCN, TimesNet, TSLANet) × 4 frozen forecasters (Toto2,
Chronos, Sundial, TimesFM3), seed = missing_seed = 42. This is the long
pole: heavy forecaster cells dominate; on a single 24 GB GPU plan for the
batch to run for many hours — use --resume across sessions.
python run_expman.py --config configs/ablation.yaml --output-dir runs/ablation
B1/B2/B3/B4/B5/B7 on the representative cell
current_velocity_10T / ModernTCN / Chronos, seed 42, sharing the staged
warm-up — each branch is one continuation epoch from the same checkpoint.
python run_expman.py --config configs/sensitivity.yaml --output-dir runs/sensitivity
python run_expman.py --config configs/sensitivity_refined.yaml --output-dir runs/sensitivity_refined
python run_expman.py --config configs/sensitivity_reference.yaml --output-dir runs/sensitivity_reference
One-factor-at-a-time arms around the base point (subspace rank, ε, gate δ, cone δ, probe-batch count), all on the representative cell with the B1 reference in-batch for paired comparison.
python scripts/run_pilot_controlled.py --output-dir runs/controlled_pilot
Fully analytic, paired-seed single-step experiments over fixed finite window
populations (the mechanism's six-panel figure, conflict boundary,
finite-difference scaling, finite-population coverage), CPU-only and fast.
--population-seeds and --probe-draws control the coverage diagnostics.
runs/<name>/metrics.sqlite3, metric path
["test", "pred_mse"] / ["test", "imp_mse"] at stage [2] (latest
attempt). Join with experiments/<id>/config.pkl for
(dataset, imputer, predictor, stage, seed).[1] rows include per-epoch
["loss", "train/val_imp_loss"], ["loss", "train/val_pred_loss"] and
["pilot", <branch>, ...] diagnostics (projection rate, robust margin,
correction ratio, finite-population coverage statistics).["test", ...] rows above.seed controls initialization, batch order, dropout and all probing
randomness; missing_seed (equal to seed here) fixes the missingness
realization. B1/B2 pairs share both.uv.lock,
expman revision) and downloaded once into the HF cache.The paper's runs were produced on a single-card host with: NVIDIA GeForce
RTX 4090 (24 GB, driver 595.71.05, CUDA 13.2), Python 3.12.3, PyTorch
2.12.1+cu130, uv 0.9.26, Linux 5.15 x86_64 (128 cores). Dependencies are
locked in uv.lock; the expman scheduler is vendored verbatim at
revision 84a4d47 under third_party/expman/ and consumed as a local path
dependency, so no external repository access is needed. A one-epoch B1/B2 smoke on Current Velocity /
ModernTCN / Chronos at commit ffbf0be finished 1/1 on both branches.
Probe-based Imputation via Local Subspace Optimization for Time-Series Forecasting.
PILOT trains an imputer for a frozen, inference-only time-series foundation model
(TSFM) without backpropagating through the forecaster. At each optimizer step it
builds a local subspace from the optimizer's candidate descent direction and
recently applied directions, measures forecast-gradient coordinates in that
subspace with 2r symmetric finite-difference queries over the current batch
plus a without-replacement probe sample, forms a finite-population estimator
of the full-training gradient mean, and applies a safe-release rule: the raw
update is released only when it is robustly non-harming for a confidence
ellipsoid of the unobserved full-training mean; otherwise a minimum-change
robust correction is applied. A shared reconstruction warm-up trains the
imputer alone first; the B1/B2 branches then continue from the same captured
checkpoint (identical weights, optimizer, and RNG state), so the two arms
differ only in whether the forecast mechanism intervenes.
This repository contains the code for the controlled mechanism validation (RQ0), the main imputer–forecaster–dataset matrix (RQ2–RQ3), the component ablation, and the parameter-sensitivity scans.
Only the following are needed to reproduce the paper's experiments:
run_expman.py Entry point (expman batch runner; --staged protocol)
configs/
experiment.yaml Main matrix: 2 stages x 5 datasets x 5 imputers x 4 predictors (200 runs)
ablation.yaml Component ablation: 6 branches on one representative cell (6 runs)
ablation_base.yaml Shared protocol for ablation/sensitivity (cell + pilot + staged protocol)
sensitivity.yaml OFAT sensitivity, 19 one-factor B2 arms + B1 reference (20 runs)
sensitivity_refined.yaml Additional one-factor points (epsilon, probe count)
sensitivity_reference.yaml Alignment point for the refined scan
dataset/ imputer/ predictor/ Named defaults resolved by expman (source, item, model hparams)
src/
pipeline.py expman stages: PrepareData, Warmup (shared-checkpoint early stop), ContinueTraining, Evaluate
pipeline.py expman pipeline (data prep, config, run tags)
experiments.py B1--B5, B7 branch definitions and training drivers
models/jointModel.py The PILOT mechanism: subspace, finite-difference probing,
finite-population ellipsoid, safe-release gate, robust correction
data_loader.py Windowing, shuffled 7:1:2 split, block-missing generation
evaluator.py Test/validation metrics (forecast MSE on observed future positions)
scripts/
run_pilot_controlled.py Paired single-step controlled experiment (population seeds, probe draws)
runs/, results/, logs/, cache/ and tmp/ are runtime artifacts and are
git-ignored. scripts/check_scoping.py is a development utility, not part of
the reproduction path.
Requirements: Linux, Python 3.12 (managed by uv), one CUDA GPU (the paper's
runs used a single 24 GB consumer card), and Hugging Face access for the
datasets (Real-TSF/TIME) and the frozen forecasters (TimesFM 3, Chronos-2;
Toto2 ships with the toto-models package; Sundial is a local module under
src/models/predictors/Sundial).
uv sync # root environment (includes the TimesFM 3 PyTorch runtime)
For users in mainland China, point Hugging Face at a mirror before the first run (the runner also sets this by default):
export HF_ENDPOINT=https://hf-mirror.com
run_expman.py pins BLAS/torch thread pools and idle-pool waiting policy for
its worker processes automatically; no manual tuning is needed.
All three expman matrices below run through the staged pipeline
(src/pipeline.py, stages in execution order):
missing_seed fixes the realization across seeds).patience: 5 epochs without a new minimum validation forecast
loss; captures the best checkpoint including optimizer and all RNG
states.continuation_epochs: 1 epoch.
B2 engages the forecast mechanism (subspace rank 3, ε=1e-3 FD probes over
the current batch + 4 probe batches, safe_release gate with Hotelling
ellipsoid radius, cone_delta 0.3, preserve_norm so corrected steps keep
the optimizer's displacement norm).Branches (see src/experiments.py: variants): B1 reconstruction only;
B2 full PILOT; B3 random subspace basis; B4 point projection
(no uncertainty margin); B5 mean gate (no safe-release margin);
B7 direct application of the estimated forecast-gradient direction
(norm-preserving). The paired design means every B1/B2 difference at the
endpoint is attributable to the mechanism alone.
Every batch is resumable: re-running the same command, or
python run_expman.py --resume --output-dir runs/<name>, skips completed runs.
Progress lines print to stdout; per-run metrics accumulate in
runs/<name>/metrics.sqlite3 (table metrics, JSON metric_path such as
["test", "pred_mse"], stage [1] = continuation training, [2] = evaluation;
run metadata in runs/<name>/experiments/<id>/config.pkl).
python run_expman.py --config configs/experiment.yaml --output-dir runs/main
2 stages × 5 datasets (EWELD_Load_15T, Coastal_T_S_15T, Port_Activity_D,
Supply_Chain_Customer_D_Amazon, current_velocity_10T) × 5 imputers
(SAITS, DeepMVI, ModernTCN, TimesNet, TSLANet) × 4 frozen forecasters (Toto2,
Chronos, Sundial, TimesFM3), seed = missing_seed = 42. This is the long
pole: heavy forecaster cells dominate; on a single 24 GB GPU plan for the
batch to run for many hours — use --resume across sessions.
python run_expman.py --config configs/ablation.yaml --output-dir runs/ablation
B1/B2/B3/B4/B5/B7 on the representative cell
current_velocity_10T / ModernTCN / Chronos, seed 42, sharing the staged
warm-up — each branch is one continuation epoch from the same checkpoint.
python run_expman.py --config configs/sensitivity.yaml --output-dir runs/sensitivity
python run_expman.py --config configs/sensitivity_refined.yaml --output-dir runs/sensitivity_refined
python run_expman.py --config configs/sensitivity_reference.yaml --output-dir runs/sensitivity_reference
One-factor-at-a-time arms around the base point (subspace rank, ε, gate δ, cone δ, probe-batch count), all on the representative cell with the B1 reference in-batch for paired comparison.
python scripts/run_pilot_controlled.py --output-dir runs/controlled_pilot
Fully analytic, paired-seed single-step experiments over fixed finite window
populations (the mechanism's six-panel figure, conflict boundary,
finite-difference scaling, finite-population coverage), CPU-only and fast.
--population-seeds and --probe-draws control the coverage diagnostics.
runs/<name>/metrics.sqlite3, metric path
["test", "pred_mse"] / ["test", "imp_mse"] at stage [2] (latest
attempt). Join with experiments/<id>/config.pkl for
(dataset, imputer, predictor, stage, seed).[1] rows include per-epoch
["loss", "train/val_imp_loss"], ["loss", "train/val_pred_loss"] and
["pilot", <branch>, ...] diagnostics (projection rate, robust margin,
correction ratio, finite-population coverage statistics).["test", ...] rows above.seed controls initialization, batch order, dropout and all probing
randomness; missing_seed (equal to seed here) fixes the missingness
realization. B1/B2 pairs share both.uv.lock,
expman revision) and downloaded once into the HF cache.The paper's runs were produced on a single-card host with: NVIDIA GeForce
RTX 4090 (24 GB, driver 595.71.05, CUDA 13.2), Python 3.12.3, PyTorch
2.12.1+cu130, uv 0.9.26, Linux 5.15 x86_64 (128 cores). Dependencies are
locked in uv.lock; the expman scheduler is vendored verbatim at
revision 84a4d47 under third_party/expman/ and consumed as a local path
dependency, so no external repository access is needed. A one-epoch B1/B2 smoke on Current Velocity /
ModernTCN / Chronos at commit ffbf0be finished 1/1 on both branches.