Mamba: The State Space Model That Challenges the Transformer

The Transformer's Long-Sequence Problem

Standard attention scales as O(N²) in both memory and compute with sequence length N. A 1M-token context with standard attention requires roughly 1 trillion attention operations per layer. Even with FlashAttention's IO efficiency, this remains compute-expensive and limits practical context lengths. Mamba (arXiv:2312.00752) proposes a fundamentally different architecture.

State Space Models: The Foundation

SSMs model sequences as continuous dynamical systems. The core recurrence is:

h'(t) = Ah(t) + Bx(t) y(t) = Ch(t)

where h is a hidden state, x is input, y is output, and A/B/C are learned matrices. Discretized to sequences, this becomes a linear recurrence that can be computed as a convolution during training (parallel, like attention) but executed as a recurrence during inference (constant compute per step, unlike attention which must attend to all previous tokens).

The Selective Mechanism in Mamba

Prior SSMs like S4 used time-invariant parameters - A, B, C were fixed regardless of the input token. This made them less effective at content-aware selection (remembering the right things and forgetting others). Mamba makes B and C input-dependent: the model learns to selectively update its hidden state based on the current input. This selectivity is what enables Mamba to match Transformer quality on language tasks.

Hardware-Parallel Scan Algorithm

Making B and C input-dependent breaks the convolution computation path. Mamba uses a parallel scan algorithm (also known as prefix sum) to compute the recurrence in O(N log N) operations in a way that is highly parallelizable on modern GPUs. The scan is implemented with custom CUDA kernels similar in spirit to FlashAttention's tiling approach.

# Simplified Mamba block structure
import torch
import torch.nn as nn

class MambaBlock(nn.Module):
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2):
        super().__init__()
        d_inner = d_model * expand
        self.in_proj = nn.Linear(d_model, d_inner * 2)
        self.conv1d = nn.Conv1d(d_inner, d_inner, d_conv, groups=d_inner, padding=d_conv-1)
        self.x_proj = nn.Linear(d_inner, d_state * 2 + 1)  # B, C, dt
        self.out_proj = nn.Linear(d_inner, d_model)

    def forward(self, x):
        # x: (batch, seq_len, d_model)
        xz = self.in_proj(x)
        x, z = xz.chunk(2, dim=-1)
        # Apply selective SSM (simplified)
        x = self.conv1d(x.transpose(1, 2)).transpose(1, 2)
        return self.out_proj(x * torch.sigmoid(z))

Mamba-2 Improvements

Mamba-2 (arXiv:2405.21060) showed that the selective SSM is a special case of a broader class called structured state space duality (SSD), which connects SSMs and linear attention. Mamba-2 uses a simplified block structure (scalar instead of matrix selective parameters) that runs 2-8x faster than Mamba-1 on modern hardware.

Where Mamba Wins vs Transformers

Mamba excels at: very long sequences (>100k tokens), streaming inference (constant memory per step), tasks requiring selective state tracking (DNA sequences, audio, time series). Transformers remain superior for tasks requiring precise in-context learning and retrieval from a fixed context window.

Mamba in Production

Jamba (AI21 Labs) is a hybrid model combining Transformer layers and Mamba layers, using attention for local context and SSM for long-range dependencies. This hybrid approach captures the strengths of both architectures.

Further Reading

• Mamba paper arXiv:2312.00752
• Mamba-2 arXiv:2405.21060
• Mamba GitHub