| """ |
| FinJEPA: Financial Joint-Embedding Predictive Architecture |
| A JEPA-based world model for portfolio optimization over a separated action space. |
| """ |
|
|
| import math |
| from typing import Optional, Tuple, Dict, List |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import numpy as np |
|
|
|
|
| def off_diagonal(x): |
| n, m = x.shape |
| assert n == m |
| return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten() |
|
|
|
|
| def sigreg_loss(z_a, z_b, sim_coeff=0.0, std_coeff=16.0, cov_coeff=8.0): |
| sim_loss = torch.tensor(0.0, device=z_a.device) |
| std_z_a = torch.sqrt(z_a.var(dim=0) + 1e-4) |
| std_z_b = torch.sqrt(z_b.var(dim=0) + 1e-4) |
| std_loss = torch.mean(F.relu(1 - std_z_a)) + torch.mean(F.relu(1 - std_z_b)) |
| N = z_a.size(0) |
| z_a = z_a - z_a.mean(dim=0) |
| z_b = z_b - z_b.mean(dim=0) |
| cov_z_a = (z_a.T @ z_a) / (N - 1) |
| cov_z_b = (z_b.T @ z_b) / (N - 1) |
| cov_loss = off_diagonal(cov_z_a).pow_(2).sum() / z_a.size(1) + off_diagonal(cov_z_b).pow_(2).sum() / z_b.size(1) |
| loss = sim_coeff * sim_loss + std_coeff * std_loss + cov_coeff * cov_loss |
| return loss, sim_loss, std_loss, cov_loss |
|
|
|
|
| class TimeSeriesTokenizer(nn.Module): |
| def __init__(self, in_features: int, embed_dim: int = 128, patch_size: int = 4): |
| super().__init__() |
| self.patch_size = patch_size |
| self.embed_dim = embed_dim |
| self.proj = nn.Conv1d(in_features, embed_dim, kernel_size=patch_size, stride=patch_size) |
| self.max_patches = 1024 |
| self.pos_embed = nn.Parameter(torch.zeros(1, self.max_patches, embed_dim)) |
| nn.init.normal_(self.pos_embed, std=0.02) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| B, T, F = x.shape |
| x = x.transpose(1, 2) |
| x = self.proj(x) |
| x = x.transpose(1, 2) |
| N = x.size(1) |
| pos = self.pos_embed[:, :N] |
| x = x + pos |
| return x |
|
|
|
|
| class RotaryPositionEmbedding(nn.Module): |
| def __init__(self, dim: int, max_seq_len: int = 2048, base: float = 10000.0): |
| super().__init__() |
| inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) |
| self.register_buffer("inv_freq", inv_freq) |
| self.max_seq_len = max_seq_len |
| self.dim = dim |
|
|
| def forward(self, x: torch.Tensor, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]: |
| device = x.device |
| t = torch.arange(seq_len, device=device).type_as(self.inv_freq) |
| freqs = torch.einsum("i,j->ij", t, self.inv_freq) |
| emb = torch.cat((freqs, freqs), dim=-1) |
| cos = emb.cos()[None, :, :] |
| sin = emb.sin()[None, :, :] |
| return cos, sin |
|
|
|
|
| def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| x1, x2 = x[..., ::2], x[..., 1::2] |
| rotated = torch.stack([-x2, x1], dim=-1).flatten(-2) |
| return x * cos + rotated * sin |
|
|
|
|
| class AdaLN(nn.Module): |
| def __init__(self, dim: int, action_dim: int): |
| super().__init__() |
| self.norm = nn.LayerNorm(dim, elementwise_affine=False) |
| self.scale_shift = nn.Linear(action_dim, dim * 2) |
| nn.init.zeros_(self.scale_shift.weight) |
| nn.init.zeros_(self.scale_shift.bias) |
| self.scale_shift.bias.data[:dim] = 1.0 |
|
|
| def forward(self, x: torch.Tensor, a_emb: torch.Tensor) -> torch.Tensor: |
| x = self.norm(x) |
| scale, shift = self.scale_shift(a_emb).chunk(2, dim=-1) |
| return x * (1 + scale) + shift |
|
|
|
|
| class FinJEPATransformerBlock(nn.Module): |
| def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0, action_dim: int = 128, dropout: float = 0.0): |
| super().__init__() |
| self.dim = dim |
| self.num_heads = num_heads |
| self.head_dim = dim // num_heads |
| assert dim % num_heads == 0 |
| self.adaln1 = AdaLN(dim, action_dim) |
| self.qkv = nn.Linear(dim, dim * 3) |
| self.proj = nn.Linear(dim, dim) |
| self.dropout = nn.Dropout(dropout) |
| self.adaln2 = AdaLN(dim, action_dim) |
| mlp_dim = int(dim * mlp_ratio) |
| self.mlp = nn.Sequential( |
| nn.Linear(dim, mlp_dim), nn.GELU(), nn.Dropout(dropout), |
| nn.Linear(mlp_dim, dim), nn.Dropout(dropout), |
| ) |
| self.rope = RotaryPositionEmbedding(self.head_dim) |
|
|
| def forward(self, x: torch.Tensor, a_emb: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, mask=None): |
| B, N, D = x.shape |
| x_norm = self.adaln1(x, a_emb) |
| qkv = self.qkv(x_norm).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) |
| q, k, v = qkv[0], qkv[1], qkv[2] |
| q = apply_rotary(q, cos, sin) |
| k = apply_rotary(k, cos, sin) |
| attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) |
| if mask is not None: |
| attn = attn.masked_fill(mask == 0, float('-inf')) |
| attn = F.softmax(attn, dim=-1) |
| attn = self.dropout(attn) |
| out = (attn @ v).transpose(1, 2).reshape(B, N, D) |
| out = self.proj(out) |
| x = x + out |
| x_norm = self.adaln2(x, a_emb) |
| x = x + self.mlp(x_norm) |
| return x |
|
|
|
|
| class FinJEPATransformerEncoder(nn.Module): |
| def __init__(self, dim=128, depth=4, num_heads=4, mlp_ratio=4.0, dropout=0.0): |
| super().__init__() |
| self.blocks = nn.ModuleList([ |
| FinJEPATransformerBlock(dim, num_heads, mlp_ratio, action_dim=dim, dropout=dropout) |
| for _ in range(depth) |
| ]) |
| self.register_buffer("dummy_action", torch.zeros(1, 1, dim)) |
| self.norm = nn.LayerNorm(dim) |
| self.rope = RotaryPositionEmbedding(dim // num_heads) |
|
|
| def forward(self, x, mask=None): |
| B, N, D = x.shape |
| cos, sin = self.rope(x, N) |
| a = self.dummy_action.expand(B, N, D) |
| for block in self.blocks: |
| x = block(x, a, cos, sin, mask) |
| x = self.norm(x) |
| return x |
|
|
|
|
| class ActionEmbedder(nn.Module): |
| def __init__(self, n_assets=10, signal_vocab_size=3, hidden_dim=128, out_dim=128): |
| super().__init__() |
| self.n_assets = n_assets |
| self.weight_proj = nn.Sequential( |
| nn.Linear(n_assets, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, hidden_dim), |
| ) |
| self.signal_embed = nn.Embedding(signal_vocab_size, hidden_dim // n_assets) |
| self.hedge_embed = nn.Embedding(2, hidden_dim // 2) |
| fusion_in = hidden_dim + n_assets * (hidden_dim // n_assets) + hidden_dim // 2 |
| self.fusion = nn.Sequential( |
| nn.Linear(fusion_in, out_dim), nn.GELU(), nn.Linear(out_dim, out_dim), |
| ) |
|
|
| def forward(self, weights, signals, hedge=None): |
| w_emb = self.weight_proj(weights) |
| s_emb = self.signal_embed(signals).flatten(1) |
| if hedge is None: |
| h_emb = torch.zeros(weights.size(0), self.fusion[0].in_features - w_emb.size(1) - s_emb.size(1), device=weights.device) |
| else: |
| h_emb = self.hedge_embed(hedge).squeeze(1) |
| a_emb = self.fusion(torch.cat([w_emb, s_emb, h_emb], dim=-1)) |
| return a_emb |
|
|
|
|
| class FinJEPAPredictor(nn.Module): |
| def __init__(self, dim=128, depth=6, num_heads=4, mlp_ratio=4.0, action_dim=128, dropout=0.0, max_target_tokens=64): |
| super().__init__() |
| self.dim = dim |
| self.max_target_tokens = max_target_tokens |
| self.target_queries = nn.Parameter(torch.zeros(1, max_target_tokens, dim)) |
| nn.init.normal_(self.target_queries, std=0.02) |
| self.blocks = nn.ModuleList([ |
| FinJEPATransformerBlock(dim, num_heads, mlp_ratio, action_dim, dropout) |
| for _ in range(depth) |
| ]) |
| self.norm = nn.LayerNorm(dim) |
| self.action_proj = nn.Linear(action_dim, dim) |
| self.rope = RotaryPositionEmbedding(dim // num_heads) |
|
|
| def forward(self, x, action_emb, n_target_tokens=None, mask=None): |
| B, N_ctx, D = x.shape |
| n_target = n_target_tokens if n_target_tokens is not None else 1 |
| queries = self.target_queries[:, :n_target].expand(B, -1, -1) |
| x = torch.cat([x, queries], dim=1) |
| N_total = x.size(1) |
| a_seq = self.action_proj(action_emb).unsqueeze(1).expand(B, N_total, D) |
| cos, sin = self.rope(x, N_total) |
| for block in self.blocks: |
| x = block(x, a_seq, cos, sin, mask) |
| x = self.norm(x) |
| target_pred = x[:, N_ctx:] |
| return target_pred |
|
|
|
|
| class InverseDynamicsModel(nn.Module): |
| def __init__(self, z_dim, n_assets, signal_vocab_size=3, hidden_dim=256): |
| super().__init__() |
| self.n_assets = n_assets |
| self.mlp = nn.Sequential( |
| nn.Linear(z_dim * 2, hidden_dim), nn.GELU(), |
| nn.Linear(hidden_dim, hidden_dim), nn.GELU(), |
| ) |
| self.weight_head = nn.Linear(hidden_dim, n_assets) |
| self.signal_head = nn.Linear(hidden_dim, n_assets * signal_vocab_size) |
|
|
| def forward(self, z_t, z_tp1): |
| h = self.mlp(torch.cat([z_t, z_tp1], dim=-1)) |
| weights = self.weight_head(h) |
| signals_logits = self.signal_head(h).reshape(-1, self.n_assets, 3) |
| return {"weights": weights, "signals_logits": signals_logits} |
|
|
|
|
| class FinJEPA(nn.Module): |
| def __init__( |
| self, in_features=7, n_assets=10, patch_size=4, embed_dim=128, |
| encoder_depth=4, encoder_heads=4, predictor_depth=6, predictor_heads=4, |
| action_hidden_dim=128, signal_vocab_size=3, mlp_ratio=4.0, |
| dropout=0.0, ema_decay=0.996, use_idm=True, use_td_branch=False, |
| ): |
| super().__init__() |
| self.embed_dim = embed_dim |
| self.ema_decay = ema_decay |
| self.use_idm = use_idm |
| self.use_td_branch = use_td_branch |
| self.n_assets = n_assets |
|
|
| self.tokenizer = TimeSeriesTokenizer(in_features, embed_dim, patch_size) |
| self.context_encoder = FinJEPATransformerEncoder(embed_dim, encoder_depth, encoder_heads, mlp_ratio, dropout) |
| self.target_encoder = FinJEPATransformerEncoder(embed_dim, encoder_depth, encoder_heads, mlp_ratio, dropout) |
| for p in self.target_encoder.parameters(): |
| p.requires_grad = False |
|
|
| self.action_embedder = ActionEmbedder(n_assets, signal_vocab_size, action_hidden_dim, embed_dim) |
| self.predictor = FinJEPAPredictor(embed_dim, predictor_depth, predictor_heads, mlp_ratio, embed_dim, dropout) |
|
|
| if use_idm: |
| self.idm = InverseDynamicsModel(embed_dim, n_assets, signal_vocab_size, hidden_dim=256) |
| else: |
| self.idm = None |
|
|
| if use_td_branch: |
| self.task_encoder = FinJEPATransformerEncoder(embed_dim, encoder_depth, encoder_heads, mlp_ratio, dropout) |
| for p in self.task_encoder.parameters(): |
| p.requires_grad = False |
| self.task_predictor = FinJEPAPredictor(embed_dim, predictor_depth, predictor_heads, mlp_ratio, embed_dim, dropout) |
| else: |
| self.task_encoder = None |
| self.task_predictor = None |
|
|
| self._init_target_encoder() |
|
|
| def _init_target_encoder(self): |
| self.target_encoder.load_state_dict(self.context_encoder.state_dict()) |
| for p in self.target_encoder.parameters(): |
| p.requires_grad = False |
|
|
| def update_target(self): |
| with torch.no_grad(): |
| for param_s, param_t in zip(self.context_encoder.parameters(), self.target_encoder.parameters()): |
| param_t.data.mul_(self.ema_decay).add_(param_s.data, alpha=1 - self.ema_decay) |
|
|
| def encode_context(self, x, mask=None): |
| tokens = self.tokenizer(x) |
| z = self.context_encoder(tokens, mask) |
| return z |
|
|
| @torch.no_grad() |
| def encode_target(self, x, mask=None): |
| tokens = self.tokenizer(x) |
| z = self.target_encoder(tokens, mask) |
| return z |
|
|
| def forward(self, context_series, target_series, weights, signals, hedge=None, context_mask=None, target_mask=None): |
| z_context = self.encode_context(context_series, context_mask) |
| with torch.no_grad(): |
| z_target = self.encode_target(target_series, target_mask) |
| action_emb = self.action_embedder(weights, signals, hedge) |
| n_target = z_target.size(1) |
| z_pred = self.predictor(z_context, action_emb, n_target_tokens=n_target, mask=context_mask) |
| idm_out = None |
| if self.use_idm and self.idm is not None: |
| z_t = z_context[:, -1] |
| z_next = z_target[:, 0] |
| idm_out = self.idm(z_t, z_next) |
| return {"z_pred": z_pred, "z_target": z_target, "action_emb": action_emb, "idm": idm_out} |
|
|
| def predict_next_state(self, state_series, weights, signals, hedge=None, mask=None): |
| z = self.encode_context(state_series, mask) |
| action_emb = self.action_embedder(weights, signals, hedge) |
| z_next = self.predictor(z, action_emb, n_target_tokens=1, mask=mask) |
| return z_next |
|
|
|
|
| class FinJEPALoss(nn.Module): |
| def __init__(self, pred_loss="l1", alpha=2.0, beta=1.0, delta=4.0, omega=0.5, gamma=0.5, use_sigreg=True): |
| super().__init__() |
| self.pred_loss = pred_loss |
| self.alpha = alpha |
| self.beta = beta |
| self.delta = delta |
| self.omega = omega |
| self.gamma = gamma |
| self.use_sigreg = use_sigreg |
|
|
| def forward(self, outputs, actions_gt=None, rollout_outputs=None): |
| z_pred = outputs["z_pred"] |
| z_target = outputs["z_target"] |
| idm_out = outputs.get("idm") |
|
|
| if self.pred_loss == "l1": |
| l_pred = F.l1_loss(z_pred, z_target) |
| else: |
| l_pred = F.mse_loss(z_pred, z_target) |
|
|
| B, N, D = z_pred.shape |
| z_pred_flat = z_pred.reshape(B * N, D) |
| z_target_flat = z_target.reshape(B * N, D) |
| l_reg, l_sim, l_std, l_cov = sigreg_loss( |
| z_pred_flat, z_target_flat, sim_coeff=0.0, std_coeff=self.alpha, cov_coeff=self.beta |
| ) |
|
|
| if z_pred.size(1) > 1: |
| z_diff = (z_pred[:, 1:] - z_pred[:, :-1]).pow(2).mean() |
| else: |
| z_diff = torch.tensor(0.0, device=z_pred.device) |
| l_temporal = self.delta * z_diff |
|
|
| l_idm = torch.tensor(0.0, device=z_pred.device) |
| if idm_out is not None and actions_gt is not None: |
| l_w = F.mse_loss(idm_out["weights"], actions_gt["weights"]) |
| l_s = F.cross_entropy(idm_out["signals_logits"].reshape(-1, 3), actions_gt["signals"].reshape(-1)) |
| l_idm = self.omega * (l_w + l_s) |
|
|
| l_rollout = torch.tensor(0.0, device=z_pred.device) |
| if rollout_outputs is not None and len(rollout_outputs) > 0: |
| for ro in rollout_outputs: |
| if self.pred_loss == "l1": |
| l_rollout += F.l1_loss(ro["z_pred"], ro["z_target"]) |
| else: |
| l_rollout += F.mse_loss(ro["z_pred"], ro["z_target"]) |
| l_rollout = self.gamma * (l_rollout / len(rollout_outputs)) |
|
|
| total = l_pred + l_reg + l_temporal + l_idm + l_rollout |
| return { |
| "loss": total, |
| "loss_pred": l_pred.item(), |
| "loss_reg": l_reg.item(), |
| "loss_sim": l_sim.item(), |
| "loss_std": l_std.item(), |
| "loss_cov": l_cov.item(), |
| "loss_temporal": l_temporal.item(), |
| "loss_idm": l_idm.item() if isinstance(l_idm, torch.Tensor) else 0.0, |
| "loss_rollout": l_rollout.item() if isinstance(l_rollout, torch.Tensor) else 0.0, |
| } |
|
|
