fin-ai-lab/tfwm-timemae

TFWM encoder — TimeMAE

Code: fin-ai-lab/tfwm trains, loads and scores these encoders. Its examples/ retrain two of them from scratch and check the result against these weights.

SSL. TimeMAE (Cheng et al.): masked-patch reconstruction and representation alignment, 60% of patches masked.

One of 18 encoders compared in Towards Financial World Modeling (TFWM). All 18 share the same backbone and are trained for 12 passes over the same six-month spans, so they differ mainly in the training objective. See the TFWM Pre-Trained Encoders collection for the others.

Checkpoints

One checkpoint per evaluation month. Each was trained on the six months immediately before it and never saw the evaluation month.

Folder	Trained on (6 months)	Evaluated on	Note
2019-09/	2019-03 → 2019-08	2019-09	Training span is outside the released data (Market-1T covers 2019-07 → 2020-12); this encoder cannot be retrained from it.
2020-01/	2019-07 → 2019-12	2020-01
2020-08/	2020-02 → 2020-07	2020-08
2020-09/	2020-03 → 2020-08	2020-09
2020-12/	2020-06 → 2020-11	2020-12

Each folder holds config.json + model.pt, plus train_meta.json (the full resolved training config, the training span and the view-normalisation settings).

Architecture and training

Backbone	Transformer, 12 layers, width 384, 6 heads, MLP 1536, patch 8, sinusoidal positions (~22M parameters)
Input	1 Hz regular-session US equity data: 9 market channels (bid_price, vwap_all, high, low, ask_price, bid_size, ask_size, volume, n) + 11 view-information channels computed at load time (per-view normalisation statistics and window geometry) = 20 channels
Training data	fin-ai-lab/Market-1T-1Hz-2019H2-2020-dense → 1Hz_mosaic_mnth/ (one ticker-day per record on the filled 1 Hz grid, shuffled within each month). uv run train.py in the code downloads the months a run needs
Schedule	12 passes over the 6-month span, base LR 0.001, weight decay 0.05, batch 128
Pooling in config/training	mean

Readout

The paper reads every encoder two ways:

• Forecasting probes: the embedding of the last patch (pool="last"), i.e. the state at the decision time.
• Latent analyses: the mean over patches (pool="mean").

Loading a checkpoint through a mode class's from_pretrained uses the pool stored in config.json (mean for every self-supervised encoder) and ignores any pool you pass in a separate config. To get the last-patch readout, set .pool = "last" on every sub-backbone after loading (backbone, and also swa_backbone for TS2Vec and freq_backbone for TF-C).

Usage

Download one month:

from huggingface_hub import snapshot_download

path = snapshot_download("fin-ai-lab/tfwm-timemae", allow_patterns=["2020-12/*"])
ckpt = f"{path}/2020-12"

With the project code:

from market_jepa.eval.checkpoints import load_encoder

encoder = load_encoder(ckpt, pool="last")   # or pool="mean"

Without it, the files are plain PyTorch state dicts:

import torch

state = torch.load(f"{ckpt}/model.pt", map_location="cpu", weights_only=True)