Upload cartel_block.py

cartel_block.py

@@ -0,0 +1,79 @@
+"""
+CARTEL Backbone: Hybrid SSM + RWKV + LTC block.
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+try:
+    from .ssm_block import SimplifiedMambaBlock
+    from .ltc_gate import LTCGate
+except ImportError:
+    from ssm_block import SimplifiedMambaBlock
+    from ltc_gate import LTCGate
+
+
+class RWKVBlock(nn.Module):
+    """RWKV-style block for spatial reasoning."""
+    def __init__(self, dim: int, n_head: int = 4):
+        super().__init__()
+        self.dim = dim
+        self.n_head = n_head
+        self.head_dim = dim // n_head
+        self.linear_qkv = nn.Linear(dim, dim * 3, bias=False)
+        self.out_proj = nn.Linear(dim, dim, bias=False)
+        self.rwkv_alpha = nn.Parameter(torch.ones(n_head) * 0.5)
+        self.beta = nn.Parameter(torch.zeros(n_head))
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            nn.GELU(),
+            nn.Linear(dim * 4, dim),
+        )
+        self.time_mix = nn.Parameter(torch.ones(dim) * 0.5)
+
+    def forward(self, x: torch.Tensor):
+        residual = x
+        x = self.norm1(x)
+        qkv = self.linear_qkv(x)
+        q, k, v = qkv.chunk(3, dim=-1)
+        B, L, D = q.shape
+        q = q.reshape(B, L, self.n_head, self.head_dim).transpose(1, 2)
+        k = k.reshape(B, L, self.n_head, self.head_dim).transpose(1, 2)
+        v = v.reshape(B, L, self.n_head, self.head_dim).transpose(1, 2)
+        alpha = torch.sigmoid(self.rwkv_alpha.view(1, self.n_head, 1, 1))
+        beta = self.beta.view(1, self.n_head, 1, 1)
+        wkv = torch.zeros(B, self.n_head, 1, self.head_dim, device=x.device, dtype=x.dtype)
+        outs = []
+        for t in range(L):
+            kt = k[:, :, t:t+1, :]
+            vt = v[:, :, t:t+1, :]
+            qt = q[:, :, t:t+1, :]
+            wkv = alpha * wkv + kt.transpose(-2, -1) @ vt
+            nom = qt @ (beta * wkv.transpose(-2, -1) + kt)
+            outs.append(nom)
+        out = torch.cat(outs, dim=2)
+        out = out.transpose(1, 2).reshape(B, L, D)
+        out = self.out_proj(out)
+        x = residual + out
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+class CARTELBlock(nn.Module):
+    """One CARTEL layer = SSM + RWKV + LTC merge"""
+    def __init__(self, dim: int, d_state: int = 16, expand: int = 2):
+        super().__init__()
+        self.ssm = SimplifiedMambaBlock(dim, d_state=d_state, expand=expand)
+        self.rwkv = RWKVBlock(dim)
+        self.ltc = LTCGate(dim)
+        self.merge = nn.Linear(dim * 2, dim)
+
+    def forward(self, x: torch.Tensor):
+        x_ssm = self.ssm(x)
+        x_rwkv = self.rwkv(x)
+        stacked = torch.cat([x_ssm, x_rwkv], dim=-1)
+        merged = self.merge(stacked)
+        gated = self.ltc(merged)
+        return gated + x