| --- |
| 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. |
|
|
