Mamba: State Space Models Challenge Transformers

Mamba represents a paradigm shift in sequence modeling—achieving transformer-quality results with linear scaling in sequence length, no attention mechanism, and constant memory during inference.

The Attention Problem

Transformers scale quadratically: O(N²) for sequence length N.

Attention complexity:
N = 1K   →  1M operations
N = 10K  →  100M operations  
N = 100K →  10B operations

Memory also scales O(N²) for KV cache

State Space Models: The Alternative

SSMs process sequences through a continuous dynamical system:

h'(t) &= Ah(t) + Bx(t) \\ y(t) &= Ch(t) + Dx(t) \end{aligned}$$ Discretized for sequences: $$\begin{aligned} h_k &= \bar{A}h_{k-1} + \bar{B}x_k \\ y_k &= Ch_k + Dx_k \end{aligned}$$ **Key insight**: This is a linear recurrence—computable in O(N) time! ## Mamba's Innovation: Selection Previous SSMs (S4, H3) used **fixed** A, B, C matrices. Mamba makes them **input-dependent**: ```python class MambaBlock(nn.Module): def __init__(self, d_model, d_state=16, expand=2): self.d_state = d_state d_inner = d_model * expand # Projections self.in_proj = nn.Linear(d_model, d_inner * 2) self.out_proj = nn.Linear(d_inner, d_model) # Selection mechanism (input-dependent) self.x_proj = nn.Linear(d_inner, d_state * 2 + 1) # Δ, B, C # Fixed parameters self.A = nn.Parameter(torch.randn(d_inner, d_state)) self.D = nn.Parameter(torch.ones(d_inner)) def forward(self, x): # x: [batch, seq_len, d_model] # Project and split xz = self.in_proj(x) x, z = xz.chunk(2, dim=-1) # Input-dependent parameters (THE KEY INNOVATION) x_params = self.x_proj(x) delta, B, C = x_params.split([1, self.d_state, self.d_state], dim=-1) delta = F.softplus(delta) # Selective SSM computation y = selective_scan(x, delta, self.A, B, C, self.D) # Gate and project out y = y * F.silu(z) return self.out_proj(y) ``` ## Selective Scan Algorithm The magic happens here—parallel prefix sum enables O(N) computation: ```python def selective_scan(x, delta, A, B, C, D): """ x: input [B, L, D] delta: step size [B, L, D, 1] A: state matrix [D, N] B, C: input-dependent [B, L, D, N] D: skip connection [D] """ batch, seq_len, d_inner = x.shape n_state = A.shape[1] # Discretize A deltaA = torch.exp(delta.unsqueeze(-1) * A) # [B, L, D, N] deltaB = delta.unsqueeze(-1) * B # [B, L, D, N] # Recurrence (can be parallelized with associative scan) h = torch.zeros(batch, d_inner, n_state) ys = [] for i in range(seq_len): h = deltaA[:, i] * h + deltaB[:, i] * x[:, i:i+1] y = (h * C[:, i]).sum(-1) ys.append(y) y = torch.stack(ys, dim=1) return y + x * D ``` ## Comparison | Aspect | Transformer | Mamba | |--------|-------------|-------| | Sequence scaling | O(N²) | O(N) | | Memory (inference) | O(N) KV cache | O(1) state | | Long-range | Excellent | Excellent | | In-context learning | Excellent | Good | | Hardware utilization | Good (matmuls) | Requires custom kernels | ## Mamba-2: Improved Architecture Mamba-2 refines the design for better hardware efficiency: ```python class Mamba2Block(nn.Module): def __init__(self, d_model, n_heads=8, d_head=64, d_state=128): # Multi-head structure (like attention) self.n_heads = n_heads self.d_head = d_head # Structured state space (SSD) layer self.ssd = SSD(d_model, n_heads, d_head, d_state) def forward(self, x): # Reshape to heads # Apply SSD (structured state space duality) # Combines linear attention and SSM benefits pass ``` ## Training Considerations ### Advantages - Linear memory scaling enables huge contexts - Efficient training with parallel scan - Strong performance on language modeling ### Challenges - Custom CUDA kernels required for efficiency - Different inductive biases than attention - Still maturing ecosystem ## Applications | Domain | Why Mamba Excels | |--------|------------------| | Long documents | Linear scaling with length | | Genomics | DNA sequences are very long | | Audio | High sample rates = long sequences | | Video | Frame sequences | | Time series | Natural fit for recurrent structure | ## Code Example (Hugging Face) ```python from mamba_ssm import Mamba # Create Mamba layer mamba = Mamba( d_model=1024, d_state=16, d_conv=4, expand=2 ) # Use like any other layer x = torch.randn(batch_size, seq_len, d_model) y = mamba(x) ``` ## Hybrid Architectures Combining Mamba with attention: ``` Jamba (AI21): ├── Mamba layers (efficient) ├── Attention layers (where needed) └── MoE layers (capacity) Result: Best of both worlds ``` ## Future Directions 1. **Better hardware support**: Native Mamba ops in frameworks 2. **Pre-trained models**: Mamba foundation models 3. **Multimodal**: Vision and audio Mamba 4. **Retrieval**: Combining with external memory ## References - [Mamba Paper](https://arxiv.org/abs/2312.00752) - [Mamba-2 Paper](https://arxiv.org/abs/2405.21060) - [S4: Efficiently Modeling Long Sequences](https://arxiv.org/abs/2111.00396) --- *Mamba challenges the assumption that attention is all you need—sometimes a well-designed recurrence is all you need.*

Mamba State Space Models