Upload vil_tracker/models/mlstm.py with huggingface_hub

vil_tracker/models/mlstm.py

@@ -0,0 +1,297 @@
+"""
+mLSTM Cell and Block for Vision-LSTM (ViL) backbone.
+
+Architecture follows the official NX-AI ViL-S implementation:
+- LinearHeadwiseExpand for Q/K/V projections (block-diagonal, ~3K params each)
+- Depthwise causal Conv1d on the mLSTM branch
+- Gates (igate, fgate) take concatenated [q, k, v] as input
+- Output gate from second half of proj_up output
+- Parallel mLSTM scan with matrix memory
+
+Reference: Beck et al., "xLSTM: Extended Long Short-Term Memory" (arXiv:2405.04517)
+           Alkin et al., "Vision-LSTM: xLSTM as Generic Vision Backbone" (arXiv:2406.04303)
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, einsum
+
+
+class LinearHeadwiseExpand(nn.Module):
+    """Block-diagonal linear projection: each head has its own small weight matrix.
+    
+    Instead of a full Linear(inner_dim, inner_dim) with inner_dim^2 params,
+    this uses num_heads independent (head_dim, head_dim) matrices.
+    For inner_dim=768, num_heads=192, head_dim=4:
+      Full linear: 768*768 = 589,824 params
+      Headwise:    192*4*4 = 3,072 params  (192x fewer!)
+    """
+    def __init__(self, in_features: int, num_heads: int, bias: bool = False):
+        super().__init__()
+        assert in_features % num_heads == 0, f"{in_features} not divisible by {num_heads}"
+        self.num_heads = num_heads
+        self.head_dim = in_features // num_heads
+        self.in_features = in_features
+        
+        # Weight: (num_heads, head_dim_out, head_dim_in)
+        self.weight = nn.Parameter(torch.empty(num_heads, self.head_dim, self.head_dim))
+        self.bias = nn.Parameter(torch.zeros(in_features)) if bias else None
+        self._reset_parameters()
+    
+    def _reset_parameters(self):
+        nn.init.normal_(self.weight, std=math.sqrt(2.0 / (5.0 * self.head_dim)))
+        if self.bias is not None:
+            nn.init.zeros_(self.bias)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x: (..., in_features)
+        x = rearrange(x, '... (nh d) -> ... nh d', nh=self.num_heads)
+        x = einsum(x, self.weight, '... nh d, nh od d -> ... nh od')
+        x = rearrange(x, '... nh od -> ... (nh od)')
+        if self.bias is not None:
+            x = x + self.bias
+        return x
+
+
+class StochasticDepth(nn.Module):
+    """Drop entire residual path with probability `drop_prob` during training."""
+    def __init__(self, drop_prob: float = 0.0):
+        super().__init__()
+        self.drop_prob = drop_prob
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.training or self.drop_prob == 0.0:
+            return x
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        mask = torch.bernoulli(torch.full(shape, keep_prob, device=x.device, dtype=x.dtype))
+        return x * mask / keep_prob
+
+
+class mLSTMCell(nn.Module):
+    """Parallel mLSTM cell with matrix memory.
+    
+    Official architecture from xLSTM/ViL:
+    - proj_up: Linear(D, 2*inner_dim) → split into mLSTM branch + output gate branch
+    - CausalConv1d on mLSTM branch (depthwise, k=4)
+    - LinearHeadwiseExpand for Q, K, V projections
+    - igate, fgate: Linear(3*inner_dim, num_heads) from concat(q,k,v)
+    - Parallel scan: C_t = f_t*C_{t-1} + i_t*(v_t ⊗ k_t), h_t = C_t*q_t
+    - Output: (h + skip*conv_act) * sigmoid(z), then proj_down
+    
+    ViL-S config: D=384, proj_factor=2.0, inner_dim=768, 
+                  qkv_proj_blocksize=4, num_heads=4
+    Per-cell params: ~920K (vs 2.66M with full Linear Q/K/V)
+    """
+    def __init__(
+        self,
+        dim: int = 384,
+        proj_factor: float = 2.0,
+        qkv_proj_blocksize: int = 4,
+        num_heads: int = 4,
+        conv_kernel: int = 4,
+        bias: bool = False,
+    ):
+        super().__init__()
+        self.dim = dim
+        # inner_dim rounded up to multiple of 64
+        self.inner_dim = math.ceil(proj_factor * dim / 64) * 64
+        self.num_heads = num_heads
+        self.head_dim = self.inner_dim // num_heads  # 768/4 = 192
+        
+        # Number of projection heads for Q/K/V (block-diagonal)
+        num_proj_heads = self.inner_dim // qkv_proj_blocksize
+        
+        # Up-projection: D -> 2*inner_dim (mLSTM branch + output gate branch)
+        self.proj_up = nn.Linear(dim, 2 * self.inner_dim, bias=bias)
+        
+        # Depthwise causal conv1d on mLSTM branch
+        self.conv1d = nn.Conv1d(
+            self.inner_dim, self.inner_dim,
+            kernel_size=conv_kernel,
+            padding=conv_kernel - 1,  # causal: pad left
+            groups=self.inner_dim,
+            bias=True,
+        )
+        self.conv_kernel = conv_kernel
+        
+        # Block-diagonal Q/K/V projections
+        self.q_proj = LinearHeadwiseExpand(self.inner_dim, num_proj_heads, bias=bias)
+        self.k_proj = LinearHeadwiseExpand(self.inner_dim, num_proj_heads, bias=bias)
+        self.v_proj = LinearHeadwiseExpand(self.inner_dim, num_proj_heads, bias=bias)
+        
+        # Gates: take concat(q, k, v) as input
+        self.igate = nn.Linear(3 * self.inner_dim, num_heads, bias=True)
+        self.fgate = nn.Linear(3 * self.inner_dim, num_heads, bias=True)
+        
+        # Output normalization (per-head group norm)
+        self.outnorm = nn.GroupNorm(num_heads, self.inner_dim, affine=True)
+        
+        # Down-projection: inner_dim -> D
+        self.proj_down = nn.Linear(self.inner_dim, dim, bias=bias)
+        
+        # Learnable skip connection and layer scale
+        self.learnable_skip = nn.Parameter(torch.ones(self.inner_dim))
+        self.layerscale = nn.Parameter(torch.ones(self.inner_dim))
+        
+        self._reset_gate_bias()
+    
+    def _reset_gate_bias(self):
+        """Initialize forget gate bias high (encourages remembering) and input gate low."""
+        with torch.no_grad():
+            nn.init.zeros_(self.igate.bias)
+            # Forget gate bias: initialize to encourage remembering
+            nn.init.constant_(self.fgate.bias, 3.0)
+    
+    def forward(self, x: torch.Tensor, reverse: bool = False) -> torch.Tensor:
+        """
+        Args:
+            x: (B, S, D) input sequence
+            reverse: if True, process sequence right-to-left (for bidirectional scanning)
+        Returns:
+            (B, S, D) output
+        """
+        B, S, D = x.shape
+        
+        if reverse:
+            x = x.flip(1)
+        
+        # 1. Up-project to 2*inner_dim
+        up = self.proj_up(x)  # (B, S, 2*inner)
+        x_mlstm = up[..., :self.inner_dim]   # mLSTM branch
+        z = up[..., self.inner_dim:]          # output gate branch
+        
+        # 2. Causal conv1d on mLSTM branch
+        x_conv = self.conv1d(x_mlstm.transpose(1, 2))  # (B, inner, S+pad)
+        x_conv = x_conv[..., :S].transpose(1, 2)       # causal: keep first S
+        x_conv_act = F.gelu(x_conv)
+        
+        # 3. Q/K/V projections (block-diagonal, very lightweight)
+        q = self.q_proj(x_conv_act)   # (B, S, inner)
+        k = self.k_proj(x_conv_act)   # (B, S, inner)
+        v = self.v_proj(x_mlstm)      # V from pre-conv branch
+        
+        # 4. Gates from concat(q, k, v)
+        qkv_cat = torch.cat([q, k, v], dim=-1)  # (B, S, 3*inner)
+        i_gate = self.igate(qkv_cat)    # (B, S, num_heads)
+        f_gate = self.fgate(qkv_cat)    # (B, S, num_heads)
+        
+        # Stabilized gates
+        i_tilde = torch.exp(i_gate)                          # (B, S, H)
+        f_tilde = torch.sigmoid(f_gate)                      # (B, S, H)
+        # Log-space stabilization
+        log_f = torch.log(f_tilde.clamp(min=1e-6))           # (B, S, H)
+        
+        # 5. Reshape Q/K/V for multi-head matrix memory
+        q = rearrange(q, 'b s (h d) -> b h s d', h=self.num_heads)  # (B, H, S, D_h)
+        k = rearrange(k, 'b s (h d) -> b h s d', h=self.num_heads)
+        v = rearrange(v, 'b s (h d) -> b h s d', h=self.num_heads)
+        
+        # 6. Parallel mLSTM computation (log-space stabilized)
+        # Cumulative sum of log forget gates for parallel scan
+        log_f_cumsum = torch.cumsum(log_f.permute(0, 2, 1), dim=-1)  # (B, H, S)
+        
+        # Compute pairwise log forget gate differences for attention-like matrix
+        # log_f_cumsum[:,:,j] - log_f_cumsum[:,:,i] gives product of f gates from i+1 to j
+        log_D = log_f_cumsum.unsqueeze(-1) - log_f_cumsum.unsqueeze(-2)  # (B, H, S, S)
+        
+        # Causal mask: only attend to past positions
+        causal_mask = torch.tril(torch.ones(S, S, device=x.device, dtype=torch.bool))
+        log_D = log_D.masked_fill(~causal_mask, -1e9)
+        
+        # Add input gate contribution
+        i_tilde_perm = i_tilde.permute(0, 2, 1)  # (B, H, S)
+        log_D = log_D + torch.log(i_tilde_perm.clamp(min=1e-6)).unsqueeze(-2)  # broadcast over queries
+        
+        # Stabilize with max trick
+        max_log_D = log_D.max(dim=-1, keepdim=True).values.clamp(min=-10)
+        D = torch.exp(log_D - max_log_D)  # (B, H, S, S)
+        D = D.masked_fill(~causal_mask, 0.0)
+        
+        # Compute attention: h = D @ v, then normalize by D @ k·q
+        # Scale queries
+        q_scaled = q / math.sqrt(self.head_dim)
+        
+        # Attention scores: (q @ k^T) * D
+        attn = torch.matmul(q_scaled, k.transpose(-1, -2)) * D  # (B, H, S, S)
+        
+        # Normalizer
+        normalizer = attn.sum(dim=-1, keepdim=True).clamp(min=1.0)
+        attn = attn / normalizer
+        
+        # Output
+        h = torch.matmul(attn, v)  # (B, H, S, D_h)
+        h = rearrange(h, 'b h s d -> b s (h d)')
+        
+        # 7. Output norm
+        h = self.outnorm(h.transpose(1, 2)).transpose(1, 2)  # GroupNorm on channel dim
+        
+        # 8. Skip connection + output gate
+        h_skip = h + self.learnable_skip * x_conv_act
+        output = h_skip * torch.sigmoid(z)  # output gate
+        
+        # 9. Down-project + layer scale
+        output = self.proj_down(output)
+        output = output * self.layerscale[:self.dim]  # Note: layerscale is inner_dim sized, we need dim
+        
+        if reverse:
+            output = output.flip(1)
+        
+        return output
+
+
+class SwiGLUMLP(nn.Module):
+    """SwiGLU MLP as used in ViL blocks.
+    
+    SwiGLU(x) = (W1·x ⊙ Swish(V·x)) then W2·hidden → output
+    """
+    def __init__(self, dim: int, mlp_ratio: float = 4.0, bias: bool = False, drop: float = 0.0):
+        super().__init__()
+        hidden_dim = int(dim * mlp_ratio)
+        # SwiGLU: two parallel projections, one gated
+        self.w1 = nn.Linear(dim, hidden_dim, bias=bias)     # value path
+        self.w2 = nn.Linear(hidden_dim, dim, bias=bias)     # down projection
+        self.v = nn.Linear(dim, hidden_dim, bias=bias)       # gate path
+        self.drop = nn.Dropout(drop)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.drop(self.w2(F.silu(self.v(x)) * self.w1(x)))
+
+
+class mLSTMBlock(nn.Module):
+    """Single ViL block: LayerNorm → mLSTMCell → residual.
+    
+    Following the official ViL-S architecture, there is NO separate MLP/FFN layer.
+    The gated output (proj_up → split → z-gate → proj_down) inside the mLSTMCell
+    already performs the role of dimension expansion + nonlinearity + projection.
+    
+    This matches ViL-S: ~0.92M params per block, 24 blocks ≈ 22M backbone.
+    """
+    def __init__(
+        self,
+        dim: int = 384,
+        proj_factor: float = 2.0,
+        qkv_proj_blocksize: int = 4,
+        num_heads: int = 4,
+        conv_kernel: int = 4,
+        mlp_ratio: float = 4.0,  # kept for API compat but unused in standard blocks
+        drop_path: float = 0.0,
+        bias: bool = False,
+    ):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, bias=False)
+        self.mlstm = mLSTMCell(
+            dim=dim,
+            proj_factor=proj_factor,
+            qkv_proj_blocksize=qkv_proj_blocksize,
+            num_heads=num_heads,
+            conv_kernel=conv_kernel,
+            bias=bias,
+        )
+        self.drop_path = StochasticDepth(drop_path)
+    
+    def forward(self, x: torch.Tensor, reverse: bool = False) -> torch.Tensor:
+        x = x + self.drop_path(self.mlstm(self.norm1(x), reverse=reverse))
+        return x