C.W.K.
Stream
Lesson 01 of 05 · published

What Mamba Is — Selectivity in One Lesson

~16 min · mamba, selectivity, s6

Level 0Observer
0 XP0/50 lessons0/14 achievements
0/100 XP to next level100 XP to go0% complete

One change: parameters become functions of the input

Mamba (Gu & Dao, December 2023, arXiv:2312.00752) is a Selective State Space Model — abbreviated S6 — and the headline change in one sentence is: B, C, and Δ become functions of the current input x_t instead of being fixed. That's it. That's the entire conceptual difference from S4/S4D/H3.

The mechanism: B and C are produced by linear projections of x_t. Δ — the discretization step size, conceptually a per-token "how much should this token influence the state" gate — is also produced from x_t (through a linear projection followed by softplus to keep it positive). When Δ is large, the model strongly incorporates the current token; when Δ is small, the existing state dominates. Δ is the selectivity gate — it lets the model say "this token matters, remember it" or "this token is filler, skip it".

Why this is the right kind of expressiveness gain

The LTI restriction wasn't just a math limitation; it had a concrete cost. Pre-Mamba SSMs couldn't do content-aware compression. Every word was processed by the same time-invariant filter, regardless of whether it was a key fact or a stop word. With selectivity, the model can route the high-information tokens into the state and let the low-information ones pass through with minimal effect.

That's analogous to what gates in LSTMs and GRUs were doing — selectivity is, in a sense, the SSM family rediscovering content-dependent gating after S4/S4D/H3 had simplified it away. The difference is that selectivity sits inside the SSM framework, so you keep the rich connection to control theory and the hardware-friendly structure, while gaining the gate that LSTMs had decades ago.

The price: the convolution view dies

There's no free lunch. Once B, C, Δ depend on x_t, you can no longer unroll the recurrence into a single global convolution kernel — there isn't a kernel anymore, there's a different one at every timestep. So the FFT-based parallel training trick from S4 doesn't work. Mamba had to invent a different way to be parallel during training, which is the next lesson: the hardware-aware selective scan.

Code

Mamba S6 step — pseudocode of the per-token update·python
# x_t: (batch, d_model) — input at timestep t
# A: (d_state, d_state) — fixed (initialized HiPPO-style)
# h_{t-1}: (batch, d_state) — previous state

B_t = linear_B(x_t)            # (batch, d_state)
C_t = linear_C(x_t)            # (batch, d_state)
delta_t = softplus(linear_d(x_t))  # (batch, 1) — the selectivity gate

A_bar = torch.exp(delta_t * A)         # discretized state transition
B_bar = delta_t * B_t                  # discretized input projection

h_t = A_bar @ h_{t-1} + B_bar * x_t    # state update (now input-dependent)
y_t = C_t * h_t                        # output

External links

Exercise

Install mamba-ssm (CUDA required) and run the Mamba-130M checkpoint from Hugging Face on a short generation task. Time prefill at 1K, 8K, 32K, 128K tokens and compare to Pythia-160M (a Transformer of similar parameter count) at the same lengths. You should see Mamba's prefill time grow ~linearly while Pythia's grows ~quadratically — the selectivity didn't change that asymptotic story.

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.