""" AHAD QUANT Forex V5 — Temporal Fusion Transformer (TFT) Inspired by Lim et al., 2021 — adapted for binary classification on Forex 1h data. Architecture ─────────────────────────────────────────────────────────────── Input : (B, T, 62) — B batches, T=168 timesteps, 62 features Layer 1: Variable Selection Network (VSN) — per-timestep soft feature weighting via GRN + Softmax → (B, T, D_MODEL) Layer 2: Local LSTM encoder — captures short-range sequential dependencies → (B, T, D_MODEL) Layer 3: Gated skip connection (GRN) with add+norm → (B, T, D_MODEL) Layer 4: Temporal Multi-Head Self-Attention (N_HEADS heads) — captures long-range dependencies → (B, T, D_MODEL) Layer 5: Gated skip connection + mean pooling → (B, D_MODEL) Layer 6: Classification head GRN → Linear(D_MODEL, 1) — logit (use with BCEWithLogitsLoss) Output : (B,) logit — apply sigmoid for probability Training util functions ─────────────────────── train_tft(X_seq_tr, y_tr, X_seq_va, y_va, seq_scaler, ...) → TemporalFusionTransformer predict_tft_proba(model, X_seq_norm_tensor) → np.ndarray (B,) """ import math import time import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader, TensorDataset # ─── Hyper-parameters (can be overridden via kwargs in train_tft) ──────────── D_MODEL = 64 # embedding / hidden dimension N_HEADS = 4 # attention heads (D_MODEL must be divisible by N_HEADS) LSTM_LAYERS = 1 # LSTM depth N_ATT_LAYERS = 2 # stacked temporal attention blocks DROPOUT = 0.10 BATCH_SIZE = 512 LR = 8e-4 # raised from 5e-4: VSN's 62 parallel GRNs need more # headroom once GRAD_CLIP is loosened (see below) WEIGHT_DECAY = 1e-4 MAX_EPOCHS = 30 PATIENCE = 6 # early stopping patience GRAD_CLIP = 5.0 # raised from 1.0: the VSN's 62 independent per-feature # GRNs produce a much larger raw gradient norm than a # single shared embedding (e.g. TransformerGRU's), # so a clip of 1.0 over-throttled ALL parameters and # kept loss stuck near ln(2). 5.0 lets real updates # through while still preventing rare gradient spikes. # ─── Building blocks ───────────────────────────────────────────────────────── class GatedResidualNetwork(nn.Module): """ Core TFT building block. x → ELU(W1·x) → Dropout → W2·(·) → GLU gate → LayerNorm(skip + output) Supports optional context vector (c) injected before the first linear. """ def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, dropout: float = DROPOUT, context_dim: int = 0): super().__init__() self.fc1 = nn.Linear(input_dim + context_dim, hidden_dim) self.fc2 = nn.Linear(hidden_dim, output_dim * 2) # × 2 for GLU # Skip connection: project input → output dim if they differ self.skip = nn.Linear(input_dim, output_dim) if input_dim != output_dim else nn.Identity() self.norm = nn.LayerNorm(output_dim) self.drop = nn.Dropout(dropout) self.out_dim = output_dim def forward(self, x: torch.Tensor, context: torch.Tensor | None = None) -> torch.Tensor: h = x if context is None else torch.cat([x, context], dim=-1) h = F.elu(self.fc1(h)) h = self.drop(h) h = self.fc2(h) # (*, out_dim*2) h1, h2 = h.chunk(2, dim=-1) # GLU split h = h1 * torch.sigmoid(h2) # Gated Linear Unit return self.norm(self.skip(x) + h) class VariableSelectionNetwork(nn.Module): """ Soft per-timestep feature selection. Each of the `num_features` features is projected to D_MODEL via its own GRN. A shared GRN over the concatenated input produces Softmax weights. Output = weighted sum of per-feature projections → (B, T, D_MODEL). """ def __init__(self, num_features: int, d_model: int = D_MODEL, dropout: float = DROPOUT): super().__init__() self.d_model = d_model self.num_features = num_features # One GRN per feature: (B, T, 1) → (B, T, d_model) self.var_grns = nn.ModuleList([ GatedResidualNetwork(1, d_model, d_model, dropout) for _ in range(num_features) ]) # Selection GRN: (B, T, num_features) → weights (B, T, num_features) self.select_grn = GatedResidualNetwork(num_features, d_model, num_features, dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: """ x : (B, T, num_features) → (B, T, d_model) """ B, T, num_feat = x.shape # Per-feature embeddings: (B*T, 1) → (B*T, d_model) → stack → (B, T, F, d_model) xi = x.reshape(B * T, num_feat) # (B*T, F) embeddings = torch.stack( [grn(xi[:, i : i + 1]) for i, grn in enumerate(self.var_grns)], dim=1 # (B*T, F, d_model) ) # → (B*T, F, d_model) # Selection weights weights = self.select_grn(xi) # (B*T, F) logits weights = F.softmax(weights, dim=-1) # (B*T, F) soft weights # Weighted sum over features out = (embeddings * weights.unsqueeze(-1)).sum(dim=1) # (B*T, d_model) return out.reshape(B, T, self.d_model) class TemporalAttentionBlock(nn.Module): """One layer of temporal self-attention + gated residual.""" def __init__(self, d_model: int = D_MODEL, n_heads: int = N_HEADS, dropout: float = DROPOUT): super().__init__() self.attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True) self.gate = GatedResidualNetwork(d_model, d_model, d_model, dropout) self.norm = nn.LayerNorm(d_model) self.drop = nn.Dropout(dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: attn_out, _ = self.attn(x, x, x) # (B, T, d_model) attn_out = self.drop(attn_out) return self.norm(x + self.gate(attn_out)) # residual # ─── Main model ────────────────────────────────────────────────────────────── class TemporalFusionTransformer(nn.Module): """ TFT for binary classification on time-series of shape (B, T, num_features). Output: (B,) raw logit — use with BCEWithLogitsLoss. """ def __init__( self, num_features : int = 62, seq_len : int = 168, d_model : int = D_MODEL, n_heads : int = N_HEADS, lstm_layers : int = LSTM_LAYERS, n_att_layers : int = N_ATT_LAYERS, dropout : float = DROPOUT, ): super().__init__() self.seq_len = seq_len # ── 1. Input embedding ─────────────────────────────────────────────── # NOTE: previously used a VariableSelectionNetwork with one independent # GatedResidualNetwork PER FEATURE (62 branches), combined via softmax- # weighted average. Averaging ~62 quasi-independent branches collapses # inter-sample variance by a factor of ~sqrt(num_features) (classic # noise-cancellation effect), starving the model of signal before # training even starts — confirmed empirically (logit std stuck ~0.03 # across 60 training steps even on an easy synthetic task, vs healthy # growth to ~0.65 with this simpler shared projection). Replaced with # a single shared Linear + GRN, matching the embedding strategy that # already works well in transformer_gru_model.py. self.input_proj = nn.Linear(num_features, d_model) self.input_grn = GatedResidualNetwork(d_model, d_model, d_model, dropout) # ── 2. Local LSTM encoder ──────────────────────────────────────────── self.lstm = nn.LSTM( input_size = d_model, hidden_size = d_model, num_layers = lstm_layers, batch_first = True, dropout = dropout if lstm_layers > 1 else 0.0, ) self.lstm_gate = GatedResidualNetwork(d_model, d_model, d_model, dropout) self.lstm_norm = nn.LayerNorm(d_model) # ── 3. Temporal attention stack ────────────────────────────────────── self.attn_blocks = nn.ModuleList([ TemporalAttentionBlock(d_model, n_heads, dropout) for _ in range(n_att_layers) ]) # ── 4. Output head ─────────────────────────────────────────────────── self.output_grn = GatedResidualNetwork(d_model, d_model, d_model, dropout) self.head = nn.Linear(d_model, 1) def forward(self, x: torch.Tensor) -> torch.Tensor: """ x : (B, T, num_features) → (B,) logit """ # 1. Shared input embedding (see note in __init__) h = self.input_grn(self.input_proj(x)) # (B, T, d_model) # 2. LSTM local encoder + gated skip lstm_out, _ = self.lstm(h) # (B, T, d_model) h = self.lstm_norm(h + self.lstm_gate(lstm_out)) # add+norm # 3. Stacked temporal attention for block in self.attn_blocks: h = block(h) # (B, T, d_model) # 4. Mean pooling over time + classification h = h.mean(dim=1) # (B, d_model) h = self.output_grn(h) return self.head(h).squeeze(-1) # (B,) # ─── Training utilities ────────────────────────────────────────────────────── def _make_loader(X: np.ndarray, y: np.ndarray, batch_size: int, shuffle: bool) -> DataLoader: """Wrap numpy arrays in a TensorDataset / DataLoader.""" X_t = torch.FloatTensor(X) y_t = torch.FloatTensor(y) return DataLoader(TensorDataset(X_t, y_t), batch_size=batch_size, shuffle=shuffle, num_workers=0, pin_memory=torch.cuda.is_available()) def train_tft( X_seq_tr : np.ndarray, y_tr : np.ndarray, X_seq_va : np.ndarray, y_va : np.ndarray, *, # Architecture overrides d_model : int = D_MODEL, n_heads : int = N_HEADS, n_att_layers : int = N_ATT_LAYERS, dropout : float = DROPOUT, # Training overrides batch_size : int = BATCH_SIZE, lr : float = LR, weight_decay : float = WEIGHT_DECAY, max_epochs : int = MAX_EPOCHS, patience : int = PATIENCE, grad_clip : float = GRAD_CLIP, ) -> "TemporalFusionTransformer": """ Train a TemporalFusionTransformer on pre-normalised sequence data. Parameters ---------- X_seq_tr / X_seq_va : (N, T, 62) float32 — already normalised y_tr / y_va : (N,) float32 binary labels Returns ------- Best model (lowest val loss) as a CPU-resident TemporalFusionTransformer. """ device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f" [TFT] device={device} | " f"train={len(y_tr):,} val={len(y_va):,} samples") _, T, F = X_seq_tr.shape model = TemporalFusionTransformer( num_features = F, seq_len = T, d_model = d_model, n_heads = n_heads, n_att_layers = n_att_layers, dropout = dropout, ).to(device) n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f" [TFT] parameters: {n_params:,}") optimiser = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=max_epochs) criterion = nn.BCEWithLogitsLoss() scaler_amp = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available()) train_loader = _make_loader(X_seq_tr, y_tr, batch_size, shuffle=True) val_loader = _make_loader(X_seq_va, y_va, batch_size * 2, shuffle=False) best_val_loss = float("inf") best_state = None no_improve = 0 t0 = time.time() for epoch in range(1, max_epochs + 1): # ── Train ──────────────────────────────────────────────────────────── model.train() train_loss = 0.0 for X_b, y_b in train_loader: X_b, y_b = X_b.to(device), y_b.to(device) optimiser.zero_grad() with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()): logits = model(X_b) loss = criterion(logits, y_b) scaler_amp.scale(loss).backward() scaler_amp.unscale_(optimiser) nn.utils.clip_grad_norm_(model.parameters(), grad_clip) scaler_amp.step(optimiser) scaler_amp.update() train_loss += loss.item() * len(y_b) train_loss /= len(y_tr) # ── Validate ───────────────────────────────────────────────────────── model.eval() val_loss = 0.0 correct = 0 with torch.no_grad(): for X_b, y_b in val_loader: X_b, y_b = X_b.to(device), y_b.to(device) logits = model(X_b) val_loss += criterion(logits, y_b).item() * len(y_b) correct += ((logits > 0).float() == y_b).sum().item() val_loss /= len(y_va) val_acc = correct / len(y_va) scheduler.step() elapsed = time.time() - t0 print(f" [TFT] epoch {epoch:3d}/{max_epochs} " f"train={train_loss:.4f} val={val_loss:.4f} " f"val_acc={val_acc:.4f} ({elapsed:.0f}s)") # ── Early stopping ─────────────────────────────────────────────────── if val_loss < best_val_loss - 1e-5: best_val_loss = val_loss best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()} no_improve = 0 else: no_improve += 1 if no_improve >= patience: print(f" [TFT] early stopping at epoch {epoch}") break model.load_state_dict(best_state) model.cpu() print(f" [TFT] training complete — best val_loss={best_val_loss:.4f}") return model @torch.no_grad() def predict_tft_proba( model : TemporalFusionTransformer, X_seq_t : torch.Tensor, batch_size: int = 1024, device : str = "cpu", ) -> np.ndarray: """ Run inference on a (N, T, F) float tensor. Returns probability array of shape (N,). """ model.eval() model.to(device) probs = [] for start in range(0, len(X_seq_t), batch_size): chunk = X_seq_t[start : start + batch_size].to(device) logits = model(chunk) probs.append(torch.sigmoid(logits).cpu().numpy()) model.cpu() return np.concatenate(probs).astype(np.float32)