State Space Models: Mamba and the Transformer Alternative

Transformers have one fatal flaw.

That self-attention mechanism from Episode 2? It scales quadratically with sequence length. Double the context, quadruple the compute. This is fine at 4K tokens. Painful at 128K. And essentially impossible at 1 million tokens without tricks.

What if there was an architecture with linear scaling? Double the context, double the compute. No attention mechanism at all. No KV cache. No quadratic bottleneck.

That's Mamba — the state space model that's challenging the transformer's throne. It processes 1 million tokens with the same per-token cost as processing 1,000. And it's not a toy — it matches transformer quality on language tasks while being dramatically faster at inference.

The Quadratic Problem

Let's make the transformer's scaling problem concrete:

Sequence Length Attention Compute Relative Cost

Self-attention computes a score between every pair of tokens. With N tokens, that's N² pairs. This is why long-context models are so expensive — and why companies charge a premium for 128K+ context.

State space models eliminate this entirely.

What Is a State Space Model?

A state space model (SSM) comes from control theory — the math behind autopilots, robot controllers, and signal processing systems. The idea: maintain a hidden state that gets updated as each new input arrives.

The Continuous Form

In continuous time (control theory), a state space model is:

dx/dt = Ax + Bu (state equation)

y = Cx + Du        (output equation)

where:

x = hidden state vector
u = input (one element at a time)
y = output
A = state transition matrix (how state evolves)
B = input projection matrix (how input affects state)
C = output projection matrix (how state becomes output)
D = skip connection (direct input-to-output)

The Discrete Form

For processing sequences (text, audio), we discretize to get:

x_k = Ā·x_{k-1} + B̄·u_k    (update state with new input)
y_k = C·x_k                   (read output from state)

where:

k = timestep (each token)

Ā, B̄ = discretized versions of A, B

Think of it like this: the model maintains a summary vector (the state) that gets updated with each new token. To produce an output, it reads from this summary. It never looks back at previous tokens — it only looks at its current state.

Token 1: "The"    → state_1 = update(state_0, "The")
Token 2: "cat"    → state_2 = update(state_1, "cat")
Token 3: "sat"    → state_3 = update(state_2, "sat")
Token 4: "on"     → state_4 = update(state_3, "on")

To predict token 5: read(state_4) → "the"

Why Linear Scaling?

Each token update is a constant-time operation — one matrix multiply regardless of how many tokens came before. Process 1,000 tokens? 1,000 matrix multiplies. Process 1,000,000 tokens? 1,000,000 matrix multiplies. Linear, not quadratic.

Architecture Per-Token Compute Memory for Context

🔧

S4: The Breakthrough

S4 (Structured State Spaces for Sequence Modeling, 2021) was the paper that proved SSMs could compete with transformers on real tasks.

The key innovation: how to parameterize the A matrix. A naive state matrix loses long-range information — early tokens get "forgotten" as the state evolves. S4 uses the HiPPO (High-Order Polynomial Projection Operator) initialization:

A = HiPPO matrix (specially structured)

This matrix is designed to continuously approximate the history

of inputs as polynomial coefficients. It "remembers" optimally.

S4 also showed that SSMs can be computed as a convolution during training (parallel, fast on GPUs) and as a recurrence during inference (sequential, memory-efficient).

Training: Convolution mode — process all tokens in parallel

Inference: Recurrence mode — process one token at a time with fixed state

This dual computation mode is a huge advantage:

Training benefits from GPU parallelism (like transformers)

Inference benefits from constant memory (unlike transformers)

📊

Mamba: The Selective SSM

Mamba (December 2023, Albert Gu and Tri Dao) is the SSM that actually threatens transformers. Its key innovation: selective state spaces — making the SSM parameters input-dependent.

The Problem with Basic SSMs

In a standard SSM, A, B, and C are fixed matrices — the same for every input token. This means the model processes "the" and "photosynthesis" with the exact same dynamics. It can't selectively focus on important tokens or ignore filler.

Mamba's Fix: Selection Mechanism

Mamba makes B, C, and Δ (the discretization step) functions of the input:

python

Standard SSM (fixed parameters)

A = fixed_matrix
B = fixed_matrix
C = fixed_matrix
Mamba (input-dependent parameters)

B = linear(input)    # Changes based on what the token is
C = linear(input)    # Changes based on what the token is
Δ = softplus(linear(input))  # Controls how much to update state

The Δ (delta) parameter is especially important — it controls the "gate" for each token:

Large Δ → focus on the current input, update state significantly

Small Δ → ignore the current input, keep existing state

This is conceptually similar to the gating mechanism in LSTMs, but implemented through the state space formalism.

The Selection Effect

Input: "The weather in Mumbai is usually hot and humid during summer months"

Δ values (learned):

"The" → small Δ (filler word, don't update much)

"weather" → medium Δ (somewhat important)

"Mumbai" → LARGE Δ (critical — location matters!)

"is" → small Δ (filler)

"usually" → small Δ (filler)

"hot" → LARGE Δ (key fact!)

"and" → small Δ (filler)

"humid" → LARGE Δ (key fact!)

...

The model learns to selectively compress the input sequence — keeping important information and discarding filler. This is why Mamba matches transformer quality despite never doing explicit attention.

Mamba Architecture

A full Mamba model stacks Mamba blocks (replacing transformer blocks):

Input tokens

↓

Token Embedding

↓

┌─ Mamba Block 1 ──────────────────────┐

│ Linear → Conv1D → SSM → Linear │

│ + Skip connection + LayerNorm │

└──────────────────────────────────────┘

↓

┌─ Mamba Block 2 ──────────────────────┐

│ (same structure) │

└──────────────────────────────────────┘

↓

... (N blocks)

↓

Output head → next token prediction

Inside a Mamba Block

python

def mamba_block(x):

# Two parallel branches
    # Branch 1: Selective SSM path
    z = linear_proj(x)              # Expand dimension
    z = silu(conv1d(z))             # Local context (kernel=4)
    z = selective_ssm(z)            # The magic — selective state space

# Branch 2: Gate
    gate = silu(linear_proj(x))

# Combine
    output = z * gate               # Element-wise gating
    output = linear_proj(output)    # Project back to model dim
    return output + x               # Residual connection

The Conv1D with kernel size 4 gives local context (similar to looking at neighboring tokens), while the SSM provides global context through its recurrent state.

Mamba vs Transformers: Head to Head

Feature Transformer Mamba

Attention O(N²) quadratic None

Where Mamba Wins

Long sequences — 1M+ tokens without the KV cache problem

Inference speed — constant per-token cost regardless of context

Memory efficiency — fixed state instead of growing KV cache

Audio/signal processing — naturally good at continuous signals

Where Transformers Win

In-context retrieval — "What was the 3rd word in paragraph 7?" requires looking back

Copying — reproducing exact sequences from context

Established ecosystem — years of optimization, tooling, and scaling research

🔬

Mamba-2 and Hybrid Models

Mamba-2 (2024) bridged the gap further by showing that selective SSMs can be expressed as a form of structured attention — unifying the two paradigms mathematically.

The most promising direction: hybrid models that use both:

Layer 1: Mamba block (efficient long-range)

Layer 2: Mamba block

Layer 3: Attention block (precise retrieval)

Layer 4: Mamba block

Layer 5: Mamba block

Layer 6: Attention block

...

Jamba (AI21 Labs) was the first production hybrid — 52B params with a mix of Mamba and attention layers. It gets the best of both worlds: Mamba's efficiency for most processing, attention's precision for retrieval tasks.

Model Architecture Context Notes

Mamba-1 (2.8B) Pure Mamba Unlimited Research model

🛡️

The Hardware Angle

Mamba requires a custom CUDA kernel called selective scan — and this is both its strength and weakness:

Strength: The selective scan kernel (written by Tri Dao, the FlashAttention author) is highly optimized and memory-efficient.

Weakness: Transformers have years of hardware optimization. Every GPU vendor optimizes for matrix multiplications and attention. Mamba's selective scan is newer and less optimized.

Transformer training: cuBLAS (GEMM) + FlashAttention → highly optimized

Mamba training: cuBLAS (GEMM) + Selective Scan → newer, less optimized

The gap is closing, but transformers have a hardware ecosystem advantage.

📦

Will SSMs Replace Transformers?

Honest answer: probably not entirely, but they'll coexist.

The most likely future:

Pure Mamba for streaming/continuous data (audio, sensor data, very long documents)

Hybrid Mamba+Attention for general language models (best of both worlds)

Pure Transformers for tasks requiring precise retrieval and in-context learning

The transformer isn't going away — it's too good at too many things. But SSMs have shown that the quadratic attention bottleneck isn't the only path to intelligence. And for the long-context future (million-token windows, streaming audio, always-on agents), SSMs are essential.

🚀

Practical Takeaways

SSMs process sequences with constant per-token cost — linear scaling instead of quadratic

Mamba's key innovation is selection — input-dependent parameters let it focus on important tokens

Dual computation modes — convolution for parallel training, recurrence for efficient inference

No KV cache needed — fixed-size state replaces the ever-growing cache of transformers

Hybrid models are the future — combining Mamba efficiency with attention precision

Transformers aren't dead — they're better at retrieval and in-context learning, so the future is likely a mix

🔄

What's Next?

Episode 62: Test-time Compute — Why do "thinking" models like o1 and DeepSeek-R1 take longer to answer? What are thinking tokens? How does spending more compute during inference improve reasoning? The answer changes how you think about AI costs.

← Previous

Ep 60: Diffusion Models

Next →

Ep 62: Test-Time Compute

Next: Episode 62 — Test-Time Compute

Ask GPT-4 a tricky math problem. It answers in 2 seconds. Wrong. Ask o1 the same thing. It thinks for 30 seconds. Right.

This is part of a 98-episode series covering AI engineering from tokens to production deployment.