⚡ ULTRON: Sync model.py

model.py

@@ -0,0 +1,295 @@
+"""
+⚡ ULTRON Systems // Advanced Convolutional Transformer Architecture for OHLCV
+Implements:
+1. Reversible Instance Normalization (RevIN)
+2. Multi-Scale Overlapping 1D Convolutional Tokenizer (kernels 3, 7, 15, 31)
+3. Rotary Positional Embeddings (RoPE)
+4. Deep Conformer Blocks (Self-Attention + Depthwise Separable 1D Conv)
+5. Multi-Horizon Dual Forecasting Head (Return Magnitude + Directional Probability)
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import List, Tuple
+
+from config import ModelConfig
+
+class RevIN(nn.Module):
+    """
+    Reversible Instance Normalization to eliminate non-stationary distribution drift.
+    Normalizes input sequences per-instance and per-feature.
+    """
+    def __init__(self, num_features: int, eps: float = 1e-5, affine: bool = True):
+        super().__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.affine = affine
+        if self.affine:
+            self.affine_weight = nn.Parameter(torch.ones(num_features))
+            self.affine_bias = nn.Parameter(torch.zeros(num_features))
+
+    def forward(self, x: torch.Tensor, mode: str = "norm") -> torch.Tensor:
+        # x: [Batch, SeqLen, Features]
+        if mode == "norm":
+            self.mean = torch.mean(x, dim=1, keepdim=True).detach()
+            self.stdev = torch.sqrt(torch.var(x, dim=1, keepdim=True, unbiased=False) + self.eps).detach()
+            x_norm = (x - self.mean) / self.stdev
+            if self.affine:
+                x_norm = x_norm * self.affine_weight + self.affine_bias
+            return x_norm
+        elif mode == "denorm":
+            x_denorm = x
+            if self.affine:
+                x_denorm = (x_denorm - self.affine_bias) / (self.affine_weight + self.eps)
+            return x_denorm * self.stdev + self.mean
+        else:
+            raise ValueError(f"Unknown RevIN mode: {mode}")
+
+class MultiScaleOverlappingConvTokenizer(nn.Module):
+    """
+    Extracts multi-frequency candlestick representations via overlapping 1D convolutions.
+    Operates over overlapping temporal windows with kernel sizes [3, 7, 15, 31].
+    Stride S < K ensures heavy overlap between adjacent token receptive fields.
+    """
+    def __init__(self, in_features: int, d_model: int, kernels: List[int] = [3, 7, 15, 31], stride: int = 4):
+        super().__init__()
+        self.stride = stride
+        branch_dim = d_model // len(kernels)
+        self.branches = nn.ModuleList()
+        
+        for k in kernels:
+            # Symmetric causal/same padding
+            pad = k // 2
+            self.branches.append(
+                nn.Sequential(
+                    nn.Conv1d(in_features, branch_dim, kernel_size=k, stride=stride, padding=pad),
+                    nn.GELU()
+                )
+            )
+            
+        # Projection layer to fuse multi-scale features into exact d_model
+        concat_dim = branch_dim * len(kernels)
+        self.proj = nn.Linear(concat_dim, d_model)
+        self.norm = nn.LayerNorm(d_model)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [Batch, SeqLen, InFeatures] -> [Batch, InFeatures, SeqLen]
+        x_trans = x.transpose(1, 2)
+        
+        branch_outs = []
+        min_len = None
+        for branch in self.branches:
+            out = branch(x_trans)
+            if min_len is None or out.shape[2] < min_len:
+                min_len = out.shape[2]
+            branch_outs.append(out)
+            
+        # Align temporal length across all branches
+        aligned = [b[:, :, :min_len] for b in branch_outs]
+        fused = torch.cat(aligned, dim=1) # [Batch, concat_dim, Patches]
+        
+        # Transpose back: [Batch, Patches, concat_dim]
+        fused = fused.transpose(1, 2)
+        return self.norm(self.proj(fused))
+
+class RotaryEmbedding(nn.Module):
+    """Rotary Position Embeddings (RoPE) for sequence tokens."""
+    def __init__(self, dim: int, max_seq_len: int = 2048):
+        super().__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+        t = torch.arange(max_seq_len, dtype=torch.float32)
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos(), persistent=False)
+        self.register_buffer("sin_cached", emb.sin(), persistent=False)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        seq_len = q.shape[2]
+        cos = self.cos_cached[:seq_len, :].unsqueeze(0).unsqueeze(0) # [1, 1, seq_len, dim]
+        sin = self.sin_cached[:seq_len, :].unsqueeze(0).unsqueeze(0)
+        
+        def rotate_half(x):
+            x1 = x[..., :x.shape[-1] // 2]
+            x2 = x[..., x.shape[-1] // 2:]
+            return torch.cat((-x2, x1), dim=-1)
+            
+        q_rot = (q * cos) + (rotate_half(q) * sin)
+        k_rot = (k * cos) + (rotate_half(k) * sin)
+        return q_rot, k_rot
+
+class ConformerConvModule(nn.Module):
+    """
+    Depthwise Separable Convolution block to model fine-grained micro-structure
+    and local volatility gradients.
+    """
+    def __init__(self, d_model: int, kernel_size: int = 7, dropout: float = 0.1):
+        super().__init__()
+        self.norm = nn.LayerNorm(d_model)
+        # Pointwise expansion
+        self.pointwise1 = nn.Conv1d(d_model, 2 * d_model, kernel_size=1)
+        # Depthwise conv
+        self.depthwise = nn.Conv1d(
+            d_model, d_model, 
+            kernel_size=kernel_size, 
+            stride=1, 
+            padding=kernel_size // 2, 
+            groups=d_model
+        )
+        self.bn = nn.BatchNorm1d(d_model)
+        self.act = nn.SiLU()
+        # Pointwise projection
+        self.pointwise2 = nn.Conv1d(d_model, d_model, kernel_size=1)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: [Batch, Patches, d_model]
+        res = x
+        x = self.norm(x).transpose(1, 2)
+        # Pointwise expansion with Gated Linear Unit
+        x = self.pointwise1(x)
+        x = F.glu(x, dim=1) # [Batch, d_model, Patches]
+        # Depthwise
+        x = self.depthwise(x)
+        x = self.bn(x)
+        x = self.act(x)
+        x = self.pointwise2(x)
+        x = self.dropout(x)
+        return res + x.transpose(1, 2)
+
+class ConformerBlock(nn.Module):
+    """
+    Hybrid Convolution + FlashAttention Conformer Block.
+    """
+    def __init__(self, d_model: int, n_heads: int = 4, d_ff: int = 256, conv_kernel: int = 7, dropout: float = 0.1):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.head_dim = d_model // n_heads
+        
+        # 1. Multi-Head Self-Attention with RoPE
+        self.ln_attn = nn.LayerNorm(d_model)
+        self.q_proj = nn.Linear(d_model, d_model, bias=False)
+        self.k_proj = nn.Linear(d_model, d_model, bias=False)
+        self.v_proj = nn.Linear(d_model, d_model, bias=False)
+        self.out_proj = nn.Linear(d_model, d_model)
+        self.rope = RotaryEmbedding(self.head_dim)
+        self.attn_dropout = dropout
+        
+        # 2. Conformer Convolution Module
+        self.conv_module = ConformerConvModule(d_model, kernel_size=conv_kernel, dropout=dropout)
+        
+        # 3. Feed-Forward SwiGLU Network
+        self.ln_ffn = nn.LayerNorm(d_model)
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, d_ff),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(d_ff, d_model),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Step 1: Pre-LN Self Attention with FlashAttention-2 / scaled_dot_product_attention
+        res = x
+        x_norm = self.ln_attn(x)
+        B, T, C = x_norm.shape
+        
+        q = self.q_proj(x_norm).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x_norm).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x_norm).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        
+        # Apply RoPE
+        q, k = self.rope(q, k)
+        
+        # Native PyTorch FlashAttention / SDPA
+        attn_out = F.scaled_dot_product_attention(
+            q, k, v, 
+            dropout_p=self.attn_dropout if self.training else 0.0,
+            is_causal=True # Causal temporal mask ensures no future leakage
+        )
+        attn_out = attn_out.transpose(1, 2).contiguous().view(B, T, C)
+        x = res + self.out_proj(attn_out)
+        
+        # Step 2: Conformer Depthwise Convolution
+        x = self.conv_module(x)
+        
+        # Step 3: Feed-Forward
+        x = x + self.ffn(self.ln_ffn(x))
+        return x
+
+class OHLCVTransformer(nn.Module):
+    """
+    End-to-End Advanced Convolutional Transformer for Multi-Horizon OHLCV Forecasting.
+    """
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.revin = RevIN(num_features=cfg.in_features) if cfg.use_revin else None
+        
+        # Multi-scale overlapping tokenizer
+        self.tokenizer = MultiScaleOverlappingConvTokenizer(
+            in_features=cfg.in_features,
+            d_model=cfg.d_model,
+            kernels=cfg.conv_kernels,
+            stride=cfg.patch_stride
+        )
+        
+        # Conformer Layers
+        self.blocks = nn.ModuleList([
+            ConformerBlock(
+                d_model=cfg.d_model,
+                n_heads=cfg.n_heads,
+                d_ff=cfg.d_ff,
+                conv_kernel=cfg.conv_depthwise_kernel,
+                dropout=cfg.dropout
+            ) for _ in range(cfg.n_layers)
+        ])
+        
+        self.final_norm = nn.LayerNorm(cfg.d_model)
+        
+        # Multi-Horizon Heads:
+        # 1. Return Regression Head: predicts expected log return for [1m, 3m, 5m, 10m, 15m]
+        self.reg_head = nn.Sequential(
+            nn.Linear(cfg.d_model, cfg.d_ff),
+            nn.GELU(),
+            nn.Dropout(cfg.dropout),
+            nn.Linear(cfg.d_ff, cfg.num_horizons)
+        )
+        
+        # 2. Directional Classification Head: predicts binary direction probability
+        self.dir_head = nn.Sequential(
+            nn.Linear(cfg.d_model, cfg.d_ff),
+            nn.GELU(),
+            nn.Dropout(cfg.dropout),
+            nn.Linear(cfg.d_ff, cfg.num_horizons)
+        )
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        x: [Batch, SeqLen, InFeatures]
+        Returns:
+            pred_returns: [Batch, NumHorizons] (expected future log returns)
+            pred_logits:  [Batch, NumHorizons] (logits for P(return > 0))
+        """
+        # RevIN normalisation
+        if self.revin is not None:
+            x = self.revin(x, mode="norm")
+            
+        # Overlapping 1D conv tokenization
+        tokens = self.tokenizer(x) # [Batch, Patches, d_model]
+        
+        # Transformer processing
+        for block in self.blocks:
+            tokens = block(tokens)
+            
+        tokens = self.final_norm(tokens)
+        
+        # Global aggregation: use final token (most recent causal context)
+        context = tokens[:, -1, :] # [Batch, d_model]
+        
+        pred_returns = self.reg_head(context)
+        pred_logits = self.dir_head(context)
+        
+        return pred_returns, pred_logits