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