fin-ai-lab/tfwm-supervised-return

TFWM encoder — Supervised (return)

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.

Supervised. Encoder trained end-to-end with an MLP head to rank the cross-section of stocks on forward VWAP return over the next 900 s (pairwise ranking loss). Each training cell is one (day, 5-minute anchor) with 16 stocks drawn from the same cross-section; the target is computed from data strictly after the anchor.

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 backbone.pt + head.pt, plus train_meta.json (the full resolved training config, the training span and the view-normalisation settings). The heads are the trained forecasting heads the paper reports; xs_ic.json is the head's cross-sectional IC on the evaluation month.

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-daystore → 1Hz_daystore/ (day-major: one trading day per record, every ticker plus precomputed targets, so each training cell is a same-day cross-section). uv run train.py dataset.backend=days in the code downloads the months a run needs
Schedule	12 passes over the 6-month span, base LR 0.0002, weight decay 0.05, effective batch 256
Pooling in config/training	last

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 (not applicable: supervised checkpoints have no config.json) 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-supervised-return", 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}/backbone.pt", map_location="cpu", weights_only=True)