Skip to content

StrataFormer

StrataFormer(cfg)

Trading foundation model — the backbone for trading-specific LLMs.

from strata import StrataFormer, StrataFormerConfig

cfg   = StrataFormerConfig(max_seq_len=256, n_layers=4, n_heads=8)
model = StrataFormer(cfg)
print(model.summary())

forward(x, asset_ids, mask, return_all_hidden) → StrataNetOutput

x         = torch.randn(2, 128, 5)      # (B, T, 5)  variable T
asset_ids = torch.tensor([0, 1])         # (B,)  AAPL=0, SPY=1
out       = model(x, asset_ids=asset_ids)

out.action_logits   # (B, 3)   LONG / SHORT / HOLD
out.confidence      # (B, 1)   [0, 1]
out.hidden          # (B, 4)   [bias, momentum, trap_risk, uncertainty]
out.regime_logits   # (B, 4)   TRENDING / RANGING / TRANSITIONING / CHOPPY

predict_action(x, asset_ids) → dict

Variable-length inference. Auto-truncates if T > max_seq_len.

# Any sequence length — no need to match training seq_len
result = model.predict_action(torch.randn(200, 5))
# {
#   "action":      "LONG",
#   "confidence":  0.71,
#   "regime":      "TRENDING",
#   "context_len": 200,          ← bars actually used
#   "state": {
#     "bias":        0.82,
#     "momentum":    0.45,
#     "trap_risk":   0.18,
#     "uncertainty": 0.31,
#   }
# }

forward_mbm(x, asset_ids, mask) → Tensor

Masked Bar Modeling forward pass (used during pretraining only).

x    = torch.randn(4, 64, 5)
mask = torch.rand(4, 64) < 0.15    # 15% of bars masked
pred = model.forward_mbm(x, mask=mask)
# pred: (4, 64, 5) — predicted OHLCV for all positions
# Loss computed only on masked positions

save(path) / StrataFormer.load(path)

model.save("strata_base.pt")
model = StrataFormer.load("strata_base.pt")

summary() → str

print(model.summary())
# StrataFormer v1.0
#   asset           : GENERIC
#   max_seq_len     : 256  (variable, up to this limit)
#   n_assets        : 64   (multi-asset tokenizer)
#   embed_dim       : 64
#   n_layers        : 4    (causal attention blocks)
#   n_heads         : 4
#   ffn_expand      : 4x
#   hidden_dim      : 4    [bias, momentum, trap_risk, uncertainty]
#   bottleneck_dim  : 32   (attention → 4-dim compression)
#   parameters      : 92,473
#   sub-modules     :
#     StrataBarTokenizer       (OHLCV + asset_id → embed_dim)
#     StrataCausalAttentionBlock × 4  (causal self-attention)
#     StrataStateBottleneck    (embed_dim → 4-dim interpretable state)
#     StrataFormerHead         (4-dim → LONG/SHORT/HOLD + regime + confidence)
#     StrataMaskedBarHead      (embed_dim → 5  [pretraining only])

StrataFormerConfig

from strata import StrataFormerConfig

cfg = StrataFormerConfig(
    n_features     = 5,      # OHLCV features per bar
    max_seq_len    = 512,    # maximum context window (bars)
    n_assets       = 64,     # max number of distinct asset IDs
    embed_dim      = 64,     # bar → token embedding dimension
    asset_embed_dim= 16,     # asset ID embedding dimension
    n_heads        = 4,      # attention heads
    n_layers       = 4,      # transformer blocks
    ffn_expand     = 4,      # FFN hidden = embed_dim * ffn_expand
    hidden_dim     = 4,      # MUST be 4: interpretable state
    bottleneck_dim = 32,     # attention → bottleneck → 4
    head_dim       = 32,     # head MLP width
    mask_prob      = 0.15,   # fraction masked during MBM pretraining
    dropout        = 0.10,
    attn_dropout   = 0.05,
    label_smoothing= 0.05,
    asset          = "GENERIC",
)
# Tiny — testing / resource constrained
cfg_tiny  = StrataFormerConfig(n_layers=2, n_heads=2, embed_dim=32, max_seq_len=64)

# Small — single asset, ~1 year history
cfg_small = StrataFormerConfig(n_layers=2, n_heads=4, embed_dim=64, max_seq_len=128)

# Base — multi-asset, 1-2 years  (default)
cfg_base  = StrataFormerConfig(n_layers=4, n_heads=4, embed_dim=64, max_seq_len=256)

# Large — foundation model, large corpus
cfg_large = StrataFormerConfig(n_layers=6, n_heads=8, embed_dim=256, max_seq_len=512)