---
tags:
- finance
- time-series
- transformer
- mixture-of-experts
- forecasting
library_name: pytorch
---

# CassaFormer

A causal decoder Transformer with mixture-of-experts layers that forecasts the
distribution of a stock's forward returns from its sequence of daily representations.
It is the second stage of a two-stage model: each day is first encoded by the frozen
[ArgusJEPA](https://huggingface.co/dams2005/ArgusJEPA), and CassaFormer reads one stock's
chronological day vectors. Code: [github.com/daMS2005/jepa-transformer-model](https://github.com/daMS2005/jepa-transformer-model).

## Architecture

- Pre-norm decoder blocks with RMSNorm, full causal multi-head attention with QK-RMSNorm
  and RoPE (no tokenizer, no learned positions), and SwiGLU feed-forwards; the middle
  blocks are top-*k* mixture-of-experts layers with a load-balancing loss and no token
  dropping.
- **Head:** a Student-t distribution (location, scale, degrees of freedom > 2) per horizon,
  trained by negative log-likelihood.
- **Targets:** log returns from day *t*'s close to the close 1, 5 and 20 trading days later.
- **Context:** trained on windows of 20, 60, 120, 252 and 504 sessions.

| Folder | Width | Layers | Heads | SwiGLU width | Experts (top-*k*) | MoE blocks | Parameters | Pairs with |
|---|---|---|---|---|---|---|---|---|
| `m/` | 768 | 6 | 12 | 3072 | 4 (top-2) | 2, 3 | 99.7M | ArgusJEPA `m/` |
| `xs/` | 256 | 4 | 4 | 1024 | 8 (top-1) | 1, 2 | 15.3M | ArgusJEPA `xs/` |

Each folder has `transformer.safetensors` and `transformer_config.json`. Use each
CassaFormer checkpoint with the ArgusJEPA checkpoint of the same size: it was trained on
that encoder's frozen representations.

## Training data

US-listed common stocks in a point-in-time top-1000 universe, daily from 2010; trained on
core and extended companies through 2023, with holdout companies kept out of training for
model selection.

## Loading

```python
import json, torch
from safetensors.torch import load_file
from jepa_transformer.config import TransformerConfig       # from the GitHub repository
from jepa_transformer.models.transformer import Transformer

cfg = json.load(open("m/transformer_config.json"))["config"]
cfg = {**cfg, "sequence_lengths": tuple(cfg["sequence_lengths"]), "moe_layers": tuple(cfg["moe_layers"]),
       "kernel_backend": "torch"}   # "triton" on NVIDIA sm80+ GPUs
model = Transformer(TransformerConfig(**cfg))
model.load_state_dict(load_file("m/transformer.safetensors"))
forecast = model.eval()(day_vectors).pred   # day_vectors: [batch, time, 768] from ArgusJEPA m/
expected_log_return = forecast.loc          # [batch, time, 3]: 1-, 5- and 20-day horizons
```

## Limitations

Research artifact, trained on one market and period through a private feature pipeline.
Forecasts are noisy by nature; nothing here is investment advice.