fin-ai-lab/tfwm-supervised-multihead

TFWM encoder — Supervised (multihead)

⚠️ Pre-release. These weights and the code that loads them are a work in progress. Contents and layout may change without notice. The training code (market_jepa, stable_finance) is not public yet.

Supervised. One shared encoder with three heads (return, volatility change, spread change, all at 900 s), trained jointly with gradient-norm balancing. Otherwise identical to the single-task supervised encoders.

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 + heads.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). Train from it with dataset.backend=days machine.daystore_dir=hf://datasets/fin-ai-lab/Market-1T-1Hz-2019H2-2020-daystore/1Hz_daystore
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-multihead", allow_patterns=["2020-12/*"])
ckpt = f"{path}/2020-12"

With the project code (release forthcoming):

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)