TSFMx

TSFMx (TSFMx Standardizes Fusion of Multimodal exogenous features) is a framework for extending TSFMs (including TimesFM and Chronos) with multimodal inputs such as text.

Installation

pip install tsfmx[all]

Running in Docker

docker/Dockerfile builds a CUDA image with the dependencies, the Time-MMD clone, and the MM-TSFlib checkout already in place, and runs the Time-MMD split at build time.

docker build -t tsfmx -f docker/Dockerfile .

The image clones tsfmx from GitHub rather than from the build context, so the build ignores local changes and a rebuild after a new commit needs --no-cache.

docker run --gpus all -it \
  -v "$PWD/data/Fidel-TS:/tsfmx/data/Fidel-TS" \
  -v "$PWD/data/cache:/tsfmx/data/cache" \
  -v "$PWD/outputs:/tsfmx/outputs" \
  -v "$HOME/.cache/huggingface:/root/.cache/huggingface" \
  -e WANDB_API_KEY \
  tsfmx bash

The four mounts are what makes a run repeatable rather than disposable:

They are subdirectories rather than a single mount over data/, which would hide the Time-MMD clone baked into the image. --gpus all needs the NVIDIA Container Toolkit on the host. -e WANDB_API_KEY forwards the host variable, which the sweeps need in order to log; drop it if you are only running evaluation.

Time-MMD is ready inside the container, so the quick start below starts at step 2. Fidel-TS is not downloaded at build time — it lands in the mounted volume instead, so the sub-dataset choice is not baked into the image and the download survives a rebuild.

Quick Start

1. Setup

Clone the Time-MMD dataset:

./scripts/clone_time_mmd.sh

Split the dataset into train / val / test:

PYTHONPATH=. uv run python scripts/split_time_mmd_datasets.py \
    --train-ratio 0.7 \
    --val-ratio 0.1

2. Pre-compute Text Embeddings

TimesFM:

PYTHONPATH=. uv run python scripts/cache_time_mmd_datasets.py \
    --model-config examples/time_mmd/configs/models/timesfm.yml \
    --text-encoder-type english
PYTHONPATH=. uv run python scripts/cache_time_mmd_datasets.py \
    --model-config examples/time_mmd/configs/models/timesfm.yml \
    --text-encoder-type english --augment

Chronos:

PYTHONPATH=. uv run python scripts/cache_time_mmd_datasets.py \
    --model-config examples/time_mmd/configs/models/chronos.yml \
    --text-encoder-type english
PYTHONPATH=. uv run python scripts/cache_time_mmd_datasets.py \
    --model-config examples/time_mmd/configs/models/chronos.yml \
    --text-encoder-type english --augment

3. Fusion Hyperparameter Tuning

Run a W&B Sweeps search for the fusion mode (adapter frozen, fusion layer trained):

The sweep configs optimize val/best_loss, the best validation MSE reached during a trial. Selecting on test/mse instead would tune the hyperparameters on the same split the reported numbers come from, which is also why --keep-best-test-mse and --keep-best-test-mae are opt-in: the checkpoints they retain are chosen on test and are not valid to report.

text_dropout_prob zeroes a random subset of each training batch's text, resampled every step and applied to training only. Trained with text always present, the fusion branch can settle into a constant offset that the forecast then depends on — which step 6 sees as a large drop degradation with no matching shuffle one, even though no text was ever read. The sweep searches [0.0, 0.1, 0.25, 0.5], so training with text always present stays reachable.

TimesFM:

PYTHONPATH=. uv run python scripts/tune_time_mmd_fusion_sweep.py \
    --model-config examples/time_mmd/configs/models/timesfm.yml \
    --sweep-config examples/time_mmd/configs/sweeps/fusion_3layers.yml

Chronos:

PYTHONPATH=. uv run python scripts/tune_time_mmd_fusion_sweep.py \
    --model-config examples/time_mmd/configs/models/chronos.yml \
    --sweep-config examples/time_mmd/configs/sweeps/fusion_3layers.yml

To run the adapter mode (adapter fine-tuned, no fusion):

TimesFM:

PYTHONPATH=. uv run python scripts/tune_time_mmd_adapter_sweep.py \
    --model-config examples/time_mmd/configs/models/timesfm.yml \
    --sweep-config examples/time_mmd/configs/sweeps/adapter.yml

Chronos:

PYTHONPATH=. uv run python scripts/tune_time_mmd_adapter_sweep.py \
    --model-config examples/time_mmd/configs/models/chronos.yml \
    --sweep-config examples/time_mmd/configs/sweeps/adapter.yml

4. Fine-tune Hyperparameter Tuning

After fusion tuning, run a W&B Sweeps search for the finetune mode (adapter + fusion trained jointly), starting from the best fusion checkpoint:

TimesFM:

PYTHONPATH=. uv run python scripts/tune_time_mmd_finetune_sweep.py \
    --model-config examples/time_mmd/configs/models/timesfm.yml \
    --sweep-config examples/time_mmd/configs/sweeps/finetune_1layer.yml \
    --fusion-checkpoint-path outputs/sweeps/fusion/best_checkpoints/best_val_loss.pt

Chronos:

PYTHONPATH=. uv run python scripts/tune_time_mmd_finetune_sweep.py \
    --model-config examples/time_mmd/configs/models/chronos.yml \
    --sweep-config examples/time_mmd/configs/sweeps/finetune_1layer.yml \
    --fusion-checkpoint-path outputs/sweeps/fusion/best_checkpoints/best_val_loss.pt

5. Visualize Forecasts

After training, generate per-sample forecast plots from a saved checkpoint:

TimesFM:

PYTHONPATH=. uv run python scripts/visualize_time_mmd_predictions.py \
    --model-config examples/time_mmd/configs/models/timesfm.yml \
    --checkpoint-path outputs/sweeps/fusion/best_checkpoints/best_val_loss.pt \
    --output-dir outputs/visualizations/timesfm

Chronos:

PYTHONPATH=. uv run python scripts/visualize_time_mmd_predictions.py \
    --model-config examples/time_mmd/configs/models/chronos.yml \
    --checkpoint-path outputs/sweeps/fusion/best_checkpoints/best_val_loss.pt \
    --output-dir outputs/visualizations/chronos

Use --max-samples N to limit the number of plots per split, and --splits train val test to select which splits to visualize.

6. Text Ablation Analysis

Beating a unimodal baseline does not prove that a model reads its text: the fusion branch can also act as a plain regularizer, or latch onto a domain identity signal that happens to be encoded in the text embeddings. This script evaluates one checkpoint repeatedly, perturbing only the text side each time, and reports the degradation relative to the unperturbed run.

PYTHONPATH=. uv run python scripts/eval_time_mmd_text_ablation.py \
    --model-config examples/time_mmd/configs/models/timesfm.yml \
    --checkpoint-path outputs/sweeps/fusion/best_checkpoints/best_val_loss.pt \
    --output outputs/text_ablation_results.json

The telling comparison is drop against shuffle. If both degrade by a similar amount, the model is reading the text. If drop degrades but shuffle does not, the fusion branch is contributing independently of what the text actually says, and mean distinguishes the two readings. cross_domain separates one more explanation from those: shuffle leaves the domain intact, so text that only identifies the domain survives it, and only cross_domain destroys that too.

The oracles invert the question. Every other row degrades the text and asks whether the forecast notices; these hand the model text that is by construction worth reading, so they should improve on none. No improvement means the model cannot use text at all, whatever the corpus says; an improvement places the fault in the corpus rather than in the fusion mechanism. oracle_trend exists because sentence encoders represent magnitude poorly: oracle failing on its own would be ambiguous between a fusion branch that cannot carry sample-specific information and an encoder that cannot read a list of floats. Both consume the labels, so neither is a forecasting result.

Perturbations are applied per sample index rather than per batch, so results are independent of batch size and iteration order, and reproducible for a given --seed. Use --ablations to run a subset (none is always included as the reference), --domains to select domains, and --augment to evaluate on the augmented cache from step 2.

cross_domain takes its text from the next entry in --domains, cycling, and is skipped with a warning when fewer than two domains load. The pairing is positional, so it destroys the domain identity without preserving the date alignment. The oracles synthesize text and so load the text encoder named by the model config, which must be the one the cache was built with; they run by default, so pass --ablations explicitly to skip that load.

Read the deltas against the per-domain sample counts, which are logged and written to the num_samples field of the output JSON. With the default context_len and horizon_len of 32, a monthly domain's test split holds only a handful of samples, far too few to read a difference of a few percent; --augment raises that by up to patch_len times. Those added samples are overlapping windows rather than independent draws, so the confidence intervals narrow less than the raw count suggests.

Note that under the current bias-free fusion projection, drop and zeroing the text embeddings are equivalent, with or without fusion_normalize. Training-time text_dropout_prob relies on that same equivalence to withhold text without changing the batch shape.

7. Text Fusion Diagnostics

The ablations above say whether sample-specific text information reaches the forecast. When it does not, this script says where it was lost, which decides whether to fix the text pipeline or the fusion mechanism.

PYTHONPATH=. uv run python scripts/diagnose_time_mmd_text_fusion.py \
    --model-config examples/time_mmd/configs/models/chronos.yml \
    --checkpoint-path outputs/sweeps/fusion/best_checkpoints/best_val_loss.pt \
    --augment

It splits the text embeddings, the fusion projection output, and the time series embeddings each into *_constant_rms, the component shared by every sample, and *_varying_rms, the component that differs between samples. Both are root mean squares, so they are orthogonal parts of one magnitude, and *_varying_fraction reports the varying part as a share of it: 0 is fully collapsed, 1 is nothing shared. constant_vs_ts_rms, varying_vs_ts_rms, and projection_vs_ts_rms put the projection on the scale of the time series embeddings fusion adds it to, divided by their total magnitude — those embeddings vary strongly between samples and share correspondingly little, so dividing by their shared component inflates every ratio instead.

A low text_varying_fraction, or a text_mean_pairwise_cosine near 1, means the signal is already gone at the encoder and no fusion mechanism could recover it: all-MiniLM-L6-v2 truncates at 256 tokens and silently drops the tail of a multi-article patch, and a patch with no text at all is encoded as the empty string, which maps every such patch to one fixed vector. A healthy text_varying_fraction with a low projection_varying_fraction means the additive projection is discarding it instead.

A high projection_vs_ts_rms is a third failure: the text reaches the backbone intact but at a magnitude rivalling the time series representation, with no way for the model to admit less of it. The fusion_normalize option answers that, dividing out the projection's own output scale and replacing it with an explicit learned one. chronos_normalized.yml enables it; pass it to the step 3 sweep as --model-config, then re-run this script to see how far projection_vs_ts_rms fell. It defaults to off, so checkpoints trained before it are unaffected.

Finally it compares the domains against each other, for the text embeddings and for the projection output, over up to --cosine-sample-size samples per domain: the domain-by-domain mean pairwise cosine matrix, and separability, the mean within-domain cosine minus the mean cross-domain cosine. This settles a question the ablations raise but cannot answer. When mean beats none the fusion branch is contributing a per-domain constant rather than reading the text, and separability says whether the representation carries the domain identity that would make such a constant learnable at all.

Fidel-TS

Fidel-TS is a second benchmark, added because Time-MMD cannot support the control the ablations in step 6 call for. Its textual CSVs carry start_date and end_date — the period a report covers — and no publication timestamp, so there is no way to check that a document was available when the forecast was made. A report stating March's realized figure was published after March ended, yet step 6 hands it to a model forecasting April.

Fidel-TS records each report under the timestamp it was issued, and keeps the successive versions of a forecast, so the version available at prediction time can be identified. Download one sub-dataset with:

./scripts/download_fidel_ts.sh Bear_room

The argument names any sub-dataset in the fidel-ts collection: Bear_room, California_ISO, Canada_photovoltaics_plants, Germany_Renewable_Energy_Grid, Jena_Atmospheric_Physics, NYC_traffic_speed. The raw_data/ archives are skipped, holding the unprocessed dumps the dataset's own cleaning scripts already consumed. Reading the time series needs the fidel extra, for parquet support.

As-of text retrieval

FidelTsDataset retrieves text by the time it became available rather than by the period it describes. This is what lets a forecast-bearing report reach the model without any change to the fusion mechanism: every report legitimately usable at prediction time was issued at or before it, so each one falls inside the context window and lands on a context patch. A weather report issued at the last context timestamp — "the weather is expected to remain overcast" — is text about the horizon, delivered through a context-side slot.

Reports are sampled far more coarsely than the series (weather every 6 hours against 5-minute readings), so a patch rarely contains one. Each patch takes the most recent report at or before its final timestamp: the statement in force over that patch, which is also what a forecaster would have had in hand.

Static text (general_info, channel_info) is deliberately left out. It is constant per series, and a constant is what the fusion branch degenerates into when the text carries nothing else, so including it would make that degeneracy indistinguishable from success.

On Bear_room room 104 at the default context and horizon of 32, this yields 1859 samples with no empty text patch, 783 distinct weather sentences and 106 distinct control sentences over 3718 patch slots, and no single control sentence covering more than 16% of them. Time-MMD's text, by contrast, reports a text_mean_pairwise_cosine of 0.73-0.90.

Pre-compute text embeddings

Splitting happens inside the loader: each series is cut into contiguous train, val and test parts before any window is formed, so no window straddles a boundary and no training window can see a value a later split is evaluated on. Each entity and split pair is cached separately, under the entity name <entity>_<split>.

PYTHONPATH=. uv run python scripts/cache_fidel_ts_datasets.py \
    --model-config examples/time_mmd/configs/models/chronos.yml \
    --dataset-config examples/fidel_ts/configs/datasets/bear_room.yml \
    --text-encoder-type english

--entities restricts the run to a subset (it defaults to every series with a time series file, which for Bear_room is 80 rooms), --splits to a subset of splits, and --train-ratio / --val-ratio change the 70/10/20 default. Pass --augment for the augmented cache, as in step 2 of the Time-MMD quick start.

Running the Time-MMD scripts on Fidel-TS

Every script that reads a cached split takes --dataset, naming the cache to read, so steps 3 to 7 of the quick start work unchanged against either benchmark. The sweeps take --entities, whose values they suffix with _train, _val and _test:

PYTHONPATH=. uv run python scripts/tune_time_mmd_fusion_sweep.py \
    --model-config examples/time_mmd/configs/models/chronos.yml \
    --sweep-config examples/time_mmd/configs/sweeps/fusion_3layers.yml \
    --dataset fidel_ts --entities 104 105 107 108 110

PYTHONPATH=. uv run python scripts/eval_time_mmd_text_ablation.py \
    --model-config examples/time_mmd/configs/models/chronos.yml \
    --checkpoint-path outputs/sweeps/fusion/best_checkpoints/best_val_loss.pt \
    --dataset fidel_ts --domains 104 105 107 108 110

The scripts keep their time_mmd names, which no longer describe everything they read; --dataset defaults to time_mmd, so existing invocations are unaffected.

Running the experiment

Caching all 80 Bear_room entities takes a while, so start with a handful. Add --augment and cache a second time: the ablation deltas are read against the per-domain sample counts, and the unaugmented test splits are small.

Run the unimodal baseline first. Without it, a fusion result has nothing to be better than, and step 6's drop row is the only other estimate of it:

PYTHONPATH=. uv run python scripts/tune_time_mmd_adapter_sweep.py \
    --model-config examples/time_mmd/configs/models/chronos.yml \
    --sweep-config examples/time_mmd/configs/sweeps/adapter.yml \
    --dataset fidel_ts --entities 104 105 107 108 110 --keep-best-val-loss

Then the fusion head, then the ablations and the diagnostics, exactly as in steps 3, 6 and 7 with --dataset fidel_ts added.

What separates a successful migration from a repeat of the Time-MMD result is which rows move:

Three readings follow from that:

• shuffle does not degrade, but oracle_trend improves — the fusion mechanism works and the corpus is the problem.
• Neither shuffle nor oracle_trend moves — the fault is in the mechanism, not the data, and no change of benchmark will fix it.
• mean beats none again — the branch has settled back into a constant. Check whether the selected trial had text_dropout_prob at 0.0.

The diagnostics then say how far the input side actually moved: text_mean_pairwise_cosine against Time-MMD's 0.73-0.90 measures whether the encoder now separates samples at all.

Benchmark Comparison with MM-TSFlib

MM-TSFlib is cloned under third_party/MM-TSFlib (not tracked by git). MM-TSFlib is run on its own pre-processed Time-MMD CSVs; tsfmx is evaluated on the raw Time-MMD data split 70/10/20. Both cover the same underlying domains and split ratio.

./scripts/setup_mm_tsflib.sh

1. Run MM-TSFlib benchmark

./scripts/run_mm_tsflib_benchmark.sh 0 Autoformer YOUR_HF_TOKEN

Requires a HuggingFace token with access to LLaMA 3.

2. Evaluate tsfmx checkpoint

PYTHONPATH=. uv run python scripts/eval_tsfmx_checkpoint.py \
    --model-config examples/time_mmd/configs/models/timesfm.yml \
    --checkpoint-path outputs/sweeps/fusion/best_checkpoints/best_val_loss.pt

3. Compare results

PYTHONPATH=. uv run python scripts/compare_benchmark_results.py

Acknowledgments

We thank the Time-MMD team for providing the multimodal time series dataset used in our examples and experiments.

License

MIT