DL0170 Mamba Selective SSM vs Attention

Mamba’s S6 layer is a linear recurrence that carries a fixed-size hidden state per channel and reads it out with a learned projection, exactly like the earlier S4 state space model. The one change that matters is selection: the timestep and the input and output maps are linear functions of the current token instead of fixed parameters, which turns the discretized decay into an input-dependent forget gate. That single change gives the model the ability to skip filler tokens and write informative ones into the state, but it destroys linear time invariance, so the layer can no longer be evaluated as one long FFT convolution and instead needs a hardware-aware parallel associative scan. On complexity the comparison is lopsided in Mamba’s favor: sequence mixing is instead of , and decoding needs a constant state rather than a KV cache that grows with context. On recall the comparison reverses, because a bounded state is a lossy summary of the past, so tasks that need verbatim retrieval of an arbitrary earlier span degrade in a way that attention’s exact, if expensive, cache does not.

(1) Selection Is The Whole Idea: become projections of , while stays a learned diagonal matrix whose effective decay is nonetheless input-dependent.
(2) Is A Forget Gate: small gives and the token is ignored, large drives toward zero and overwrites the state with the new input.
(3) Time Invariance Is Lost: S4 could precompute one global convolution kernel, Mamba cannot, so training uses a work-efficient associative scan in SRAM with activation recomputation instead of an FFT.
(4) Linear Compute: mixing costs with in Mamba-1, so the quadratic attention term dominates once approaches the model width .
(5) Bounded Decode State: values per layer regardless of context length, versus for an MHA KV cache, which is what makes million-token streaming decode cheap.
(6) Recall Is The Trade: fixed capacity caps exact copying and multi-query associative recall, which is precisely why production long-context models are usually hybrids with a few full-attention layers.

Figure 1: The recurrence is ordinary, the gate is not. Because is computed from the token itself, a filler word produces and passes through without disturbing the state, while a rare identifier produces a large step that writes into the state and forgets older content. The state width never depends on , which is both the efficiency win and the recall ceiling.

The engineering consequence of dropping time invariance is that the layer becomes memory-bound rather than FLOP-bound. Materializing the expanded states, which have shape , would move far more bytes through HBM than the arithmetic justifies, so the reference implementation fuses discretization, scan, and readout into one kernel that keeps states in SRAM and recomputes them during the backward pass. The recall side has a cleaner theoretical story. Attention stores every key and value, so retrieving a specific earlier token is a lookup, whereas a selective SSM must have decided at write time to keep that token, and its slots bound how many distinct key-value associations can survive. This is why Mamba matches or beats Transformers on language modeling perplexity, audio, and DNA, yet lags on induction-head style copying and needle-in-a-haystack retrieval, and why Mamba-2 raises from 16 to as much as 256 while hybrids keep a small number of attention layers to do the exact lookups.

Mathematical Formulation:

Where:

• is the layer output for token and is the hidden state held for one channel, so the full layer keeps values under the standard expansion factor of 2.
• is the input activation, indexes the sequence, and is the required initial condition.
• is a learned diagonal matrix parameterized as so every eigenvalue is negative and the discrete decay stays stable.
• is the input-dependent timestep, and are the low-rank projections that make the layer selective; is initialized so that starts in a useful timescale range.
• is the zero-order-hold discretization of , and is the simplified input term Mamba uses in place of the exact .
• is the model width, the state dimension (16 in Mamba-1), and the sequence length; the scan is work with depth.

Decode Memory At 32k Context (values per layer):

Those are element counts, so in bytes the gap narrows a little when the cache is fp16 and the SSM state is fp32, and it narrows further with grouped-query attention, which divides the cache by the query-to-key-value head ratio. The structural point survives every such adjustment: the attention term is proportional to and the SSM term is not, so past a few thousand tokens the recurrent model is decoding from a constant working set while the Transformer is streaming a cache that eventually dominates both memory and bandwidth.

Figure 2: Below roughly the wider Mamba block, which expands the channel dimension by 2, actually costs more than attention, because both are dominated by their projection matmuls. Above it the term takes over, and at 128k tokens the attention layer needs about 20 times the FLOPs and a per-layer cache about 4,000 times larger than the constant SSM state.