A different question than "how fast is inference"
A sibling chapter in this course already walks through why a selective state space model answers each new token in constant time while a Transformer's KV-cache keeps growing, and it works a throughput comparison at a fixed sequence length. Take that result as given. This chapter opens the box the other one left closed: what actually happens inside the recurrence at each timestep, why that mechanism is called "selective," and why implementing it efficiently required NVIDIA GPU engineering, not just linear algebra. If you have not read the sibling chapter yet, the two are ordered — general continuous-time math and HiPPO first, this second — but they teach genuinely different content and neither restates the other's numbers.
Start with a National Payments Corporation of India-style problem. A bank's fraud engine watches a continuous stream of UPI transactions per account and must maintain a running "risk state" that updates online, one transaction at a time, in constant memory — it cannot re-scan the entire day's history before every decision. Most transactions are routine: a ₹40 auto-rickshaw payment, a ₹15 tea-stall scan. A handful are anomalous: an unusually large transfer, a new payee at 3 a.m. The engine's state should barely move on the routine ones and jump sharply on the anomalous ones. A plain linear recurrence — the kind an S4-style state space model uses, with fixed transition dynamics — cannot do this: it applies the identical update rule to every token regardless of content, because its dynamics are frozen at training time and do not look at what actually arrived. Mamba's central contribution, introduced by Albert Gu and Tri Dao in "Mamba: Linear-Time Sequence Modeling with Selective State Spaces" (2023), is to make the recurrence's own parameters a function of the current input — so the model can choose, token by token, whether to write to memory, skip a token entirely, or reset. That choice is the "selection mechanism," and it is the actual subject of this chapter.
The scalar selective recurrence, defined precisely
Strip a single channel of Mamba down to its bones. The continuous-time state space model is h'(t) = A h(t) + B x(t), y(t) = C h(t), discretized with step size Δ. For a fixed scalar A, the exact zero-order-hold (ZOH) discretization gives Ā = exp(ΔA) and, when A ≠ 0, B̄ = (ΔA)⁻¹(exp(ΔA) − 1)·ΔB. In an ordinary S4 layer, Δ, A, B, C are all learned constants — the same for every position in every sequence. Mamba's selection mechanism replaces this with three input-dependent projections: Δₜ = softplus(Linear(xₜ)), Bₜ = Linear(xₜ), Cₜ = Linear(xₜ), so the discretized matrices Āₜ, B̄ₜ are recomputed at every timestep from the token that just arrived. A is kept fixed per channel (structured, usually diagonal and negative for stability) — only the step size and the input/output projections become selective. The softplus keeps Δₜ positive, which matters: Δ plays the role of a time-step, and a near-zero Δ means "barely advance the ODE," while a large Δ means "take a big step."
Set up the smallest possible instance that still shows the effect honestly: one scalar channel, A = −1 (fixed), C = 1 (fixed, so the output is just the state itself), and B fixed at 1 so B̄ₜ reduces to the commonly used Euler simplification B̄ₜ = Δₜ·B = Δₜ. With A = −1 the exact ZOH formula for Ā also simplifies cleanly: Āₜ = exp(−Δₜ). The only thing that varies with content is Δₜ, produced by a one-parameter gate Δₜ = softplus(xₜ − 5) — a stand-in for a learned linear layer that has come to treat "5" as the boundary between routine and anomalous activity for this channel.
Tracing the recurrence by hand
Feed the fraud channel the (already-embedded) activations x = [3, 0, 0, 5] — mildly unusual, routine, routine, highly anomalous — with h₀ = 0.
t = 1, x₁ = 3. Δ₁ = softplus(3 − 5) = softplus(−2) = ln(1 + e⁻²) = ln(1.135335) = 0.126928. Ā₁ = exp(−0.126928) = 0.880842. h₁ = Ā₁·h₀ + Δ₁·x₁ = 0.880842·0 + 0.126928·3 = 0.380784.
t = 2, x₂ = 0. Δ₂ = softplus(0 − 5) = softplus(−5) = ln(1 + e⁻⁵) = ln(1.006738) = 0.006715. Ā₂ = exp(−0.006715) = 0.993307. h₂ = 0.993307·0.380784 + 0.006715·0 = 0.378235.
t = 3, x₃ = 0. Identical gate to t = 2: Δ₃ = 0.006715, Ā₃ = 0.993307. h₃ = 0.993307·0.378235 = 0.375703.
t = 4, x₄ = 5. Δ₄ = softplus(5 − 5) = softplus(0) = ln 2 = 0.693147. Ā₄ = exp(−0.693147) = 0.5 exactly. h₄ = 0.5·0.375703 + 0.693147·5 = 0.187852 + 3.465735 = 3.653587.
Read the sequence of Ā values: 0.881, 0.993, 0.993, 0.500. On the two routine tokens Āₜ sits above 0.99 — the state is carried forward almost untouched, and the input contribution (Δₜ·xₜ = 0.0067·0 = 0) is exactly zero anyway, so h barely drifts (0.3808 → 0.3782 → 0.3757, a 1.3% decay over two full timesteps of "doing nothing"). On the anomalous token at t = 4, Āₜ collapses to 0.5 — half the old state is deliberately discounted — while the injected term Δ₄·x₄ = 3.4657 dwarfs everything that came before it, and the state jumps nearly tenfold. That is selectivity in one recurrence: the same fixed A = −1 dynamics, but a gate that decides, from the content of xₜ alone, whether to preserve memory or overwrite it. An S4 layer with fixed Δ could never produce this pattern — every token would move the state by the same proportional amount regardless of whether it was a ₹15 tea payment or a 3 a.m. anomaly.
This is exactly reproducible in code, with no helpers:
import math
def softplus(z):
return math.log1p(math.exp(z))
def selective_scan_scalar(xs, A=-1.0, w_delta=1.0, b_delta=-5.0):
h = 0.0
outputs = []
for x in xs:
delta = softplus(w_delta * x + b_delta) # input-dependent gate
A_bar = math.exp(delta * A) # exact ZOH, A fixed
B_bar = delta # Euler-simplified B_bar = delta * B, B = 1
h = A_bar * h + B_bar * x
outputs.append(h)
return outputs
print(selective_scan_scalar([3.0, 0.0, 0.0, 5.0]))
# [0.3807840331289175, 0.37823550236083775, 0.37570402852926466, 3.653587917064359]Rounded to six decimal places, this is 0.380784, 0.378235, 0.375703, 3.653587 — reproducing the four numbers derived above, because it is the identical computation: softplus is defined inline, math.exp/math.log1p are standard library, and the loop body performs exactly the ZOH-discretize-then-recur step traced by hand.
Diagram: gate strength controlling memory vs. update
Why this breaks the convolution trick — and what fixes the speed
An S4-style layer with fixed Ā, B̄, C can be unrolled into a single global convolution: yₜ = Σₖ C·Āᵏ·B̄·xₜ₋ₖ is a convolution of x against a fixed kernel built once from Ā, B̄, C, and can be computed for an entire sequence in O(L log L) via FFT — the reason S4 trains fast. The moment Āₜ and B̄ₜ depend on t (because they depend on xₜ), that trick is gone: there is no single kernel, because the effective kernel is different at every position. You are forced back to explicit recurrence, and a naive Python-style for-loop over L timesteps is exactly what GPUs are bad at — each step depends on the last, so the L steps cannot run as L independent parallel threads, and worse, materializing every intermediate state h in a tensor of shape (batch, length, channels, state-dim) means writing and re-reading a huge tensor from HBM (a GPU's off-chip high-bandwidth memory) at every step.
The fix has two independent parts, and conflating them is a common error. First, the recurrence itself is made parallel using an associative scan. Rewrite step t as an affine map: hₜ = aₜ·hₜ₋₁ + bₜ where aₜ = Āₜ, bₜ = B̄ₜ·xₜ. Two consecutive affine maps compose into one: applying (a₁, b₁) then (a₂, b₂) is equivalent to the single map (a₂a₁, a₂b₁ + b₂). This composition operator is associative, which means all L prefix-compositions — i.e., all L states h₁ … h_L — can be computed by a Blelloch-style parallel prefix scan in O(log L) sequential depth with O(L) total work, the same trick used for parallel prefix sums, and the same one Jimmy Smith, Andrew Warrington, and Scott Linderman used for SSMs slightly before Mamba in "Simplified State Space Layers for Sequence Modeling" (S5, ICLR 2023). This is a depth reduction, not a FLOP reduction — more on that below.
Second, and separately, Mamba's actual CUDA kernel fuses the discretization step, the scan, and the output projection into a single kernel so the large (batch, length, channels, state-dim) intermediate tensor never round-trips through HBM — it is computed and consumed inside GPU SRAM, which has roughly an order of magnitude more bandwidth than HBM. This is the same IO-awareness principle Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher Ré used for attention in FlashAttention (NeurIPS 2022): the bottleneck for scan-heavy operations is memory traffic between GPU memory tiers, not arithmetic, so the win comes from restructuring where data lives during the computation, not from doing less arithmetic. Mamba additionally uses recomputation in the backward pass — rather than storing every intermediate hₜ for gradient computation (which would need the same expensive HBM traffic), it recomputes them on the fly during backpropagation, trading cheap, parallel extra FLOPs for memory-bandwidth savings, exactly the gradient-checkpointing tradeoff FlashAttention also makes.
Tri Dao and Albert Gu extended this in "Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality" (Mamba-2, ICML 2024). Restricting the state matrix A to a scalar times identity lets the same recurrence be written as a multiplication by a structured (semiseparable) lower-triangular matrix — and that matrix-vector product can be computed two equivalent ways: as a linear scan (state size independent of sequence length — the SSM view) or as a block matrix multiplication with a mask, which looks like attention restricted to tensor-core-friendly matmuls. The paper reports the resulting SSD algorithm running several times faster in wall-clock training throughput than the original selective-scan kernel, while supporting a larger per-channel state dimension than Mamba-1's typical N ≈ 16. The "duality" in the name is literal: it is the same underlying computation, expressed once as a recurrence and once as attention-shaped matmuls, and you pick whichever form your hardware executes faster.
Common misconception
Because Mamba is described as a recurrence, students often conclude it must process a training sequence strictly one token after another, the way a plain RNN or LSTM does — and therefore cannot be trained in parallel across the sequence dimension the way a Transformer can. This is wrong, and the wrongness is exactly the associative-scan point above. During training, the whole sequence is known in advance, so the parallel scan computes all L states with O(log L) sequential depth: the sequence dimension parallelizes across GPU threads, not one token at a time. The strictly sequential, one-token-at-a-time behavior is real, but it belongs to autoregressive inference, not training — and there it is a deliberate design choice, not a limitation, since it is what gives Mamba constant per-token cost instead of a growing attention computation. Conflating "the mathematical form is a recurrence" with "the training loop is a Python for-loop" is the error; the associative scan is precisely the piece of engineering that makes the first true without the second following from it.
Active recall
Attempt each question before reading its answer.
Q1. Why can't an S4-style global convolution kernel be used to train a selective SSM the way FFT convolution trains plain S4?
Q2. In a naive implementation of the selective scan, is the GPU bottleneck FLOPs or something else? What is Mamba's fix, and which earlier paper's principle does it borrow?
Q3. Using the same toy model (A = −1, B = 1, C = 1, Δₜ = softplus(xₜ − 5), h₀ = 0), compute h₁, h₂, h₃ for the sequence x = [1, 0, 6].
Q4. The fraud team recalibrates the gate bias from −5 to −3, so Δₜ = softplus(xₜ − 3), keeping the original sequence x = [3, 0, 0, 5] and A = −1, B = 1, C = 1, h₀ = 0. Recompute Δ₁…Δ₄ and h₁…h₄, and explain what changed qualitatively about which transactions the model treats as anomalous.
Q5. True or false: the parallel associative scan lowers Mamba's asymptotic FLOP count compared to a naive sequential recurrence. Justify your answer.
Q6. In what precise sense does Mamba-2's "Structured State Space Duality" claim that selective SSMs and attention are the same computation?
A1. A fixed global kernel exists only because Ā and B̄ are the same at every position, so yₜ = Σₖ C·Āᵏ·B̄·xₜ₋ₖ is a true convolution against one static filter, computable via FFT. Selectivity makes Āₜ and B̄ₜ functions of xₜ, so the "kernel" applied at each output position is different — there is no single fixed filter to convolve with, and the computation must fall back to an explicit (parallel) scan over the recurrence.
A2. The bottleneck is memory bandwidth, not arithmetic: a naive implementation materializes the full (batch, length, channels, state-dim) intermediate tensor in off-chip HBM at every step. Mamba's fix is kernel fusion — computing discretization, scan, and output projection in one CUDA kernel so intermediates stay in on-chip SRAM — the same IO-awareness principle Dao, Fu, Ermon, Rudra, and Ré used in FlashAttention (2022).
A3. x₁=1: Δ₁ = softplus(1−5) = softplus(−4) = ln(1+e⁻⁴) = ln(1.018316) = 0.018150. Ā₁ = exp(−0.018150) = 0.982014. h₁ = 0.982014·0 + 0.018150·1 = 0.018150. x₂=0: Δ₂ = softplus(−5) = 0.006715 (as in the worked example). Ā₂ = 0.993307. h₂ = 0.993307·0.018150 = 0.018028. x₃=6: Δ₃ = softplus(6−5) = softplus(1) = ln(1+e¹) = ln(3.718282) = 1.313262. Ā₃ = exp(−1.313262) = 1/3.718282 = 0.268939. h₃ = 0.268939·0.018028 + 1.313262·6 = 0.004848 + 7.879572 = 7.884420.
A4. x₁=3: Δ₁ = softplus(3−3) = softplus(0) = ln2 = 0.693147; Ā₁ = exp(−0.693147) = 0.5; h₁ = 0.5·0 + 0.693147·3 = 2.079442. x₂=0: Δ₂ = softplus(0−3) = softplus(−3) = ln(1+e⁻³) = ln(1.049787) = 0.048587; Ā₂ = exp(−0.048587) = 0.952576; h₂ = 0.952576·2.079442 = 1.980827. x₃=0: Δ₃ = 0.048587, Ā₃ = 0.952576; h₃ = 0.952576·1.980827 = 1.886888. x₄=5: Δ₄ = softplus(5−3) = softplus(2) = ln(1+e²) = ln(8.389056) = 2.126928; Ā₄ = exp(−2.126928) = exp(−2)·exp(−0.126928) = 0.135335·0.880842 = 0.119209; h₄ = 0.119209·1.886888 + 2.126928·5 ≈ 0.224934 + 10.634640 ≈ 10.8596. Every one of h₁…h₄ changed, not only the terms touching x₁ or x₄ — because the recurrence is a chain, a change to the gate's bias ripples through all four steps. Qualitatively, lowering the bias from −5 to −3 raised Δ at every position, including the routine x=0 steps (Δ rose from 0.0067 to 0.0486, roughly 7×), so Ā₂ and Ā₃ dropped from about 0.993 to about 0.953 — the "routine" tokens now leak noticeably instead of being almost perfectly preserved. And x=3, previously "mildly unusual" (Δ₁=0.127), now triggers the gate almost as strongly as x=5 did under the old threshold (Δ₁=0.693 ≈ Δ₄_old=0.693) — the recalibration silently redefined what counts as anomalous for the whole stream, not just for the specific transactions that changed.
A5. False. The parallel scan performs the same asymptotic O(L) total work as the sequential recurrence — composing two affine maps costs a small constant more than one fused multiply-add, so the scan is if anything slightly more total arithmetic, not less. Its benefit is reducing sequential depth from O(L) to O(log L), which lets the L steps execute across parallel GPU threads instead of waiting on each other; it changes wall-clock time via parallelism, not the FLOP count.
A6. Restricting A to a scalar multiple of the identity lets the SSM recurrence be written as multiplying the input sequence by one particular structured (semiseparable, lower-triangular) matrix. That same matrix-vector product can be computed either as a linear scan — cost independent of state size beyond N, ideal for long sequences — or as an explicit masked matrix multiplication, which is the same computational shape as (linear) attention and runs efficiently on tensor cores. "Duality" means these are two algorithms for the identical underlying structured-matrix computation, not two different models that happen to resemble each other — Mamba-2 exploits this by choosing whichever form is faster for the hardware and sequence length at hand.
Think About It
Think about this: How would you explain state space models and mamba: linear-time sequential processing to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.
Practice Exercises
Now it is time to practice! Complete these challenges to solidify your understanding:
- Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
- Exercise 2: Find a real-world example where state space models and mamba: linear-time sequential processing is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
- Exercise 3: Create a mind-map connecting state space models and mamba: linear-time sequential processing to at least 3 other topics you have studied.
Key Takeaways — Summary and Recap
Let us recap what we covered: the core ideas behind state space models and mamba: linear-time sequential processing, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.