JEPA Transformer
ArgusJEPA learns what a single asset-day looks like from a sparse, point-in-time feature vector. It is then frozen, and its daily embeddings become the input sequence of CassaFormer, a causal Transformer (RoPE, QK-RMSNorm, Mixture-of-Experts) that forecasts returns over several horizons. Hot paths run on custom Triton kernels, kept only where they beat torch.compile.
Four sources, one vector per asset‑day
A separate data repo collects each source into PostgreSQL and publishes reviewed, point-in-time views. A value only enters a row once it was actually public on that trading day. The universe is a point-in-time set of large, liquid US common stocks. Delisted names stay in, so the history is not limited to today's survivors.
IBKR
Completed daily TRADES bars, keyed by the stable contract ID rather than the ticker.
SEC EDGAR
CompanyFacts and filings. A value becomes visible on the first session strictly after the filing's acceptance timestamp.
Fundamentals, 8-K events, business descriptionsFRED / ALFRED
Macro series aligned by vintage, i.e. the value as it was published then, not as revised later.
Joined by date as market contextAlpha Vantage options
Implied volatility, skew and put/call activity, used only once a month's option chains are complete.
Optional; masked when missing↓ point-in-time join on (contract, trading day) → concatenated into x ∈ ℝF ↓
Missing values stay NaN. An availability mask (finite ∧ public by then) goes with the vector instead of imputed values.
ArgusJEPA: predict the full view from a masked view
Each asset-day is an independent sample; this stage models no time. Half of the available features are hidden. From what is left, the context encoder and predictor must reproduce the latent that an EMA teacher computes from the full view. The loss is in embedding space, so the model learns how features relate to each other rather than reconstructing raw numbers.
Swipe sideways to see the whole diagram →
ArgusJEPA embeddings become CassaFormer's input
JEPA is trained only on the training period, then frozen. The two models never share gradients and have separate optimizers, schedules and checkpoints. The Transformer checkpoint just records which JEPA checkpoint produced its inputs.
JEPA.encode(x) runs the context encoder on the full view, with no mask and no dropoutCassaFormer: causal Transformer over daily embeddings
A decoder-only stack in which each position t forecasts that day's forward returns. It has no tokenizer and no learned positional table: a bias-free linear layer maps the JEPA vector into the residual stream, and RoPE supplies order.
Swipe sideways to see the whole diagram →
Triton only where it beats torch.compile
Every candidate kernel was A/B-tested on NVIDIA L4 and A100 GPUs against the strongest off-the-shelf baseline: torch.compile, cuBLAS, or SDPA/FlashAttention. Parity was checked before timing, and warm-up included compilation. Two kernels survived. The backend is an explicit torch / triton switch with identical state dicts.
MoE expert MLP
The forward pass never writes the up/gate/activation tensors to memory, and the host never waits on token counts.
Tile shapes are autotuned per GPU, keyed by model shape and a power-of-two token bucket that is a runtime argument, so a new batch size never recompiles.
Strict-FP32 linear layers
FP32 matmuls on tensor cores via 3×TF32: each operand is split into a TF32 high part and a remainder, and three partial products accumulate in FP32.
torch.library op, so torch.compile traces through itIt applies only to strict FP32 runs. Under autocast or with TF32 enabled the layers stay plain nn.Linear.
Measured on L4 and A100, against torch.compile
| Operation | Shipped implementation | Custom vs compile |
|---|---|---|
| MoE experts (per layer) | Triton | up to 3.6× |
| Full model, MoE kernel | Triton MoE + compiled rest | up to 2.4× |
| Linear layers, strict FP32 (A100) | Triton 3×TF32 | 1.3 – 1.5× per GEMM |
| Full model, strict FP32 (A100) | Triton 3×TF32 + MoE | 1.4 – 1.9× |
| JEPA fused first block | PyTorch | at most 1.12×, removed |
| RMSNorm, QK-norm + RoPE, SwiGLU | PyTorch | tied or lost, removed |
| Attention | SDPA / FlashAttention | baseline |
| Early sparse-gather JEPA kernel | discarded | 30 – 100× slower |
Why it is built this way
Two stages, never joint
The forecasting loss cannot reshape the representation, and JEPA is fit on the training period only. That keeps leakage boundaries simple to audit, and either model can be retrained on its own.
Per-day JEPA, temporal Transformer
JEPA learns cross-sectional structure within a single day. Sequence modelling across days belongs to the Transformer.
Mask by zeroing, not compaction
Hidden and missing features are zeroed inside the projection, so every feature keeps its own weight column and NaNs never reach the arithmetic. Missingness becomes something the model sees, not an imputation.
Latent L1 + EMA teacher
Predicting the teacher's embedding rather than raw values avoids fitting noise in heavy-tailed inputs. The EMA target and predictor head were not enough on their own: the first encoder collapsed to about 3 effective dimensions, so a VICReg variance–covariance penalty now keeps about 370 in use (section 07).
RoPE, no learned positions
Relative position comes from rotations, so cached inference stays exact: cached keys keep their original angles. No position table has to be learned from limited financial data.
QK-RMSNorm before RoPE
Normalizing Q and K per head bounds attention-logit growth, which keeps training stable. Attention is full multi-head, with no grouped-query sharing.
MoE only in the middle
The middle blocks route each day to a few of several experts, letting different market regimes use different MLPs. The first and last blocks stay dense, so input embedding and output readout are shared.
Gates not renormalized
Renormalizing the selected gates weakens the router's gradient from the forecast loss, and with a single expert cuts it off entirely. Keeping the raw softmax probabilities preserves it. Routing is dropless, with no capacity limit.
Student-t head
Returns are heavy-tailed. The distributional head learns location, a softplus scale and degrees of freedom held above 2, so the predicted variance stays finite. It trains on NLL; a point head with MSE is also available.
Leakage-safe supervision
Labels are shifted forward to each horizon before batching, and companies are split into separate groups as well as by time. A boolean valid_mask selects the labelled positions, and windows never cross assets, continuity breaks or split boundaries.
What the encoder learned, month by month
Each dot is one company on the last trading day of a month, from 2010 to 2026: 1,000 a month, sampled from roughly 2,500 US-listed stocks. The encoder sorts them into twelve clusters, drawn as islands placed by forecast volatility (calmer left, riskier right) and 12-month momentum (winners up, losers down). Inside an island the biggest companies sit at the centre. Play it and watch companies move between islands: the stress clusters swell in early 2020 and again in 2022.
The encoders capture strong risk and volatility structure, and it holds for companies kept out of training. On their own they are a map of the market, not a trading signal.
ArgusJEPA weights