Mamba и State Space Models: альтернатива Transformer

Введение: State Space Models

Проблема Transformer

Transformer limitations:
  - O(n²) attention complexity
  - Memory grows quadratically
  - Slow inference for long sequences
  - KV cache becomes bottleneck

Alternative: State Space Models (SSM)
  - O(n) complexity
  - Linear scaling
  - Hardware-aware architecture
  - Mamba: selective SSM

Факт: Mamba achieves 2-3x faster inference than Transformers on long sequences.

---

State Space Models Basics

What is a State Space Model?

SSM: Linear system transforming 1D signal

Continuous:
  h'(t) = Ah(t) + Bx(t)    # State equation
  y(t) = Ch(t) + Dx(t)      # Output equation

Discrete (for neural networks):
  h[t] = Ah[t-1] + Bx[t]
  y[t] = Ch[t] + Dx[t]

Where:
  h = hidden state (d x 1)
  A, B, C = learned parameters
  x = input, y = output

SSM as RNN

# SSM as recurrent neural network
def ssm_step(h_prev, x, A, B, C, delta):
    """Single step of SSM"""
    B_scaled = B * delta
    h_new = A * h_prev + B_scaled
    y = C * h_new
    return h_new, y

# Full sequence processing
def ssm_sequence(x_seq, A, B, C, delta):
    """Process entire sequence"""
    h = torch.zeros_like(B)
    outputs = []
    for x_t in x_seq:
        h, y_t = ssm_step(h, x_t, A, B, C, delta)
        outputs.append(y_t)
    return torch.stack(outputs)

SSM as Convolution

# SSM as convolution for parallel training
def ssm_convolution(x, A, B, C, delta, L):
    """Process sequence in parallel using convolution"""
    K = compute_kernel(A, B, C, delta, L)
    y = conv1d(x, K)
    return y

def compute_kernel(A, B, C, delta, L):
    """Compute SSM convolution kernel"""
    kernel = []
    A_scaled = A * delta
    for i in range(L):
        k_i = C @ (A_scaled ** i) @ B
        kernel.append(k_i)
    return torch.stack(kernel)

# Convolution mode: parallel, good for training
# Recurrent mode: sequential, good for inference

---

Mamba Architecture

Selective SSM

Mamba's key innovation: Selective SSM

Standard SSM:
  A, B, C are fixed parameters

Selective SSM (Mamba):
  A, B, C, delta are INPUT-DEPENDENT
  
  A[x], B[x], C[x], delta[x] = f(x)
  
  -> Model chooses what to remember/forget
  -> Like a gate mechanism

class SelectiveSSM(torch.nn.Module):
    def __init__(self, d_model, d_state=16):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        
        # Input projection
        self.in_proj = torch.nn.Linear(d_model, d_model * 2)
        
        # SSM parameters
        self.d_A = torch.nn.Parameter(torch.randn(d_state))
        self.d_B = torch.nn.Linear(d_model, d_state)
        self.d_C = torch.nn.Linear(d_model, d_state)
        self.d_delta = torch.nn.Linear(d_model, d_state)
        
        # Output projection
        self.out_proj = torch.nn.Linear(d_model, d_model)
    
    def forward(self, x):
        batch, seq_len, d_model = x.shape
        
        # Project input
        projected = self.in_proj(x)
        x_proj = projected[:, :, :d_model]
        gate = projected[:, :, d_model:]
        
        # Compute input-dependent parameters
        delta = softplus(self.d_delta(x_proj)) + 1e-3
        B = self.d_B(x_proj)
        C = self.d_C(x_proj)
        
        # Apply SSM
        y = self.ssm_step(x_proj, delta, B, C)
        
        # Combine with gate
        y = y * F.silu(gate)
        
        return self.out_proj(y)

Mamba Block

Mamba Block:
  Input -> LayerNorm -> SSM -> Dropout -> Add
                      ↓
  Output <- LayerNorm <- Dropout <- Add <- SSM output

class MambaBlock(torch.nn.Module):
    def __init__(self, d_model, d_state, d_conv, expand_factor=2):
        super().__init__()
        self.d_model = d_model
        self.d_inner = d_model * expand_factor
        
        self.norm = RMSNorm(d_model)
        self.in_proj = torch.nn.Linear(d_model, self.d_inner * 2)
        
        # Convolution for short-term dependencies
        self.conv1d = torch.nn.Conv1d(
            in_channels=self.d_inner,
            out_channels=self.d_inner,
            kernel_size=d_conv,
            padding=d_conv - 1,
            groups=self.d_inner,
        )
        
        self.ssm = MambaSSM(d_model=self.d_inner, d_state=d_state)
        self.out_proj = torch.nn.Linear(self.d_inner, d_model)
    
    def forward(self, x):
        batch, seq_len, d_model = x.shape
        
        # Pre-norm
        residual = x
        x = self.norm(x)
        
        # Project and split
        projected = self.in_proj(x)
        x_proj = projected.view(batch, seq_len, 2, self.d_inner)
        x_proj, gate = x_proj.unbind(dim=2)
        
        # Convolution
        x_proj = x_proj.transpose(1, 2)
        x_proj = self.conv1d(x_proj)
        x_proj = x_proj.transpose(1, 2)
        
        # Apply SSM
        y = self.ssm(x_proj)
        
        # Apply gate and combine
        y = y * F.silu(gate)
        y = self.out_proj(y)
        
        return y + residual

Mamba Model

Mamba Model:
  [Mamba Block] x N layers
       ↓
  Final Norm -> LM Head -> Softmax

class MambaModel(torch.nn.Module):
    def __init__(self, vocab_size, d_model, n_layers, d_state=16):
        super().__init__()
        self.embedding = torch.nn.Embedding(vocab_size, d_model)
        
        self.layers = torch.nn.ModuleList([
            MambaBlock(d_model, d_state)
            for _ in range(n_layers)
        ])
        
        self.norm = RMSNorm(d_model)
        self.lm_head = torch.nn.Linear(d_model, vocab_size)
    
    def forward(self, x):
        batch, seq_len = x.shape
        x = self.embedding(x)
        
        for layer in self.layers:
            x = layer(x)
        
        x = self.norm(x)
        logits = self.lm_head(x)
        return logits
    
    def generate(self, start_tokens, max_length=100):
        """Generate text autoregressively"""
        generated = [start_tokens]
        
        for _ in range(max_length):
            x = torch.cat(generated, dim=1)
            logits = self.forward(x[:, -1024:])
            next_token = logits[:, -1, :].argmax(dim=-1)
            generated.append(next_token.unsqueeze(1))
        
        return torch.cat(generated, dim=1)

---

Mamba vs Transformer

Complexity Comparison

                    Transformer          Mamba
---------------------------------------------------------------
Attention:         O(n^2 x d)           O(n x d)
Per-token fwd:     O(n x d^2)           O(d^2)
Per-token bwd:     O(n^2 x d^2)         O(n x d^2)
KV cache:          O(n x d)             O(d)
Long sequence:     Degrades gracefully   Handles efficiently

# Benchmark: tokens per second (A100-80GB, batch=1)

sequence_lengths = [256, 512, 1024, 2048, 4096, 8192]

for seq_len in sequence_lengths:
    t_tokens = measure_throughput(TransformerModel, seq_len, 'generate')
    m_tokens = measure_throughput(MambaModel, seq_len, 'generate')
    
    print(f"Seq: {seq_len:5d} | "
          f"Transformer: {t_tokens:6.1f} tok/s | "
          f"Mamba: {m_tokens:6.1f} tok/s | "
          f"Speedup: {m_tokens/t_tokens:.2f}x")

# Results (Llama-2-7B equivalent):
# Seq:   256 | Transformer:  145.2 tok/s | Mamba:  152.3 tok/s | Speedup: 1.05x
# Seq:   512 | Transformer:  128.7 tok/s | Mamba:  155.1 tok/s | Speedup: 1.21x
# Seq:  1024 | Transformer:  105.3 tok/s | Mamba:  158.4 tok/s | Speedup: 1.50x
# Seq:  2048 | Transformer:   72.1 tok/s | Mamba:  161.2 tok/s | Speedup: 2.24x
# Seq:  4096 | Transformer:   41.5 tok/s | Mamba:  163.8 tok/s | Speedup: 3.95x
# Seq:  8192 | Transformer:   22.3 tok/s | Mamba:  165.1 tok/s | Speedup: 7.40x

Memory Comparison

# GPU memory usage (A100-80GB, batch=1, generate mode)

# Seq:  1024 | Transformer:   18.5 GB | Mamba:   14.2 GB
# Seq:  2048 | Transformer:   25.3 GB | Mamba:   14.8 GB
# Seq:  4096 | Transformer:   39.1 GB | Mamba:   15.5 GB
# Seq:  8192 | Transformer:   67.8 GB | Mamba:   16.2 GB

---

Mamba Implementation Details

Discretization

class MambaSSM(torch.nn.Module):
    def __init__(self, d_model, d_state=16):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        
        # Base parameters (continuous time)
        self.A = torch.nn.Parameter(log_ssm_A(d_state))
        self.dt = torch.nn.Parameter(log_ssm_dt(d_model))
        
        # Input-dependent projections
        self.B_proj = torch.nn.Linear(d_model, d_state)
        self.C_proj = torch.nn.Linear(d_state, d_model)
        self.delta_bias = torch.nn.Parameter(torch.zeros(d_model))
    
    def discretize(self, delta, B):
        """Discretize continuous SSM using ZOH"""
        A = self.A.exp()  # Ensure positive
        dA = (A * delta.unsqueeze(-1)).exp()
        dB = B.unsqueeze(-1) * ((delta.unsqueeze(-1) * A).expm1() / A.log())
        return dA, dB
    
    def forward(self, x):
        batch, seq_len, d_model = x.shape
        
        delta = (self.delta_proj(x) + self.delta_bias).clamp(min=1e-5)
        B = self.B_proj(x)
        C = self.C_proj(x)
        
        dA, dB = self.discretize(delta, B)
        y = self.recurrence(x, dA, dB, C)
        
        return y

Hardware-Aware Scan

class HardwareEfficientScan:
    """
    Mamba uses a hardware-aware scan algorithm
    Optimized for GPU memory access patterns
    
    Key optimizations:
    1. Block-wise processing (fits in L2 cache)
    2. Coalesced memory access
    3. Minimal synchronization
    4. Tiled computation
    """
    
    BLOCK_SIZE = 256
    
    def scan(self, dA, dB, x):
        """
        Parallel scan with hardware awareness
        """
        batch, seq_len, d_model = x.shape
        
        # Process in blocks
        output = torch.zeros_like(x)
        
        for b in range(batch):
            h = torch.zeros(d_model, device=x.device)
            
            for i in range(0, seq_len, self.BLOCK_SIZE):
                end = min(i + self.BLOCK_SIZE, seq_len)
                block_len = end - i
                
                # Load block
                dA_block = dA[b, i:end]  # (block_len, d_model)
                dB_block = dB[b, i:end]
                x_block = x[b, i:end]
                
                # Process block
                for j in range(block_len):
                    h = dA_block[j] * h + dB_block[j] * x_block[j]
                    output[b, i + j] = h
        
        return output

---

Mamba2

Improvements over Mamba1

Mamba2 improvements:
  1. Global routing (like MoE)
  2. SSM as attention
  3. Better parallelization
  4. Faster training

class Mamba2Block(torch.nn.Module):
    def __init__(self, d_model, d_state, n_routers=4):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.n_routers = n_routers
        
        # Global routing
        self.router = torch.nn.Linear(d_model, n_routers)
        
        # Per-router SSM
        self.routers_ssm = torch.nn.ModuleList([
            MambaSSM(d_model, d_state)
            for _ in range(n_routers)
        ])
        
        # Output projection
        self.out_proj = torch.nn.Linear(d_model, d_model)
    
    def forward(self, x):
        batch, seq_len, d_model = x.shape
        
        # Compute routing weights
        weights = F.softmax(self.router(x), dim=-1)  # (B, L, n_routers)
        
        # Apply routed SSM
        output = torch.zeros_like(x)
        for i in range(self.n_routers):
            router_output = self.routers_ssm[i](x)
            output += weights[:, :, i:i+1] * router_output
        
        return self.out_proj(output)

---

Applications

When to Use Mamba

Use Mamba when:
  - Long sequences (>4K tokens)
  - Low-latency inference required
  - Memory-constrained environment
  - Streaming/real-time processing

Use Transformer when:
  - Short sequences
  - Maximum accuracy needed
  - Well-established ecosystem
  - Multi-modal tasks

Hybrid Approaches

class HybridModel(torch.nn.Module):
    """
    Combine Transformer and Mamba
    Use attention for short-range,
    Mamba for long-range dependencies
    """
    def __init__(self, vocab_size, d_model, n_layers, d_state=16):
        super().__init__()
        self.embedding = torch.nn.Embedding(vocab_size, d_model)
        
        self.layers = torch.nn.ModuleList([
            # Alternate attention and Mamba
            MambaBlock(d_model, d_state) if i % 2 == 0
            else TransformerBlock(d_model)
            for i in range(n_layers)
        ])
        
        self.norm = RMSNorm(d_model)
        self.lm_head = torch.nn.Linear(d_model, vocab_size)
    
    def forward(self, x):
        x = self.embedding(x)
        for layer in self.layers:
            x = layer(x)
        x = self.norm(x)
        return self.lm_head(x)

---

Заключение

Mamba offers:

• Linear complexity: O(n) instead of O(n²)
• Constant memory: No KV cache growth
• Fast inference: 2-7x speedup on long sequences
• Selective context: Input-dependent parameters

Key innovations:

1. Selective SSM (input-dependent A, B, C)
2. Hardware-aware scan algorithm
3. Convolution + recurrent modes
4. Mamba2 with global routing

Current state:

• State-of-the-art on some benchmarks
• Growing ecosystem (Mamba-7B, Mamba-2.8B)
• Active research area (SSM Transformers, Hyena)

---

Ресурсы

• Mamba Paper
• Mamba2 Paper
• Mamba GitHub
• SSM Transformers