← Back to Table of Contents

Chapter 30 — Beyond Transformers — SSMs & Alternatives

“Attention is all you need — until your sequence is a million tokens long and the O(T²) cost makes you need something else.”

The Quadratic Attention Problem

Standard attention computes a T×T score matrix:

• Compute: O(T² · d)
• Memory: O(T²) for the attention matrix (O(T) with Flash Attention, but compute remains O(T²))

At T = 1M tokens, even with Flash Attention, the FLOPs for attention alone become prohibitive. This motivates architectures with linear or sub-quadratic complexity.

State Space Models (SSMs)

SSMs map an input sequence to an output sequence through a latent state:

\(h_t = \bar{A} h_{t-1} + \bar{B} x_t \quad \text{(state update)}\) \(y_t = C h_t + D x_t \quad \text{(output)}\)

Where:

• $h_t \in \mathbb{R}^N$ — hidden state (per channel)
• $\bar{A}, \bar{B}$ — discretized system matrices
• $C, D$ — output projection matrices

S4: Structured State Spaces (2022)

The breakthrough that made SSMs competitive:

• HiPPO initialization for matrix A — enables remembering long-range dependencies
• Parallel scan — the recurrence can be computed as a convolution during training (parallelizable on GPUs)
• During training: O(T log T) via FFT-based convolution
• During inference: O(1) per step via the recurrence

Mamba: Selective State Spaces (2023)

Mamba (Gu & Dao) made SSMs competitive with transformers by making the state matrices input-dependent (selective):

Key innovation: standard SSMs have fixed A, B, C matrices (same for every input). Mamba makes B, C, and the discretization step Δ input-dependent, allowing the model to selectively focus on or ignore different parts of the input.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# Simplified selective SSM (conceptual)
def selective_ssm(x, A, D, dt_proj, B_proj, C_proj, conv1d):
    """
    x: [B, T, d_inner] — input after expansion
    """
    # Input-dependent parameters (SELECTIVE)
    delta = F.softplus(dt_proj(x))    # [B, T, d_inner] — controls forgetting
    B = B_proj(x)                      # [B, T, N] — input-to-state
    C = C_proj(x)                      # [B, T, N] — state-to-output
    
    # Discretize
    A_bar = torch.exp(delta.unsqueeze(-1) * A)  # [B, T, d_inner, N]
    B_bar = delta.unsqueeze(-1) * B.unsqueeze(2) # [B, T, d_inner, N]
    
    # Parallel selective scan (efficient CUDA kernel)
    y = selective_scan(x, A_bar, B_bar, C, D)  # [B, T, d_inner]
    return y

Mamba-1 3B matches Transformer 3B on language modeling, with 5× higher throughput on long sequences.

Mamba-2 (2024)

Mamba-2 showed that selective SSMs can be viewed as a form of structured linear attention:

• 2–8× faster than Mamba-1 via SSD (Structured State Space Duality) algorithm
• Connections to linear attention make it easier to reason about theoretically
• Still O(Td) but with better hardware utilization

Linear Attention

Replace the softmax attention kernel with a linear kernel:

By computing $\phi(K)^T V$ first (an N×d matrix), we avoid the T×T attention matrix:

• Compute: O(T · d · N) where N is the feature dimension
• Memory: O(T · d) — no attention matrix stored

Variants: RetNet, HGRN2, GLA (Gated Linear Attention).

RWKV: RNN-Transformer Hybrid

RWKV combines the parallelizable training of transformers with the constant-memory inference of RNNs:

• Training: attention-like parallelism using WKV operator
• Inference: recurrent — process one token at a time with fixed state
• No KV-cache — state size is fixed regardless of context length

Hybrid Architectures

The most promising direction may be combining attention and SSMs:

Jamba (AI21, 2024)

Jamba 52B supports 256K context with much less memory than a pure attention model, because most layers don’t have a KV-cache.

Other Hybrids

Complexity Comparison

Current State

As of early 2026:

• Transformers dominate: the ecosystem, tooling, and proven scaling make them the default
• Mamba/SSMs: competitive at small scale, promising at long context, but haven’t dethroned transformers at frontier scale
• Hybrids win on efficiency: Jamba, Samba show that mixing architectures gives the best memory-quality tradeoff
• The jury is still out: no alternative has beaten a well-trained transformer at the same compute budget on standard benchmarks

What’s Next

Architecture innovations help, but there’s another way to improve LLM output quality: spend more compute at inference time. The next chapter covers reasoning models — from chain-of-thought prompting to o1-style test-time compute scaling.

← Previous: Chapter 33 — Mixture of Experts · Next: Chapter 35 — Reasoning Models →

Last updated: April 2026