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, and CassaFormer reads one stock's chronological day vectors. Code: 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

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.