Before We Begin

The first article in this series, on the Linear State-Space Layer (LSSL), traced the development of architectures that define effective latent spaces for continuous sequential data. HiPPO supplied the key theoretical basis for modeling a hidden state over sequences of arbitrary length. LSSL then extended that idea into a layer and showed that learning the HiPPO matrices could improve performance. The second article covered S4, which improves training stability and numerical rigor while making the model more scalable in both data and model size.

Building on those state-space models, this article introduces Mamba, a model designed to avoid the unavoidable growth in computation associated with Transformers.

The Transformer's Biggest Problem

The deep-learning architecture used most often today, especially in foundation models, is the Transformer. First introduced for machine translation, it spread across computer vision, NLP, audio, and nearly every other modality. Among current architectures, it exhibits model scaling—the tendency for performance to improve as model and data size increase—particularly clearly.

Its largest problem is that it receives a bounded sequence—typically tokens produced by a tokenizer and converted to embeddings—and spends memory to process their attention in parallel. The length that can be handled at once is consequently limited.

A short sentence can simply be tokenized and processed. A paper or other long document must be split because computation is finite. More importantly, as the title “Attention Is All You Need” implies, a Transformer does not learn and carry a compact hidden state in the manner of an RNN. It performs its reasoning over the sequence supplied in the current context. Querying an entire paper for the passages relevant to a question therefore performs poorly unless enough computation and memory are available. This limitation has driven extensive work on more efficient attention.

Why Use a Transformer at All?

Why insist on a Transformer? Researchers have explored linear attention, gated convolutions, RNNs, and structured state-space models (SSMs) to improve efficiency, but on data such as language none consistently surpassed Transformer attention. The promise of Transformers also attracted an enormous body of work on hardware-aware algorithms and training methods. In such a mature ecosystem, an alternative with moderate compute but mediocre performance has difficulty gaining attention. SSMs, despite their recent focus on long-range dependencies and reasoning, performed poorly on information-dense tasks such as text.

I Choose You, Mamba

Mamba focuses on SSMs. An SSM fundamentally models a system; Mamba brings the “gate” concept from attention and RNNs into that system as selection. Attention itself is not mystical: for the current token, it distinguishes which accessible parts of the sequence to emphasize and which to ignore. An SSM can likewise use information selectively.

A selective SSM, however, loses the efficient convolutional machinery available to an ordinary time-invariant SSM. It must use the basic recurrent form. Mamba compensates with a hardware-aware algorithm that executes this recurrence efficiently. In short, it combines selective SSMs with an efficient implementation and eliminates both attention and the conventional MLP block. The paper's crucial contribution is therefore as much about GPU utilization—much like FlashAttention—as about the SSM itself.

Selection Mechanism

Previous SSMs lack a selective algorithm conditioned on the input.

An SSM can model sequences independently of their length. Viewed another way, however, every input at every time consults a state modeled in the same manner. The system cannot distinguish older or later input that ought to be ignored. The strength of an SSM is its ability to compress information from an arbitrarily long input, but the figure reveals its difference from attention.

Suppose a passage begins with “I have had a cat for five years,” contains an unrelated story about living in a country house, and states the cat's name near the end. Asked “How old is OOO likely to be?”, an attention model can focus on “The cat's name is OOO” and “I have had the cat for about five years.” A conventional SSM incorporates the intervening country-house story into the same state. It may still answer correctly, but the point is that attention remains exceptionally efficient and precise for information-dense data in which answering requires concentrating on particular regions.

Mamba adds selection by making the SSM parameters—and thus the influence of previous information—vary with the input signal.

Hardware-Aware Algorithm

Earlier efficient SSM algorithms do not apply to selective SSMs because they assume a linear time-invariant (LTI) system whose dynamics are independent of both time and input. Once the convolutional formulation is unavailable, only recurrent computation remains. Mamba therefore designs that recurrence around the hardware. The implementation achieves complexity linear in sequence length, making recurrence more effective than earlier convolution-based algorithms with pseudo-linear cost.

Modeling

The model architecture is simple. Treat the selective SSM as a module and use it in place of the attention and projection portions that make up a Transformer's mixer block. The previous articles provide the background on ordinary SSMs. Briefly, a structured SSM uses four parameters $(\Delta,A,B,C)$ and their discretized counterparts $(\overline{A},\overline{B},C)$. Classical SSMs assume LTI dynamics. Mamba gives up the corresponding computational convenience because invariance limits content-aware modeling, and adopts a selective SSM instead. The following sections show how.

Selection for Compression

Every sequence model confronts the same question: how can a context of arbitrary length be compressed into a small hidden state or latent? Each architecture makes a tradeoff. A Transformer performs almost no context compression. During autoregressive inference, it retains and consults the key–value context, so memory use and computation grow with sequence length; full attention is quadratic in training and prefill. A recurrent model such as an RNN has a fixed-size state and is computationally efficient, but may fail to summarize enough context within that limited state.

The generation tasks in the figure make this tradeoff easy to see at the level of an LTI SSM.

The task on the left copies a contiguous, marked portion of the input signal. This is a basic delay operation, something an ordinary LTI system can perform easily.

The upper-right task must distinguish relevant marked signals from irrelevant white signals and reproduce only the relevant items in input order. Unlike simple copying, this requires a time-varying or nonlinear system. The lower-right induction-head task is associated with in-context learning in modern LLMs. The model must infer a pattern from an input sequence and, when a particular black token appears again, output the contextually appropriate blue token. Because the correct successor depends on the observed input pattern, the system must reason over context.

The point is not limited to these toy generation tasks. Every sequence model must learn to understand context selectively. Existing methods balance context capacity against efficiency: attention filters individual features from an uncompressed context, while recurrence constructs a compact state that summarizes it. SSMs need an analogous mechanism, so Mamba introduces selection as a means of context compression.

Adding Selection to an SSM

Mamba's core question is how to insert selection into an SSM. The authors make parameters that directly affect the signal—analogous to recurrent dynamics in an RNN or filters in a CNN—depend on the current input.

Comparing the diagrams, $B$, $C$, and $\Delta$ are no longer independent of the batch input or its position in the sequence. They acquire dimensions corresponding to the input and output signals, changing the model from a time-invariant to a time-varying system.

The right diagram also shows how they are determined: $s_B(x)$, $s_C(x)$, and $s_\Delta(x)$ are functions of the input. In the simplest parameterization,

$$s_B(x)=\text{Linear}_N(x),\qquad s_C(x)=\text{Linear}_N(x),$$

where a linear module maps an input of dimension $D$ to an output of state dimension $N$. The discrete step size is

$$s_\Delta(x)=\text{Broadcast}_D(\text{Linear}_1(x)),$$

which produces a scalar-like value and broadcasts it across channels. The SSM dynamics have thus been parameterized as functions that vary with the input at each time.

Hardware-Aware Computation

Convolution and attention models have hardware-aware implementations. A convolution kernel has a fixed receptive field independent of input size, enabling predictable memory use. Attention still grows with sequence length, but algorithms such as FlashAttention exploit SRAM rather than repeatedly materializing data in high-bandwidth memory. FlashAttention author Tri Dao is also an author of Mamba.

Even though a selective SSM can no longer use the LTI convolution form, parts of its training can still be made efficient. The problem can be summarized as follows:

1. Recurrent models such as SSMs must balance representational state size against speed. The goal is high capacity without severe slowdown.
2. Recurrence is more flexible than convolution, but naïvely constructing latent states costs work over dimensions $(B,L,D,N)$. Models such as S4 proposed ways around this.
3. Earlier LTI state models used a dual recurrent–convolutional form so they could afford a sufficiently large state dimension $N$.

Selection prevents use of the LTI convolutional shortcut, so its inefficiency must be addressed directly. Two observations help:

• A straightforward recurrence costs $O(BLDN)$ operations, while the convolutional form costs $O(BLD\log L)$. As sequences grow, recurrence can actually require fewer operations when hidden dimension $N$ is moderate.
• Its main problems are sequential execution, which resists parallelization, and memory use. The latter can be reduced because, as with convolution, we need not materialize every complete state $h$ in high-bandwidth memory.

The idea is therefore practical rather than exotic: compute hidden state $h$ efficiently on the GPU using kernel fusion, parallel scans, and recomputation.

Rather than materialize discretized $\overline{A}$ and $\overline{B}$ in HBM, the implementation loads SSM parameters $A,B,C,\Delta$ into SRAM, performs discretization there, and writes only the necessary outputs back. A parallel-scan algorithm reduces the sequential bottleneck. This minimizes costly traffic between HBM and the smaller, faster SRAM. Backpropagation uses recomputation in essentially the same hardware-aware manner as FlashAttention.

Mixing Mamba into a Neural Network

Like the structured SSM in S4, Mamba's selective SSM is a sequence-transformation module and can be placed inside a neural network. Understanding the complete block benefits from knowing H3 and gated units.

H3, nicknamed Hungry Hungry Hippos, tries to reproduce the effect of Transformer attention with SSMs. It separates a shifting SSM from a recalling SSM and connects them multiplicatively. Input information corresponding loosely to $Q$, $K$, and $V$ is shifted so that earlier inputs can be referenced. When the current signal matches something remembered by the shifting path, the value path returns the corresponding output. A gated MLP similarly attempts to incorporate attention-like multiplicative interaction into an MLP architecture.

Because Mamba also uses SSMs, it shares H3's aim of reproducing the useful effects of attention. Its internal SSM computation resembles H3. The difference is architectural: H3 retains the $Q$, $K$, $V$ organization of linear attention, whereas Mamba discards attention entirely and wraps its selective SSM path in a gated-MLP-style block.

Effects of the Components

The main effects of the resulting design can be summarized as follows:

• Variable spacing: On discrete data such as language, the model can ignore signals irrelevant to context. Setting an effective gate $g_t$ near zero excludes irrelevant tokens and preserves a richer useful context.
• Context filtering: The importance of context can change with the situation. An LTI system cannot selectively discard part of its history, but a selective SSM can filter context dynamically.
• A Transformer can attend to contextually related phrases at discontinuous positions. An LTI system treats context as one continuous stream and lacks this property. In a selective SSM, $\Delta$ and an effective gate $g_t$ that can suppress the previous hidden state vary over time, potentially learning the same kind of discontinuous selection.
• The step size $\Delta$ controls how strongly to focus on the current input. When $\Delta\rightarrow\infty$, the discretized dynamics forget previous state and emphasize the present input. A small $\Delta$ preserves the preceding state more strongly.
• Parameter $A$ constructs the hidden-state dynamics, as in the original system. With a varying $\Delta$, its discretized transition gains selective behavior.
• Parameters $B$ and $C$ determine how much information from current input $x_t$ enters the state and how much state information contributes to output $y_t$:

$$\begin{aligned} \mathbf{h}_t^\prime &= A\mathbf{h}_{t-1} + B\mathbf{x}_t \newline \mathbf{y}_t &= C\mathbf{h}_t+D\mathbf{x}_t. \end{aligned}$$

Experimental Results

Mamba substantially outperforms earlier SSM baselines on both selective copying (left) and induction heads (right).

The improvement in perplexity is also striking. As compute and parameter count grow, context generation improves clearly. Earlier models through H3 had not matched attention on this dimension; Mamba closes much of that gap.

Zero-shot results across downstream tasks show decreasing perplexity and increasing average performance as parameter count grows. The model remains strong with relatively few parameters and also performs well on other sequential modalities, including DNA and audio.

Mamba's most important practical advantage emerges as context grows: its computational cost and inference time scale favorably. Even compared with FlashAttention, an optimized attention implementation, Mamba shows better inference time and throughput at long sequence lengths.