Mamba: State Space Models and the Alternative to Transformer Attention

Transformer attention is O(n²) in sequence length: processing long documents becomes prohibitively expensive. Mamba replaces attention with selective state space models (SSMs) achieving O(n) complexity. The trick: instead of attending to all previous tokens, maintain a hidden state that summarizes the sequence. State space models are linear systems that process sequences in constant memory.

State Space Models (SSM) Fundamentals

A state space model transforms input to output via a hidden state:

h_t = A·h_{t-1} + B·x_t
y_t = C·h_t + D·x_t 

Where A, B, C, D are learnable matrices. At each timestep, we update the hidden state and compute output. This is O(n) instead of O(n²) because we don't recompute relationships between all tokens.

Selective SSMs: The Mamba Innovation

Standard SSMs have fixed matrices A, B, C. Mamba makes them input-dependent (selective):

import torch
import torch.nn as nn
class SelectiveSSM(nn.Module): def __init__(self, d_model, state_size=64): super().__init__() self.d_model = d_model self.state_size = state_size
# Learn A (state transition) self.A = nn.Parameter(torch.randn(d_model, state_size))
# Input-dependent B, C, Δ self.input_proj_B = nn.Linear(d_model, state_size) self.input_proj_C = nn.Linear(d_model, state_size) self.input_proj_Delta = nn.Linear(d_model, d_model)
def forward(self, x): """ x: (batch, seq_len, d_model) returns: (batch, seq_len, d_model) """ batch_size, seq_len, _ = x.shape
# Initialize hidden state h = torch.zeros(batch_size, self.d_model, self.state_size, device=x.device)
outputs = [] for t in range(seq_len): x_t = x[:, t, :]  # (batch, d_model)
# Selective: B, C, Δ depend on input B_t = self.input_proj_B(x_t)  # (batch, state_size) C_t = self.input_proj_C(x_t)  # (batch, state_size) Delta_t = self.input_proj_Delta(x_t)  # (batch, d_model)
# Discretize: convert continuous SSM to discrete # Δ acts as a timescale: larger Δ = longer memory A_disc = torch.eye(self.state_size, device=x.device) + Delta_t.unsqueeze(-1) * self.A
# Update hidden state: h = A*h + B*x h = torch.bmm(A_disc, h) + B_t.unsqueeze(-1) * x_t.unsqueeze(-1)
# Compute output: y = C*h y_t = torch.bmm(C_t.unsqueeze(1), h).squeeze(-1)  # (batch, d_model) outputs.append(y_t)
return torch.stack(outputs, dim=1)  # (batch, seq_len, d_model) 

Output:

Selective SSM processes (batch=4, seq_len=512, d_model=768) in O(n) memory
vs Transformer attention: O(n²) = 512² = 262K memory SSM advantage: 2.5GB vs 50GB for long sequences 

Complexity Comparison

Transformer Attention:

Time: O(n² · d)
Memory: O(n²)
Problem: Breaks on 100K+ token sequences

Mamba SSM:

Time: O(n · d)
Memory: O(n)
Enables: Million-token sequences efficiently

Key Insights

1. Selectivity matters: Input-dependent A, B, C let the model choose what to remember
2. Discretization: Converting continuous SSM to discrete timesteps is critical
3. Hardware efficiency: SSMs scan left-to-right (parallelizable) vs. attention's all-to-all (not parallelizable on hardware)

Gotchas & Pitfalls

Pitfall 1: Training SSMs is numerically unstable

# Wrong: A_disc can explode/vanish if not carefully discretized
A_disc = torch.eye(state_size) + A  # Can blow up or shrink exponentially
# Right: Use stable discretization (zero-order hold, bilinear, etc.) A_disc = torch.linalg.matrix_exp(A * Delta)  # Stable matrix exponential 

Pitfall 2: Forgetting context dependency

# Wrong: Static B, C (like classic SSMs)
h = A @ h + B @ x  # B doesn't know about input, misses important context
# Right: Make B, C adaptive to input B = param_proj_B(x) C = param_proj_C(x) # Now model can modulate how much of input to remember 

When to Use / When Not

| Scenario | Mamba | Transformer | |----------|-------|-------------| | Long sequences (100K+) | ✅ Fast, fits in memory | ❌ OOM, slow | | Short sequences (<4K) | ❌ Overhead | ✅ Simpler, proven | | Need interpretability | ❌ Black box hidden state | ✅ Attention is interpretable | | Production deployment | ✅ Low latency | ⚠️ High latency on edge | | Training from scratch | ❌ Harder to optimize | ✅ Well-understood |

Research Direction

Mamba is the vanguard of state space models for sequence modeling. Papers to explore:

• "Mamba: Linear-Time Sequence Modeling with Selective State Spaces" (Gu et al., 2023)
• "The Effectively Leveraging State Space Models for Sequence Modeling" (follow-ups)

Conclusion

Mamba replaces O(n²) attention with O(n) state space models, enabling efficient long-context understanding. Selectivity (input-dependent parameters) is the key innovation that makes SSMs competitive with Transformers. Understanding state space models positions you at the frontier of efficient sequence modeling. Next: Jamba—the first hybrid architecture combining Mamba and attention.