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