Mamba part 3 - Details of Mamba and Structured State Space

Mamba’s core pitch is that sequence modeling can be made both fast and selective without attention’s quadratic cost. The approach builds on state space models (SSMs)—specifically the S4 line of work—then adds “selectivity” so the system’s internal transition behavior changes with the input. That combination aims to keep computation efficient on GPUs while still letting the model treat different tokens differently, which matters for language where some words carry more context than others.

The explanation starts with S4’s state space framing. Inputs (denoted U) are blended with a hidden state (X) that acts like a memory of prior context. In continuous time, the hidden state evolves according to a differential equation; discretization turns that into a step-by-step recurrence. A key practical advantage comes from two equivalent views: one that updates the hidden state sequentially (good for inference), and another that “unrolls” the recurrence into a form that enables parallel gradient computation across long sequences (good for training). This is why S4 is described as fast at inference with one-step updates, and fast at training by parallelizing across the whole sequence.

However, S4 struggled when moved from modalities like signals (the example given is echocardiograms) to language. The intuition offered is that language doesn’t behave like a repeating physical cycle: some tokens are crucial while others are effectively background, and the model needs context-dependent transitions rather than time-position-only dynamics. That’s where Mamba’s selective state space (S6 / “Mamba part 3” discussion) enters.

The transcript describes improvements associated with the S6 direction: a technique inspired by “Hunger Hippos” (H3) to collapse large matrix structure into a more efficient diagonal form, reducing computational burden. It also introduces token-dependent parameters—each token can get its own effective matrices (A, B, C, D variants), and the system performs an input-driven selection mechanism (described as “attention-free” selection rather than standard attention). The result is that the model can modulate how strongly it “forgets” old state versus “remembers” new input depending on what token arrives.

A major engineering theme then follows: making the recurrence run efficiently on GPUs. The discussion emphasizes that performance bottlenecks often come from memory movement, not raw arithmetic. Mamba’s implementation uses custom CUDA kernels, parallel scan strategies (referencing S5), and kernel fusion to avoid slow round-trips to main GPU memory and to compute recurrences in a way that fits GPU hardware well. The outcome is linear scaling with sequence length during training and constant-time-per-step behavior during autoregressive inference, without maintaining a key-value cache like Transformers.

Finally, the transcript addresses why this doesn’t collapse into a vanilla RNN. The state space formulation constrains the allowed dynamics through the continuous-time-to-discrete derivation, producing stable “forget/remember” behavior rather than arbitrary recurrent weights. The discussion also touches on vanishing gradients: while long-range influence still decays, the model’s well-conditioned transitions and gating-like structure aim to keep training numerically stable, and residual connections across layers help optimization. Overall, Mamba is presented as a middle path between Transformers’ full-history attention and RNN-style recurrence—trading quadratic attention for selective, hardware-optimized state space computation that can still capture long context.