State Space Models (SSM) and Mamba: The Post-Transformer Architecture

import torch
import torch.nn as nn

class ContinuousSSM(nn.Module):
    """
    Continuous-time State Space Model.
    
    The dynamics: dh/dt = Ah + Bu
    The output: y = Ch + Du
    """
    def __init__(self, d_model, d_state=16):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        
        # State transition matrix A (discretized later)
        self.A = nn.Parameter(torch.randn(d_model, d_state))
        self.B = nn.Parameter(torch.randn(d_model, d_state))
        self.C = nn.Parameter(torch.randn(d_model, d_state))
        self.D = nn.Parameter(torch.zeros(d_model))  # Usually zero
        
    def continuous_dynamics(self, h, u):
        """Compute continuous state derivative: dh/dt = Ah + Bu"""
        # Project state to state space
        state_proj = h @ self.A.unsqueeze(0)  # [batch, d_state]
        # Compute derivative
        dh = state_proj + u @ self.B.unsqueeze(0)
        return dh
    
    def output(self, h):
        """Compute output: y = Ch + Du"""
        # Project state to output
        output = h @ self.C.unsqueeze(0)
        return output

def discretize_ssm(A, B, delta_t):
    """
    Discretize continuous SSM parameters using ZOH.
    
    Args:
        A: Continuous state transition matrix [d_model, d_state]
        B: Continuous input matrix [d_model, d_state]
        delta_t: Time step
    
    Returns:
        discretized A_bar, B_bar
    """
    # For simplicity, use first-order approximation
    # In practice, use more sophisticated discretization
    A_bar = A * delta_t + torch.eye(A.size(0), device=A.device)
    B_bar = B * delta_t
    
    return A_bar, B_bar


class DiscreteSSM(nn.Module):
    """
    Discrete-time State Space Model.
    """
    def __init__(self, d_model, d_state=16, delta_t=0.1):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.delta_t = delta_t
        
        # Learn continuous parameters
        self.A_log = nn.Parameter(torch.randn(d_model, d_state))
        self.B = nn.Parameter(torch.randn(d_model, d_state))
        self.C = nn.Parameter(torch.randn(d_model, d_state))
        
    def forward(self, u):
        """
        Forward pass through discrete SSM.
        
        Args:
            u: Input sequence [batch, seq_len, d_model]
        
        Returns:
            output: [batch, seq_len, d_model]
        """
        batch, seq_len, _ = u.shape
        
        # Discretize
        A = torch.exp(self.A_log)  # Ensure stability
        A_bar, B_bar = discretize_ssm(A, self.B, self.delta_t)
        
        # Initialize state
        h = torch.zeros(batch, self.d_state, device=u.device)
        
        outputs = []
        for t in range(seq_len):
            # State update: h = A_bar * h + B_bar * u
            h = h @ A_bar.T + u[:, t] @ B_bar.T
            
            # Output: y = C * h
            y = h @ self.C.T
            outputs.append(y)
            
        return torch.stack(outputs, dim=1)

class SelectiveSSM(nn.Module):
    """
    Selective State Space Model - the core of Mamba.
    
    Key innovation: Input-dependent parameters for selective filtering.
    """
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.d_conv = d_conv
        self.d_inner = int(expand * d_model)
        
        # Input projection
        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)
        
        # Convolutional preprocessing
        self.conv1d = nn.Conv1d(
            in_channels=self.d_inner,
            out_channels=self.d_inner,
            kernel_size=d_conv,
            padding=d_conv - 1,
            groups=self.d_inner
        )
        
        # SELECTIVE: Input-dependent parameters
        # This is the key innovation of Mamba
        self.x_proj = nn.Linear(self.d_inner, d_state * 2, bias=False)
        
        # State space parameters (selective)
        self.dt_proj = nn.Linear(d_state, self.d_inner, bias=True)
        
        # A and B matrices (selective)
        self.A_log = nn.Parameter(torch.randn(self.d_inner, d_state))
        self.D = nn.Parameter(torch.ones(self.d_inner))
        
        # Output projection
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
        
    def forward(self, x, state=None):
        """
        Forward pass with selective scanning.
        
        Args:
            x: Input [batch, seq_len, d_model]
            state: Optional previous state [batch, d_inner, d_state]
        
        Returns:
            output: [batch, seq_len, d_model]
        """
        batch, seq_len, _ = x.shape
        device = x.device
        
        # Input projection and split
        xz = self.in_proj(x)  # [batch, seq_len, d_inner * 2]
        x_inner, z = xz.chunk(2, dim=-1)
        
        # Convolution preprocessing
        x_conv = x_inner.transpose(1, 2)  # [batch, d_inner, seq_len]
        x_conv = self.conv1d(x_conv)[:, :, :seq_len]
        x_conv = x_conv.transpose(1, 2)  # [batch, seq_len, d_inner]
        x_conv = torch.nn.functional.silu(x_conv)
        
        # SELECTIVE: Compute input-dependent parameters
        # This is the key difference from standard SSMs
        s_params = self.x_proj(x_conv)  # [batch, seq_len, d_state * 2]
        B, C = s_params.chunk(2, dim=-1)  # Each: [batch, seq_len, d_state]
        
        # Selective dt (time step) - depends on input
        dt = self.dt_proj(torch.tanh(C))  # [batch, seq_len, d_inner]
        
        # Ensure positive and make learnable
        dt = torch.nn.functional.softplus(dt)
        
        # Discretize with selective parameters
        A = -torch.exp(self.A_log)  # Ensure stability (negative for decay)
        
        # Selective scan (core computation)
        y = self.selective_scan(
            x_conv, dt, A, B, C, self.D, z
        )
        
        # Output projection with gating
        output = self.out_proj(y * torch.nn.functional.silu(z))
        
        return output
    
    def selective_scan(self, x, dt, A, B, C, D, z):
        """
        Selective scan algorithm - core of Mamba.
        
        This is where the magic happens - computing SSM
        with input-dependent parameters efficiently.
        """
        # Simplified implementation
        # In practice, use CUDA kernels for efficiency
        batch, seq_len, d_inner = x.shape
        _, _, d_state = B.shape
        
        # Initialize state
        h = torch.zeros(batch, d_inner, d_state, device=x.device)
        
        outputs = []
        for t in range(seq_len):
            # Input at this timestep
            x_t = x[:, t]  # [batch, d_inner]
            
            # Compute state update
            # h' = (I + dt*A) * h + dt*B*x
            dtA = dt[:, t].unsqueeze(-1) * A.unsqueeze(0)  # [batch, d_inner, d_state]
            dtB = dt[:, t].unsqueeze(-1) * B[:, t].unsqueeze(1)  # [batch, d_inner, d_state]
            
            h = h + dtB * x_t.unsqueeze(-1)  # [batch, d_inner, d_state]
            h = h * (1 + dtA)  # State update with decay
            
            # Output: y = C * h + D * x
            y_t = torch.einsum('bds,bsd->bd', h, C[:, t]) + D * x_t
            outputs.append(y_t)
            
        return torch.stack(outputs, dim=1)

class MambaBlock(nn.Module):
    """
    Complete Mamba Block.
    
    Combines:
    1. Input projection
    2. Depthwise convolution
    3. Selective SSM
    4. Output projection with gating
    """
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2, 
                 dropout=0.0, bias=False):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.d_conv = d_conv
        self.d_inner = int(expand * d_model)
        
        # Input projection
        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=bias)
        
        # Depthwise convolution
        self.conv1d = nn.Conv1d(
            in_channels=self.d_inner,
            out_channels=self.d_inner,
            kernel_size=d_conv,
            padding=d_conv - 1,
            groups=self.d_inner
        )
        
        # Selective SSM
        self.ssm = SelectiveSSM(d_model, d_state, d_conv, expand)
        
        # Output projection with residual connection
        self.out_proj = nn.Linear(self.d_inner, d_model, bias=bias)
        
    def forward(self, x, residual=None):
        """
        Forward pass.
        
        Args:
            x: [batch, seq_len, d_model]
            residual: Optional residual connection
        
        Returns:
            output: [batch, seq_len, d_model]
        """
        batch, seq_len, _ = x.shape
        
        # Input projection and split
        xz = self.in_proj(x)
        x_inner, z = xz.chunk(2, dim=-1)
        
        # Convolution (in conv format)
        x_conv = x_inner.transpose(1, 2)
        x_conv = self.conv1d(x_conv)[:, :, :seq_len]
        x_conv = torch.nn.functional.silu(x_conv)
        x_conv = x_conv.transpose(1, 2)
        
        # SSM
        y = self.ssm(x_conv)
        
        # Output projection with gating
        output = self.out_proj(y * torch.nn.functional.silu(z))
        
        # Residual if provided
        if residual is not None:
            output = output + residual
            
        return output

def parallel_scan_associative(op, init, tensor):
    """
    Parallel scan using associative scan operation.
    
    Computes: y[i] = op(x[0], x[1], ..., x[i])
    
    Uses divide-and-conquer for O(log n) parallel time.
    
    Args:
        op: Associative binary operation
        init: Initial value
        tensor: Input tensor [batch, seq_len, ...]
    
    Returns:
        Scanned tensor
    """
    seq_len = tensor.size(1)
    device = tensor.device
    
    # Pad to power of 2
    n = 2 ** (seq_len - 1).bit_length()
    padded = torch.zeros(tensor.size(0), n, *tensor.shape[2:], 
                        device=device)
    padded[:, :seq_len] = tensor
    
    # Downsweep phase of parallel scan
    # This is a simplified version - real implementation is more complex
    for d in range(n.bit_length()):
        stride = 2 ** (d + 1)
        idx = torch.arange(0, n, stride, device=device)
        
        # Even elements
        left = padded[:, idx]
        # Odd elements  
        right = padded[:, idx + stride // 2]
        
        # Combine
        combined = op(left, right)
        
        # Write back
        padded[:, idx + stride // 2] = combined
        padded[:, idx] = left  # Keep original for odd positions
        
    return padded[:, :seq_len]


def ssm_scan(A, B, C, D, x):
    """
    Compute SSM output using parallel scan.
    
    y[i] = C @ A^i @ B @ x[0] + C @ A^(i-1) @ B @ x[1] + ... + C @ B @ x[i]
    
    Can be computed as parallel scan of (A, B) with initial state.
    """
    # This is simplified - actual implementation uses
    # specialized CUDA kernels
    batch, seq_len, d_model = x.shape
    _, d_state = B.shape
    
    # Compute AB = A @ B for each timestep
    AB = torch.einsum('bnd,bde->bne', B, A.unsqueeze(0))
    
    # Scan using cumulative sum (simplified)
    # Real implementation: use associative scan with (A, B) as state
    h = torch.zeros(batch, d_model, d_state, device=x.device)
    outputs = []
    
    for t in range(seq_len):
        h = A.unsqueeze(0) @ h + AB[:, t].unsqueeze(1) * x[:, t].unsqueeze(-1)
        y = torch.einsum('bds,bsd->bd', h, C[:, t])
        outputs.append(y)
        
    return torch.stack(outputs, dim=1)

def complexity_comparison(seq_len, d_model=512, d_state=16):
    """Compare computational complexity."""
    
    # Transformer
    transformer_attention = seq_len ** 2 * d_model
    transformer_ffn = seq_len * d_model ** 2
    
    # Mamba
    mamba_ssm = seq_len * d_model * d_state
    mamba_conv = seq_len * d_model * 4
    
    print(f"Sequence Length: {seq_len}")
    print(f"Transformer Attention: {transformer_attention:,} ops")
    print(f"Mamba SSM: {mamba_ssm:,} ops")
    print(f"Speedup: {transformer_attention/mamba_ssm:.1f}x")
    
# Example output for seq_len=4096:
# Transformer Attention: 8,589,934,592 ops
# Mamba SSM: 33,554,432 ops
# Speedup: 256x

class ArchitectureChooser:
    """
    Guide for choosing between Transformer and Mamba.
    """
    
    @staticmethod
    def should_use_mamba(sequence_length, modality='text'):
        """
        Recommend architecture based on task.
        """
        if modality == 'text':
            # Short sequences: Transformer often better
            if sequence_length < 512:
                return "Transformer"
            # Long sequences: Mamba advantages emerge
            else:
                return "Mamba"
                
        elif modality == 'audio' or modality == 'genomics':
            # Long-range dependencies important
            return "Mamba"
            
        elif modality == 'code':
            # Complex attention patterns
            return "Transformer"  # or hybrid
            
        return "Mamba"
    
    @staticmethod
    def hybrid_approach(seq, transformer_layers, mamba_layers):
        """
        Combine both architectures.
        
        Common pattern:
        - Early layers: Transformer (better local patterns)
        - Later layers: Mamba (better long-range)
        """
        pass

class MambaModel(nn.Module):
    """
    Complete Mamba Language Model.
    """
    def __init__(self, vocab_size, d_model=256, n_layers=24, 
                 d_state=16, d_conv=4, expand=2):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        
        # Embedding
        self.embedding = nn.Embedding(vocab_size, d_model)
        
        # Mamba layers
        self.layers = nn.ModuleList([
            MambaBlock(d_model, d_state, d_conv, expand)
            for _ in range(n_layers)
        ])
        
        # Output
        self.norm = nn.RMSNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        
        # Weight tying
        self.lm_head.weight = self.embedding.weight
        
    def forward(self, input_ids, position_ids=None):
        """
        Forward pass.
        
        Args:
            input_ids: [batch, seq_len]
        
        Returns:
            logits: [batch, seq_len, vocab_size]
        """
        x = self.embedding(input_ids)
        
        # Through Mamba layers
        for layer in self.layers:
            x = layer(x)
            
        x = self.norm(x)
        logits = self.lm_head(x)
        
        return logits
    
    def generate(self, prompt, max_length=100, temperature=1.0):
        """
        Autoregressive generation.
        """
        self.eval()
        
        with torch.no_grad():
            # Encode prompt
            input_ids = torch.tensor([[self.vocab_size - 1]] + 
                                    encode(prompt)).unsqueeze(0)
            
            for _ in range(max_length):
                logits = self.forward(input_ids)
                next_token_logits = logits[:, -1] / temperature
                
                # Sample
                probs = torch.softmax(next_token_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
                
                input_ids = torch.cat([input_ids, next_token], dim=1)
                
                if next_token.item() == EOS_TOKEN:
                    break
                    
        return decode(input_ids[0])

class MambaTrainer:
    """Trainer for Mamba models."""
    
    def __init__(self, model, lr=3e-4, weight_decay=0.1):
        self.model = model
        self.optimizer = torch.optim.AdamW(
            model.parameters(), 
            lr=lr, 
            weight_decay=weight_decay
        )
        
    def train_step(self, batch):
        """Single training step."""
        input_ids = batch['input_ids'].to(model.device)
        labels = batch['labels'].to(model.device)
        
        # Forward
        logits = self.model(input_ids)
        
        # Loss
        loss = F.cross_entropy(
            logits.view(-1, self.model.vocab_size),
            labels.view(-1),
            ignore_index=-100
        )
        
        # Backward
        self.optimizer.zero_grad()
        loss.backward()
        
        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
        
        self.optimizer.step()
        
        return loss.item()

class Mamba2Block(nn.Module):
    """
    Mamba 2 introduces:
    1. State space duality (SSM ↔ Attention)
    2. Parallelization improvements
    3. Better hardware utilization
    """
    def __init__(self, d_model, d_state=128, d_head=64):
        super().__init__()
        # Multiple heads for state space
        self.n_heads = d_model // d_head
        self.d_head = d_head
        self.d_state = d_state
        
        # Similar to Mamba but with head dimension
        # ...

class HybridModel(nn.Module):
    """
    Combine Transformer and Mamba layers.
    """
    def __init__(self, vocab_size, d_model=512):
        super().__init__()
        
        # Early layers: Transformer
        self.transformer_layers = nn.ModuleList([
            TransformerBlock(d_model) for _ in range(6)
        ])
        
        # Later layers: Mamba  
        self.mamba_layers = nn.ModuleList([
            MambaBlock(d_model) for _ in range(18)
        ])
        
        # Output
        self.lm_head = nn.Linear(d_model, vocab_size)
        
    def forward(self, x):
        # Transformer first (better local modeling)
        for layer in self.transformer_layers:
            x = layer(x)
            
        # Then Mamba (efficient long-range)
        for layer in self.mamba_layers:
            x = layer(x)
            
        return self.lm_head(x)

class VisionMamba(nn.Module):
    """
    Mamba for Vision (ViM).
    """
    def __init__(self, image_size=224, patch_size=16, d_model=768):
        super().__init__()
        
        # Patch embedding
        self.patch_embed = nn.Conv2d(
            3, d_model, 
            kernel_size=patch_size, 
            stride=patch_size
        )
        
        # Positional embedding
        n_patches = (image_size // patch_size) ** 2
        self.pos_embed = nn.Parameter(torch.randn(1, n_patches, d_model))
        
        # Mamba blocks
        self.blocks = nn.ModuleList([
            MambaBlock(d_model) for _ in range(24)
        ])
        
        # Classification head
        self.norm = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, 1000)
        
    def forward(self, x):
        # Patch embedding
        x = self.patch_embed(x)  # [B, C, H, W]
        x = x.flatten(2).transpose(1, 2)  # [B, N, C]
        
        # Add positional embedding
        x = x + self.pos_embed
        
        # Through Mamba blocks
        for block in self.blocks:
            x = block(x)
            
        x = self.norm(x)
        
        # Global average pooling
        x = x.mean(dim=1)
        
        return self.head(x)

class LongContextLM:
    """Language modeling with 100k+ context length."""
    
    def __init__(self, model_path):
        self.model = MambaModel.from_pretrained(model_path)
        self.model.to('cuda')
        
    def summarize(self, text, max_length=5000):
        """Summarize very long documents."""
        # Mamba handles 100k+ context efficiently
        tokens = tokenize(text)[:100000]
        
        # Use sliding window for very long
        if len(tokens) > 50000:
            # Process in chunks
            summary = ""
            for i in range(0, len(tokens), 40000):
                chunk = tokens[i:i+50000]
                # Overlap for continuity
                summary = self._process_with_overlap(chunk, summary)
        else:
            summary = self.model.generate(tokens)
            
        return summary

class TimeSeriesMamba(nn.Module):
    """Mamba for time series forecasting."""
    
    def __init__(self, n_features, d_model=128, n_layers=4):
        super().__init__()
        
        # Input embedding
        self.input_proj = nn.Linear(n_features, d_model)
        
        # Mamba layers for temporal modeling
        self.layers = nn.ModuleList([
            MambaBlock(d_model) for _ in range(n_layers)
        ])
        
        # Output projection
        self.output_proj = nn.Linear(d_model, n_features)
        
    def forecast(self, x, horizon=24):
        """
        Forecast future values.
        
        Args:
            x: Historical data [batch, seq_len, n_features]
            horizon: Number of steps to forecast
        
        Returns:
            predictions: [batch, horizon, n_features]
        """
        # Encode history
        x = self.input_proj(x)
        
        for layer in self.layers:
            x = layer(x)
            
        # Use last representation to start forecasting
        last = x[:, -1]
        
        predictions = []
        for _ in range(horizon):
            # Project and add time encoding
            pred = self.output_proj(last)
            predictions.append(pred)
            
            # Update (simplified - real impl would be more sophisticated)
            
        return torch.stack(predictions, dim=1)

class GenomicsMamba(nn.Module):
    """Mamba for DNA sequence analysis."""
    
    def __init__(self, vocab_size=4, d_model=512, d_state=64):
        super().__init__()
        
        # DNA embeddings
        self.embedding = nn.Embedding(vocab_size, d_model)
        
        # Mamba layers
        self.layers = nn.ModuleList([
            MambaBlock(d_model, d_state=d_state) 
            for _ in range(32)
        ])
        
        # Task heads
        self.norm = nn.LayerNorm(d_model)
        self.classifier = nn.Linear(d_model, 2)  # e.g., promoter/not
        
    def identify_promoters(self, dna_sequence):
        """Identify promoter regions in DNA."""
        x = self.embedding(dna_sequence)
        
        for layer in self.layers:
            x = layer(x)
            
        x = self.norm(x.mean(dim=1))
        return self.classifier(x)

def init_mamba(model):
    """Proper initialization for Mamba models."""
    for name, param in model.named_parameters():
        if 'A_log' in name:
            # Initialize A for stable dynamics
            nn.init.constant_(param, 0)
        elif 'B' in name or 'C' in name:
            # Xavier initialization
            nn.init.xavier_uniform_(param)
        elif 'dt_proj' in name:
            # Small initial time steps
            nn.init.uniform_(param, -0.1, 0.1)

DEFAULT_MAMBA_CONFIG = {
    'd_model': 512,
    'd_state': 128,  # Larger = more capacity, slower
    'd_conv': 4,     # Convolution kernel size
    'expand': 2,      # FFN expansion factor
    'dropout': 0.0,
}