Featured image of post Mamba: Why Remembering Less Can Understand More

Mamba: Why Remembering Less Can Understand More

Mamba replaces a growing KV cache with a fixed recurrent state. This article explains selective state spaces, hardware-aware scans, and the empirical limits that make hybrid models compelling.

Transformer attention makes an expensive but cautious choice. It does not decide too early which history matters; instead, it keeps every old token’s key and value in a KV cache and lets each new query compare against them. Mamba adopts the opposite contract: history is compressed into a fixed-size state as it passes, and the original list is no longer available for arbitrary look-back.

That difference is more consequential than replacing $O(L^2)$ with $O(L)$. Compressing history is easy. Preserving the right information after compression is hard. Earlier linear time-invariant state space models (LTI SSMs) applied the same dynamics to every token. They behaved like filters that always updated in the same way: useful for continuous signals, but poorly suited to language decisions such as “retain this name, ignore that filler word.” Mamba made the important state-update parameters depend on the current input, then built a hardware-aware parallel scan for the resulting model after it had lost the convenient convolution path.

The most useful way to understand Mamba is therefore not “a model without attention,” but as three coupled choices: compress context into finite state, make the compression policy input-dependent, and recover GPU-friendly training through an efficient scan.

Two Context Contracts: Archive and Working Memory

Let the sequence length be $L$. Standard self-attention explicitly constructs pairwise token interactions during training, so compute and the attention matrix grow with $L^2$. Autoregressive generation computes only one new query per step, but it stores and reads a KV cache that grows with $L$. Its strength comes from the same design: the old representations remain available, so a later query can retrieve them in a new way.

A state space model uses a different form. After discretization, a minimal recurrence is:

$$h_t=\bar A h_{t-1}+\bar Bx_t,\qquad y_t=Ch_t$$

$h_t$ is a compressed state for everything seen through position $t$. Each token combines the previous state and new input into the next state. During generation, each layer keeps only the current state, so cache size does not grow with context and per-step work does not increase with prefix length. The trade-off is just as clear: if some information did not enter $h_t$, a later step cannot return to the original token as attention can.

Figure 1: attention keeps an addressable KV archive; Mamba continuously compresses history into a fixed state

This creates two meanings of “long context.” For attention, a longer window primarily means more addressable records, paid for with a larger cache and more reads. A recurrent SSM processes a longer sequence with linear total work, but its fixed state does not automatically gain capacity with a larger window. Being able to process one million tokens is not the same as remembering one million tokens losslessly. The original Mamba paper observed continued gains up to length one million on a DNA task. That is evidence for a particular dataset and objective, not proof of infinite memory.

Selectivity Does Not Drop Tokens; It Changes the Update

In a traditional LTI SSM, $Δ,A,B,C$ remain fixed across positions. The model can be viewed either as a recurrence or expanded into a global convolution, which makes parallel training convenient. Fixed dynamics, however, cannot alter the memory policy according to content. The paper illustrates the gap with Selective Copying: relevant symbols appear at varying positions, so the model must decide when to write based on token identity rather than wait for a fixed offset.

Mamba keeps a structured $A$ but makes $Δ_t$, $B_t$, and $C_t$ functions of the current input $x_t$:

$$Δ_t=s_\Delta(x_t),\qquad B_t=s_B(x_t),\qquad C_t=s_C(x_t)$$

After discretization, the transition itself varies by position:

$$h_t=\bar A_t h_{t-1}+\bar B_t x_t,\qquad y_t=C_t h_t$$

The three terms answer different questions. $B_t$ determines how the current input is written into state. $C_t$ determines what is read from state now. $Δ_t$ changes the timescale between preserving the old state and accepting the current input. In the paper’s mechanical interpretation, a larger $Δ_t$ tends to reset state and focus on the current input, while a smaller $Δ_t$ tends to preserve state and ignore a transient input. This is not a discrete keep/delete label assigned before processing. It is a continuous change to information flow at every step.

Figure 2: input-dependent Δ, B, and C govern forgetting/updating, write direction, and readout

That is why Mamba should not be equated with SSMs in general. In the paper’s ablations, replacing a non-selective SSM with the selective S6 layer substantially improved language-model perplexity. $Δ$ was the most important individual selective parameter, and making $Δ$, $B$, and $C$ selective together worked best. The crucial repair for discrete language was not “state space” by itself, but content-aware compression.

Selectivity Removes the Convolution Shortcut; Scan Restores Parallelism

An LTI SSM with fixed parameters can express an entire sequence as one convolution and train in parallel. Once the parameters vary with each token, there is no single stationary kernel for the sequence. A naive implementation would execute $t=1,2,\ldots,L$ in order, leaving a GPU underused by many small serial operations.

Mamba exploits the associativity of composing recurrences. Each position can be represented as a local transform $(\bar A_t,\bar B_tx_t)$. Neighboring transforms can be combined first and then combined with other intervals. Training can therefore use a parallel scan with a tree-shaped reduction, while generation uses the genuine one-step recurrence. One equation receives two execution plans: sequence parallelism for batch training and constant-size state for autoregressive decoding.

Figure 3: the same selective recurrence uses parallel scan for training and one-step state updates for generation

The algorithm alone is insufficient. A selective SSM conceptually produces expanded state with shape roughly $B\times L\times D\times N$. Materializing those intermediates in high-bandwidth memory could allow memory traffic to consume the apparent linear-time advantage. The paper’s CUDA implementation fuses discretization, scan, and following operations in on-chip SRAM, avoids writing the large expanded state, and recomputes some intermediates during the backward pass. Under the paper’s benchmark setup, selective scan became faster than FlashAttention-2 beyond sequence length 2K and 20–40 times faster than a standard PyTorch scan; end-to-end autoregressive throughput reached up to five times that of a similarly sized Transformer. These measurements depend on hardware, tensor shapes, and implementation. They are not universal deployment multipliers.

Why a Mamba Block Also Absorbs the MLP

Original Mamba does more than replace a Transformer attention layer with an SSM. A block expands its input into two paths. The main path passes through a short causal 1D convolution, SiLU, and the selective SSM. The other path forms a gate. Their outputs are multiplied and projected back to the model dimension. This merges what earlier architectures separated into a sequence-mixing layer and an independent MLP. When the paper says the network has “no MLP blocks,” it means there is no separately stacked FFN sublayer, not that the block lacks linear projections or nonlinearities.

The short convolution mixes nearby tokens, the SSM propagates state over the sequence, and the gate controls output channels. Together they explain why studying only the recurrence misses much of the actual model. Mamba’s behavior comes from the combination of selective state, local convolution, gating architecture, and the hardware implementation—not one isolated equation.

The simplification also creates an engineering threshold. The official repository’s fast path depends on custom CUDA kernels, and supported state dimensions, data types, and devices affect realized speed. Falling back to a reference implementation on mismatched hardware does not automatically convert linear asymptotics into lower latency. Complexity moved from a growing KV cache into kernels and scan execution; it did not vanish.

Pure SSM Limits Became Clearer in Hybrid Models

The original results were strong. Under its training and evaluation setup, Mamba-3B outperformed same-size Transformers and matched Transformers roughly twice its size, across work that also covered language, audio, and genomics. A later controlled 8B study coauthored by Mamba’s authors and NVIDIA researchers supplied a more informative boundary. It trained Mamba, Mamba-2, and Transformer models on the same data for as many as 3.5 trillion tokens. Pure SSMs matched or exceeded Transformers on many tasks, but lagged on tasks demanding strong copying, in-context learning, or long-context reasoning.

That result follows directly from the context contract. Attention retains raw K/V records and can point to an early detail in response to a question that appears much later. Fixed state must anticipate what deserves storage. Selectivity improves that decision but cannot remove finite-state capacity. Copying a phone number, reproducing an exact span, or learning a new mapping demonstrated only in the prompt are precisely the tasks where the future query may be unknowable at write time.

The same study used a Mamba-2-Hybrid made of 43% Mamba-2, 7% attention, and 50% MLP layers. It reported that the 8B hybrid beat the 8B Transformer by an average 2.65 points across 12 standard tasks and was predicted to generate up to eight times faster. Across 23 additional long-context tasks, the hybrid closely matched or exceeded the Transformer on average. “Predicted up to 8×” is not a complete production benchmark, but the architectural direction is important: a small amount of attention can serve as exact addressable memory while SSM layers perform cheap sequential compression.

Mamba-2’s State Space Duality (SSD) further showed that structured SSMs and certain constrained forms of attention are not unrelated worlds. They are alternative computations over the same class of structured semiseparable matrices. Its core layer was 2–8 times faster than Mamba’s selective SSM in the paper’s benchmarks. The more accurate industry conclusion is not “SSMs will replace Transformers,” but: attention and state space layers expose different memory interfaces, and models can combine them by layer and task.

The Real Change Is the Inference Memory Budget

MLA, discussed in the previous article, still compresses the KV cache inside attention. Mamba goes further by changing what is cached. Instead of storing addressable K/V for every historical token, it maintains a state per layer that is rewritten as the sequence advances. Context changes from an ever-thickening archive into fixed-size working memory. That opens a different systems design space for streaming audio, genomics, sensors, and long generation: state memory and per-step latency need not grow with elapsed sequence length.

The budget must still be matched to the task’s memory demands. Attention’s explicit archive remains valuable for exact quotation, arbitrary look-back, and temporary learning inside a prompt. Selective state is attractive for continuous streaming, low cache usage, and stable per-step cost. Hybrid models are natural not because they are a timid compromise, but because summarizing memory and addressable memory are genuinely different capabilities.

Mamba therefore leaves a question more durable than any benchmark: how much of the past should a sequence model keep, and when should it decide to forget? Attention postpones the choice until read time; Mamba makes it at write time. Its contribution was to show that the latter, once equipped with content selection, parallel scan, and the right kernels, no longer has to remain an efficient but less capable substitute. It can become one of the main components from which we design a model’s memory system.

References

  1. Gu & Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces, 2023/2024.
  2. Dao & Gu, Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality, 2024.
  3. Waleffe et al., An Empirical Study of Mamba-based Language Models, 2024.
  4. State Spaces Models team, Mamba official implementation.
  5. NVIDIA Megatron-LM, Mamba and Mamba-2 model implementation.