C.W.K.
Stream
Lesson 11 of 11 · published

State-Space Alternatives: Mamba, RWKV, and What Comes Next

~12 min · mamba, rwkv, ssm

Level 0Token
0 XP0/94 lessons0/10 achievements
0/120 XP to next level120 XP to go0% complete

Not every sequence architecture is a Transformer. State-space models (SSMs) and a few RNN-Transformer hybrids offer O(n) complexity alternatives that are getting close on language quality while being meaningfully cheaper at long context.

Mamba (Gu & Dao, 2023)

Mamba replaces self-attention with selective state-space models: input-dependent state transition matrices that decide what information to keep or forget. O(n) compute and O(1) state, vs O(n²) for attention. Mamba 1 demonstrated 5× throughput and 3× wall-clock speedup at long sequences vs comparable Transformers.

Mamba 2 (2024) reframes SSMs as structured linear attention, allowing the same computation to be expressed as matrix multiplications. Result: 2-8× faster training than Mamba 1 with simpler implementation.

RWKV

An RNN-Transformer hybrid: trains in parallel like a Transformer, runs inference in O(1) memory like an RNN. Through versions V4-V8, it has accumulated tricks: multi-headed matrix states, dynamic LoRA-based token shifting, attention-free mixing layers. Comparable quality on many benchmarks; dramatically lower inference memory.

ArchitectureTraining computeInference memoryLong context behavior
Transformer (dense attention)O(n²)O(n)1× (baseline)
Mamba 1O(n)O(1)5× throughput
Mamba 2O(n)O(1)2-8× faster training than Mamba 1
RWKV (V6+)O(n) train, O(1) inferenceO(1)Comparable quality, smaller memory

Despite their theoretical edges, SSMs have not displaced Transformers for production language tasks as of 2026. Attention's direct any-position connectivity remains powerful for complex reasoning, and the engineering ecosystem (Flash Attention, vLLM, training frameworks) has tuned Transformers to extract every drop of efficiency. Hybrids — alternating attention and SSM layers — are an active research direction (e.g., Jamba, Zamba).

Code

Conceptual SSM (selective state space)·python
# Vastly simplified — real Mamba is O(n) via clever scan algorithms.
def selective_ssm_step(x, h, A, B, C):
    # Input-dependent A, B, C (computed from x)
    h = A * h + B * x        # update state
    y = C * h                # produce output
    return y, h

# Iterating over the sequence:
h = torch.zeros(state_size)
outs = []
for t in range(seq_len):
    A_t, B_t, C_t = make_matrices(x[t])  # input-dependent!
    y, h = selective_ssm_step(x[t], h, A_t, B_t, C_t)
    outs.append(y)
# In Mamba, this scan is computed in parallel with O(n) total compute.

External links

Exercise

Read Mamba 2's introduction. Identify the single most important property the authors claim sets it apart. Then find an open-source Mamba 2 model on Hugging Face and run it on a long-context retrieval task. Where does it match Transformers? Where does it lag?

Progress

Progress is local-only — sign in to sync across devices.
Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)
Sign inPlease sign in to comment.

No comments yet — be the first.