File size: 3,412 Bytes
4b288af | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | ---
tags:
- finance
- time-series
- jepa
- self-supervised-learning
- representation-learning
library_name: pytorch
---
# 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](https://huggingface.co/dams2005/CassaFormer), a causal Transformer that
reads sequences of these day vectors. Code: [github.com/daMS2005/jepa-transformer-model](https://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
```python
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.
|