ArgusJEPA

A Joint-Embedding Predictive Architecture (JEPA) that learns one dense representation per stock-day from point-in-time daily features. It is the first stage of a two-stage forecaster; the second stage is CassaFormer, a causal Transformer that reads sequences of these day vectors. Code: github.com/daMS2005/jepa-transformer-model.

How it works

• Input: one asset-day as a 191-dimensional vector of point-in-time features, NaN where a value is unavailable on that day.
• Context encoder: a residual MLP (Linear → GELU → LayerNorm blocks) over the day vector with a random half of the available feature units hidden. The 64 components of the text-description embedding are one masking unit, hidden or shown together.
• Target encoder: an exponential moving average of the context encoder that sees the full day; its momentum ramps linearly from 0.996 to 1.0 over training, as in I-JEPA.
• Objective: a predictor maps the masked-view representation to the target encoder's full-view representation (L1 loss in latent space).
• Use: after training the context encoder is frozen; encode() maps a day to its vector.

Checkpoints

Folder	Representation	Encoder	Parameters (full JEPA)
m/	768	3 residual blocks, 768 wide	5.0M
xs/	256	3 residual blocks, 256 wide	0.6M

Each folder has jepa.safetensors (context encoder, target encoder and predictor) and jepa_config.json (architecture and training settings, input width and feature groups). Both were trained for 30 epochs.

Inputs

191 daily channels in ten groups, in this order:

Group	Channels
price	11
risk	3
liquidity	3
fundamentals	31
market	8
sectors	12
macro	35
options	7
events_8k	15
sec_descriptions	66

Channels are normalized as in the repository's data layer (training-window z-scores or same-day cross-sectional ranks, with explicit 0/1 mask channels). The individual feature definitions are not published; the repository's configs/data/example.toml shows the specification format.

Training data

US-listed common stocks in a point-in-time top-1000 universe (delisted names included for the days they were members), daily from 2010. Companies are split into core, extended and holdout groups; only core and extended companies through 2023 were used for training.

Loading

import json, torch
from safetensors.torch import load_file
from jepa_transformer.models.jepa import JEPA   # from the GitHub repository

cfg = json.load(open("m/jepa_config.json"))["config"]
jepa = JEPA(cfg["input_dim"], cfg["output_dim"], hidden_dim=cfg["hidden_dim"], depth=cfg["depth"],
            dropout=0.0, residual=cfg["residual"], mask_units=tuple(cfg["mask_units"]))
jepa.load_state_dict(load_file("m/jepa.safetensors"))
day_vectors = jepa.eval().encode(features)   # features: [N, 191], NaN where unavailable

Limitations

Research artifact. Representations are trained on one market and period and on a private feature pipeline; inputs must follow the same feature order and normalization. Nothing here is investment advice.