Mamba: Linear-Time Sequence Modeling with Selective State Spaces¶

Authors: Albert Gu, Tri Dao (Carnegie Mellon, Princeton) Published: December 2023 Paper: arxiv.org/abs/2312.00752

Why This Matters¶

Mamba is the first serious challenger to Transformer dominance in 7 years:

⚡ Linear time complexity - O(n) vs Transformer's O(n²)
🚀 5× faster inference at long sequences (vs Transformers)
📈 Matches Transformer quality - Same perplexity on language modeling
💾 Constant memory for generation - vs growing KV cache
🔮 Promising alternative - May replace Transformers for long contexts

Real-world impact: - First viable Transformer alternative since 2017 - Enables truly long sequences (100k+ tokens efficiently) - Active research and production exploration - Striped Hyena, Mamba-2 building on this

The insight: Selective state spaces can match attention's expressiveness while being linear-time.

The Transformer Bottleneck¶

Quadratic Complexity is Limiting¶

Attention is O(n²):

Sequence length: n
Attention matrix: n × n
Memory: O(n²)
Compute: O(n²)

n = 1,000   → 1M operations
n = 10,000  → 100M operations
n = 100,000 → 10B operations

Even with FlashAttention: - Still fundamentally quadratic - Long sequences get expensive - KV cache grows linearly with sequence length

The Dream: Linear Time Attention¶

What we want: - O(n) time complexity - O(1) memory for generation - Same quality as Transformers - Works for language modeling

Previous attempts fell short: - Linear attention (2020): Lower quality - Performers (2021): Approximations, quality loss - S4 (2022): Good but not quite there

State Space Models (SSMs)¶

Background: S4 (2022)¶

Structured State Space Models (S4): - Linear time complexity ✅ - Continuous-time formulation - Good for some tasks - But: Struggled on language modeling ❌

Why S4 wasn't enough:

Key limitation: Time-invariant
- Same dynamics for every input
- Can't be selective about what to remember
- Attention wins because it's content-based

Mamba's Innovation: Selective SSMs¶

Key idea: Make the state space input-dependent

S4 (time-invariant):

Same transition matrix A for all inputs
[Memory is fixed, can't adapt to content]

Mamba (selective):

Different dynamics based on input content
[Can choose what to remember, like attention!]

How Mamba Works¶

1. State Space Formulation¶

Continuous-time system:

h'(t) = Ah(t) + Bx(t)  [State evolution]
y(t)  = Ch(t)          [Output]

Where:
h(t) = hidden state (size N)
x(t) = input
y(t) = output
A, B, C = parameters

Discretized for computers:

h_t = Ā h_{t-1} + B̄ x_t
y_t = C h_t

Where Ā, B̄ are discrete versions of A, B

2. The Selective Mechanism¶

Make A, B, C input-dependent:

# Instead of fixed parameters
A = fixed_matrix

# Make them functions of input
B = Linear(x)  # Input-dependent
C = Linear(x)  # Input-dependent
Δ = Softplus(Linear(x))  # Time-dependent discretization

This is the key innovation!

Now the model can:
- Decide what information to keep (via B)
- Decide what to output (via C)
- Decide how fast to update (via Δ)

3. Efficient Computation¶

Naive approach would be slow:

# Sequential (slow, but flexible)
h = 0
for t in range(sequence_length):
    h = A @ h + B[t] @ x[t]  # Can't parallelize!
    y[t] = C[t] @ h

Mamba's solution: Hardware-aware scan - Fused CUDA kernels (like FlashAttention) - Parallel prefix scan for training - Sequential scan for generation - Result: Fast on GPUs!

Architecture¶

Mamba Block¶

class MambaBlock(nn.Module):
    def __init__(self, d_model, d_state=16):
        self.d_model = d_model  # Model dimension
        self.d_state = d_state  # SSM state dimension

        # Linear projections
        self.in_proj = nn.Linear(d_model, d_model * 2)

        # SSM parameters (input-dependent)
        self.x_proj = nn.Linear(d_model, d_state * 2)
        self.dt_proj = nn.Linear(d_model, d_model)

        # Output
        self.out_proj = nn.Linear(d_model, d_model)

        # State space parameters
        self.A_log = nn.Parameter(torch.randn(d_model, d_state))
        self.D = nn.Parameter(torch.randn(d_model))

    def forward(self, x):
        """
        x: (batch, seq_len, d_model)
        """
        batch, seq_len, dim = x.shape

        # Split into two paths (like Transformer FFN)
        x_and_res = self.in_proj(x)
        x, res = x_and_res.split([dim, dim], dim=-1)

        # Apply nonlinearity
        x = F.silu(x)

        # Selective SSM
        y = self.selective_scan(x)

        # Gating (like in Transformers)
        y = y * F.silu(res)

        # Output projection
        return self.out_proj(y)

    def selective_scan(self, x):
        """The core selective SSM."""
        batch, seq_len, dim = x.shape

        # Compute input-dependent parameters
        delta = F.softplus(self.dt_proj(x))  # (B, L, D)

        B_C = self.x_proj(x)  # (B, L, 2*N)
        B, C = B_C.split([self.d_state, self.d_state], dim=-1)

        # Discretize A (state transition)
        A = -torch.exp(self.A_log)  # (D, N)

        # Efficient selective scan (fused CUDA kernel)
        y = selective_scan_fn(x, delta, A, B, C, self.D)

        return y

Full Mamba Model¶

class Mamba(nn.Module):
    def __init__(
        self,
        vocab_size=50257,
        d_model=768,
        n_layers=12,
        d_state=16
    ):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, d_model)

        self.layers = nn.ModuleList([
            nn.ModuleDict({
                'mamba': MambaBlock(d_model, d_state),
                'norm': RMSNorm(d_model)
            })
            for _ in range(n_layers)
        ])

        self.norm_f = RMSNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)

    def forward(self, input_ids):
        x = self.embedding(input_ids)

        for layer in self.layers:
            # Pre-norm like GPT
            x = layer['mamba'](layer['norm'](x)) + x

        x = self.norm_f(x)
        logits = self.lm_head(x)

        return logits

Comparison: Mamba vs Transformer¶

Complexity¶

Aspect	Transformer	Mamba
Time (forward)	O(n²d)	O(nd)
Memory (training)	O(n²+nd)	O(nd)
KV cache (inference)	O(nd) per token	O(d) constant
Sequence length scaling	Quadratic	Linear

Perplexity (Language Modeling)¶

Pile benchmark (5B token dataset):

Model	Size	Perplexity
Transformer	125M	13.2
Mamba	125M	13.1
Transformer	350M	10.8
Mamba	350M	10.6
Transformer	1.3B	8.7
Mamba	1.3B	8.7

Mamba matches or beats Transformers!

Inference Speed¶

Tokens/second at different sequence lengths:

Sequence Length	Transformer	Mamba	Speedup
512	1000	1100	1.1×
2048	800	1200	1.5×
8192	400	1400	3.5×
32768	100	1500	15×

Massive speedup at long contexts!

Results¶

Language Modeling¶

The Pile (validation perplexity): - Mamba 1.4B: 8.69 - Transformer 1.4B: 8.71 - Statistically tied!

Scaling: Mamba follows same scaling laws as Transformers

DNA Modeling¶

Task: Predict next nucleotide in long DNA sequences

Model	Sequence Length	Accuracy
HyenaDNA	1M	85.2%
Mamba	1M	86.1%

First model to handle 1M token sequences effectively!

Audio Modeling¶

SC09 (speech commands): - Mamba: 91.2% accuracy - S4: 90.4% - Better than predecessor S4

Downstream Tasks¶

Zero-shot evaluation (after pretraining):

Task	Transformer	Mamba
LAMBADA	64.2%	63.1%
HellaSwag	71.0%	70.9%
PIQA	78.1%	77.8%
ARC-Easy	68.8%	68.4%

Nearly identical performance

Practical Usage¶

Using Mamba¶

from mamba_ssm import Mamba

# Create model
model = Mamba(
    d_model=768,        # Model dimension
    d_state=16,         # SSM state expansion factor
    d_conv=4,           # Local convolution width
    expand=2,           # Block expansion factor
)

# Input: (batch, length, d_model)
x = torch.randn(1, 1024, 768)

# Forward pass
y = model(x)  # (1, 1024, 768)

# Fast autoregressive generation
logits = model.generate(
    input_ids,
    max_length=100,
    temperature=0.7
)

Training Mamba¶

import torch
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel

# Initialize model
model = MambaLMHeadModel(
    d_model=1024,
    n_layer=24,
    vocab_size=50257,
)

# Training loop (same as Transformer)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

for batch in dataloader:
    input_ids = batch['input_ids']

    # Forward
    logits = model(input_ids)

    # Loss
    loss = F.cross_entropy(
        logits[:, :-1].reshape(-1, vocab_size),
        input_ids[:, 1:].reshape(-1)
    )

    # Backward
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

Why Selective SSMs Work¶

Content-Based Reasoning¶

Transformer attention:

Attention lets model focus on relevant tokens
Can ignore irrelevant information
Content-based selection

Previous SSMs (S4):

Fixed dynamics
Must process everything equally
Time-invariant = limited expressiveness

Mamba (selective SSM):

Input-dependent parameters
Can filter irrelevant information
Content-based like attention!

Example: Copying Task¶

Task: Copy important tokens, ignore filler

Input:

Copy this: "hello" [filler] [filler] ... [filler]

S4 behavior:

Equal weighting to all tokens
"hello" gets diluted by filler
Fails to maintain important info

Mamba behavior:

B, C, Δ adapt to input
Large Δ for "hello" → remember longer
Small Δ for filler → forget quickly
Successfully copies "hello"!

Advantages Over Transformers¶

1. Linear Scaling¶

# Transformer
time(n) = c * n²  # Quadratic
memory(n) = c * n²

# Mamba
time(n) = c * n  # Linear
memory(n) = c * n

Impact: Can handle much longer sequences

2. Constant Memory Inference¶

# Transformer: Must store KV cache
cache_size = n * d  # Grows with sequence length

# Mamba: Fixed-size state
state_size = d_state  # Constant!

Impact: 5-10× less memory for long generation

3. Faster Long Sequence Processing¶

At n=32k tokens:
Transformer: Limited by attention
Mamba: Still fast (linear time)

4. Better for Streaming¶

# Mamba can process token-by-token efficiently
state = init_state()
for token in stream:
    output, state = mamba.step(token, state)
    # Constant time per token!

Limitations¶

1. Recall Performance¶

Transformers slightly better at: - Exact copying of earlier tokens - Random access to specific earlier info - Some associative recall tasks

Why: - Attention has explicit access to all tokens - Mamba must compress into fixed-size state

2. Implementation Maturity¶

Custom CUDA kernels required
Less tooling than Transformers
Fewer pretrained models available

3. Architecture Search¶

Optimal hyperparameters still being found
d_state, d_conv, etc. less understood
Transformers have 7 years of tuning

4. Hybrid Models May Be Better¶

# Some evidence that mixing is best
model = [
    MambaLayer,  # Cheap, long-range
    MambaLayer,
    AttentionLayer,  # Expensive, precise
    MambaLayer,
    MambaLayer,
    AttentionLayer,
]

Modern Developments¶

1. Mamba-2 (2024)¶

Improvements: - Faster: 2-8× over Mamba-1 - Better quality - Theoretical connections to attention

2. Hybrid Models¶

Jamba (AI21 Labs):

Mix of Mamba + Attention layers
- Mamba for efficiency
- Attention for recall
- Best of both worlds?

Striped Hyena:

Hyena operators (similar to Mamba)
Competitive with Transformers

3. Mamba in Production¶

Together.ai: - Offering Mamba inference - Faster than Transformers for long contexts

Cartesia: - Building Mamba-based models - Focusing on real-time applications

When to Use Mamba¶

✅ Good Fit¶

1. Long Sequences - DNA/protein modeling (1M+ tokens) - Long document processing - Time series (multivariate, long horizon)

2. Real-Time Streaming - Audio processing - Video analysis - Continuous monitoring

3. Resource-Constrained - Limited GPU memory - Long inference sequences - Need low latency

4. Research/Exploration - Alternative to Transformers - Architecture research - New application domains

❌ Less Ideal¶

1. Short Sequences - Transformer overhead is manageable - Mamba's benefits minimal - Ecosystem more mature for Transformers

2. Tasks Requiring Precise Recall - Question answering with exact quotes - Copying specific earlier content - Some reasoning tasks

3. Production Critical Paths - Less battle-tested than Transformers - Fewer pretrained models - Risk tolerance needed

Implementation Tips¶

1. Use Official Implementation¶

# pip install mamba-ssm
from mamba_ssm import Mamba

# Requires CUDA for speed

2. Tune State Dimension¶

# Trade-off: capacity vs speed
d_state = 16   # Default, fast
d_state = 32   # More capacity
d_state = 64   # Expensive

3. Consider Hybrid Architectures¶

# Example: Mamba + occasional attention
layers = [
    MambaBlock(...),
    MambaBlock(...),
    MambaBlock(...),
    AttentionBlock(...),  # Every 4th layer
    MambaBlock(...),
    # ...
]

4. Profile on Your Data¶

# Compare on your specific workload
time_mamba = benchmark(mamba_model, data)
time_transformer = benchmark(transformer_model, data)

if time_mamba < time_transformer:
    use_mamba()

Key Takeaways¶

First real Transformer alternative - Matches quality with better scaling
Linear time complexity - O(n) vs O(n²)
Selective SSMs - Content-based reasoning like attention
Great for long sequences - 100k+ tokens feasible
Still maturing - Transformers still dominant for now

Bottom line: Mamba is the most promising Transformer alternative to date, offering linear-time sequence modeling that matches Transformer quality for the first time.