RugvAI Labs

State Space Models & Alt-Architectures

Attention has an O(N²) ceiling. At 1M tokens, just the attention matrix is the entire compute budget — which is why every frontier transformer caps out somewhere between 200K and 2M context. Mamba, RWKV, and Jamba bet on a completely different architecture: compress the past into a fixed-size state, never quadratic. Albert Gu's 2023 Mamba paper was the first SSM to match transformer quality at language modeling, and 2024–2026 hybrid models (Jamba, Zamba, Samba) are now serving million-token contexts at fractions of transformer cost.

Learning Objectives

After this lesson, you will be able to:

Explain why attention's O(N²) cost is the architectural ceiling that motivates linear-time alternatives — and what 'linear time' actually buys you in training and inference
Walk through the discretization that turns a continuous state space x'(t) = Ax + Bu into the Mamba selective scan you can actually run on a GPU
Pick the right architecture for a workload — pure transformer for retrieval-heavy tasks, pure SSM for streaming low-memory inference, or hybrid (Jamba/Zamba) for long-context cost-sensitive deployments
Recognize the strengths and weaknesses of Mamba, RWKV, RetNet, and Hyena well enough to read a 2025 paper without getting lost in the alphabet soup

Prerequisites

Before You Start

Why

Why This Matters

Why Albert Gu's Mamba paper became the most-talked-about ICML 2024 alternative to attention -- when "Mamba: Linear-Time Sequence Modeling with Selective State Spaces" landed in late 2023 it was the first SSM that matched transformer language-modeling quality at small-to-medium scale while running with O(N) inference and constant memory; researchers spent 2024 stress-testing whether the result generalizes
How AI21's Jamba and Zyphra's Zamba use Mamba+attention hybrids to ship million-token context cheaper than pure transformers -- Jamba (March 2024) interleaves Mamba layers with a few attention layers to keep recall ability while paying mostly linear cost; Zamba does similar; both run inference at fractions of the GPU memory of comparable transformers
Why every long-context inference benchmark in 2025 has at least one SSM in the top 3 -- when you need to serve 200K+ token contexts cheaply, the math forces you toward sub-quadratic architectures; the question stopped being "if" and became "how much hybridization with attention is needed"
Build this --> Implement a 50-line selective SSM scan in PyTorch, run it on a 100K-token sequence, and watch GPU memory stay flat while a same-parameter attention model OOMs at 32K

Don't worry if "state space" sounds like control theory — that is exactly where it came from, and the connection back to Kalman filters and signal processing is part of why these models work. The math looks daunting but the core idea is just "an RNN you can train in parallel."

#Why We Even Need Alternatives

Intuition

Build Your Intuition

Attention is content-based and exact: every token sees every other token. That is its superpower for retrieval and in-context learning, and its curse for long sequences. Doubling the context quadruples the FLOPs. Doubling again is 16x. By 1M tokens, attention alone is the entire compute budget.

A state space model takes a different bet: compress the past into a fixed-size state. The state is updated each step based on the new input. To predict the next token you only need the current state and the new input — never the full history. That is O(1) per token at inference and constant memory. The price you pay is that the state has to forget — it cannot store every previous token verbatim.

The key insight from Mamba (2023) is that the forgetting policy itself can be learned and made input-dependent. If the input is "the password is 12345", the state should remember that. If the input is the 47th boilerplate sentence in a Reddit comment, the state should forget aggressively. Earlier SSMs had a fixed forgetting rate; Mamba made it selective.

Anecdote

Historical Note

State space models go back to Rudolf Kálmán's 1960 filter — the same math that landed Apollo on the moon was tracking the lunar module's position from noisy radar with x'(t) = Ax(t) + Bu(t). Decades of control theory and signal processing later, Albert Gu's S4 paper at ICLR 2022 ("Efficiently Modeling Long Sequences with Structured State Spaces") showed that a carefully initialized linear SSM could compete with transformers on long-range arena benchmarks. Mamba (Gu & Dao, ICML 2024) added input-dependent selectivity and matched transformer LM quality. Mamba-2 (2024) unified attention and SSM under "state space duality." RWKV (Peng, EMNLP 2023) independently arrived at parallel-trainable RNNs from the LLM-from-scratch direction. RetNet (Sun, Microsoft 2023) added a third path. By 2024–2025, hybrid Transformer-Mamba models (Jamba, Zamba) became the practical default for long-context, cost-sensitive deployments.

#The Mathematical Setup

A continuous-time state space model is two linear ODEs — one for state, one for output:

x'(t) = A\,x(t) + B\,u(t) \\[4pt] y(t) = C\,x(t) + D\,u(t)

To run on a GPU you discretize — replace the derivative with a finite step of size Δ:

x_t = \bar{A}\,x_{t-1} + \bar{B}\,u_t \\[4pt] y_t = C\,x_t \\[6pt] \bar{A} = \exp(\Delta A), \quad \bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I)\,\Delta B

The convolutional view rewrites the recurrence as a global convolution with a structured kernel:

y = u \ast K, \quad K = (C\bar{B},\; C\bar{A}\bar{B},\; C\bar{A}^2\bar{B},\; \ldots,\; C\bar{A}^{N-1}\bar{B})

The trick that made S4 work — and that Mamba inherited — is the HIPPO initialization of A. HIPPO comes from approximation theory: it is the matrix that produces the best polynomial reconstruction of past inputs. Random initialization fails badly; HIPPO-initialized SSMs train cleanly and capture long-range dependencies up to 16K tokens on the long-range arena benchmark.

#Mamba: Selectivity Changes Everything

S4 was elegant but linear and time-invariant — the matrices A, B, C are fixed, not data-dependent. So an S4 state forgets at the same rate whether the input is "remember this password" or "irrelevant filler". Mamba's contribution: make B, C, and Δ a function of the current input, breaking time-invariance.

\bar{A}_t,\; \bar{B}_t,\; C_t,\; \Delta_t \;=\; f_\theta(u_t)

The cost is that the convolutional view no longer applies — the kernel changes every step. Mamba solves this with a hardware-aware parallel scan (Blelloch scan) implemented as a custom CUDA kernel that does the recurrence in O(log N) parallel passes per chunk, using shared memory aggressively. The result: training as fast as a comparable transformer, inference O(1) per token with constant memory.

Intuition

Build Your Intuition

Compare two long-context strategies for a 1M-token document:

Transformer: every new token attends to all 1M previous tokens. Each step is O(N) — for every token in the sequence. Total inference: O(N²). KV cache grows linearly: 1M tokens × layers × dim ≈ 100GB+ of VRAM at production scale.

Mamba: every new token updates a fixed-size state of (say) 16 channels × 64 dim = 1024 floats. Each step is O(1). Total inference: O(N). State memory: ~4KB per layer regardless of context length. The cost: information older than ~10K tokens is mostly forgotten, recoverable only through what survived in the state.

Hybrid models (Jamba/Zamba) sprinkle a few attention layers among many Mamba layers — the attention layers provide exact recall when needed; the Mamba layers carry the long-range gist cheaply.

Step a token through both pipelines to contrast SSM's fixed-size rolling state with attention's growing all-to-all lookup.

Loading visualization...

#The Family Tree of Alt-Architectures

Model	Core idea	Training	Inference	Strength	Weakness
S4 (2022)	HIPPO-init linear SSM	O(N log N) conv	O(N) recur	Long-range arena	LM quality below transformers
Mamba (2023)	Input-dependent SSM (selective scan)	O(N) parallel scan	O(1) per token	LM-quality competitive at small/mid scale	Recall/copying tasks
Mamba-2 (2024)	State-space duality with attention	Same	Same	Better quality, simpler kernel	Newer, less battle-tested
RWKV (2023)	RNN with linear-attention parallel training

Hybrid 2024–2025 systems typically use ~1 attention layer per 7–8 Mamba layers — enough to handle in-context learning and copy/recall tasks, while keeping most of the cost on the linear path.

What Do You Think?

You have to serve 10M-token contexts (e.g. lifetime medical history) on commodity hardware. Pure transformer is impossible — KV cache alone exceeds VRAM. Which architecture earns the slot?

Pure Mamba sounds right — and it works for many tasks — but production teams overwhelmingly pick the hybrid because real workloads include in-context learning, copying, and few-shot patterns where pure SSMs underperform. The hybrid pays a small cost relative to pure Mamba but recovers most of attention's recall ability. Jamba and Zamba both ship this design. Sliding-window transformer with 4K is a non-answer for 10M tokens; Linformer's fixed-length projection loses too much information at extreme lengths.

#Building a Tiny Selective SSM

pythonplayground.py · Pyodide

# Minimal selective SSM in PyTorch — Mamba's selective scan, simplified
import torch
import torch.nn as nn
import torch.nn.functional as F

class TinySelectiveSSM(nn.Module):
  """A 50-line selective state space model.
  
  Real Mamba uses a hand-tuned CUDA kernel for the parallel scan.
  This naive version is O(N) sequential — fine for understanding,
  too slow for production training. Run with batch=1 for the demo.
  """
  def __init__(self, d_model=64, d_state=16):
      super().__init__()
      self.d_state = d_state
      # The "structured" A matrix: diagonal with negative real parts (for stability)
      # This is the HIPPO-style decay; in real Mamba it is learnable per-channel
      log_A = torch.log(torch.linspace(0.1, 1.0, d_state))
      self.log_A = nn.Parameter(log_A.expand(d_model, d_state).clone())
      # Input-dependent projections that produce B_t, C_t, Delta_t per step
      self.proj_BCdt = nn.Linear(d_model, 2 * d_state + 1)
      # In-projection and out-projection (real Mamba has more structure)
      self.in_proj = nn.Linear(d_model, d_model)
      self.out_proj = nn.Linear(d_model, d_model)
  
  def forward(self, u):
      # u: (batch, seq_len, d_model)
      B, T, D = u.shape
      u = self.in_proj(u)
      # Compute per-step B_t, C_t, delta_t
      bcd = self.proj_BCdt(u)  # (B, T, 2*d_state + 1)
      B_t, C_t, dt = bcd.split([self.d_state, self.d_state, 1], dim=-1)
      dt = F.softplus(dt)  # delta must be positive
      # The state. (B, D, d_state) — one state per channel (per-channel SSM)
      x = torch.zeros(B, D, self.d_state, device=u.device)
      ys = []
      # A is fixed in this toy version; in real Mamba it is also gated
      A = -torch.exp(self.log_A)  # (D, d_state), negative for stable decay
      for t in range(T):
          # Discretize: A_bar = exp(dt * A), B_bar = dt * B
          A_bar = torch.exp(dt[:, t] * A.unsqueeze(0))  # (B, D, d_state)
          B_bar = dt[:, t].unsqueeze(-1) * B_t[:, t].unsqueeze(1)  # (B, 1, d_state) -> (B, D, d_state)
          # Update state
          x = A_bar * x + B_bar * u[:, t].unsqueeze(-1)  # (B, D, d_state)
          # Read out via C_t, summed over state dim
          y = (x * C_t[:, t].unsqueeze(1)).sum(dim=-1)  # (B, D)
          ys.append(y)
      y = torch.stack(ys, dim=1)
      return self.out_proj(y)

# Demo: process a 4096-length sequence, watch memory stay flat
torch.manual_seed(0)
model = TinySelectiveSSM(d_model=64, d_state=16)
u = torch.randn(1, 4096, 64)

print(f"Input shape:    {tuple(u.shape)}")
print(f"Parameters:     {sum(p.numel() for p in model.parameters()):,}")

with torch.no_grad():
  y = model(u)
print(f"Output shape:   {tuple(y.shape)}")
print(f"State size:     {16} per channel (constant, regardless of T)")
print()
print("Compare to attention on the same sequence:")
print("  - Attention: O(T^2) = 4096^2 = 16.7M ops + O(T^2) memory")
print("  - This SSM:  O(T)   = 4096   ops + O(d_state) memory")
print()
print("CHALLENGE: Replace the python for-loop with a parallel scan")
print("(torch.cumsum-based associative scan). Real Mamba uses a custom")
print("CUDA kernel; an associative scan via prefix sums is the textbook approach.")
print("Hint: log-cumsum-exp gives you the running A_bar product.")

Tests · The model should produce output shape (1, 4096, 64) and use roughly the same VRAM as a 64-dim shape regardless of sequence length, while a same-d_model attention model would balloon.

#Key Takeaways

State space models are linear-time alternatives to attention. They compress the past into a fixed-size state and update it each step. Inference is O(1) per token with constant memory; training is O(N) parallelizable via convolution or scan.
HIPPO initialization is what made S4 work. Random init fails badly. The structured initialization is non-negotiable; modern selective SSMs (Mamba) inherit HIPPO-style structure even after making A partly learnable.
Selectivity is what made Mamba match transformers. Input-dependent B_t, C_t, Δ_t let the state forget aggressively when the input is filler and remember strongly when it is informative. This breaks time-invariance and requires a hardware-aware parallel scan.
Hybrids beat purity in production. Pure SSMs underperform on copying, recall, and in-context learning. Jamba and Zamba intersperse a few attention layers among many Mamba layers; this hybrid is the practical default for cost-sensitive long-context deployment.
The architecture race is not over. Transformers are not the final answer; they are the current local optimum that combines content-based retrieval with hardware-friendly matmul. SSMs, RWKV, RetNet, and Hyena are all credible alt-architectures with active research traction.

#Quick Check

Quick Check1 / 4

What is the asymptotic inference cost per token for a selective SSM (Mamba) vs a vanilla transformer (without KV cache)?

State space models are the most credible challenger transformers have faced. Whether they win, lose, or merge into hybrids by 2030, an ML engineer in 2026 needs to recognize them on sight. Next up: how transformers got into vision in the first place — Vision Transformers, where image patches became the new tokens.