Mamba: Linear-Time Sequence Modeling with Selective State Spaces

A small state has to choose what survives

Suppose a sequence contains a few useful signals separated by distractions. Keeping every earlier token makes later retrieval straightforward, but the stored history grows with the sequence. Keeping only a small array makes storage predictable. The harder question is what that array should retain when the next input arrives.

Mamba makes the update depend on the input. A token can change how much old state survives, what gets written and how the state is read. The paper motivates this with selective copying, where useful symbols appear at variable distances. The mechanism must respond to content as well as position. Paper v2, §3.1

Start with one input channel and three state coordinates. At position t, the update has two contributions: retained memory from the previous position and a write from the current input. A readout turns the updated array into an output.

The explorer below supplies incoming memory [0.2, −0.1, 0.3] before position zero. Step through the six inputs and inspect the retained and written terms coordinate by coordinate. Switch to fixed coefficients to keep the same inputs while changing the update rule.

The two paths in the figure give the state update and readout equations:

These are calculated values from authored coefficient presets. The changing coefficients illustrate selection; they are not learned saliency, pretrained activations or a successful trained copy task. In particular, a small retention factor does not automatically mean “ignore this token”: it suppresses earlier memory, while the separate write can still add new information. Understanding that interaction requires the parameterization behind a and w.

The input changes the coefficients, while A stays learned

For the selected real-valued diagonal implementation, Mamba uses a learned negative continuous parameter A and an input-dependent positive time scale delta. The transition to the next discrete state is:

Here c indexes expanded channels and n indexes state coordinates. With batch b, sequence length T, expanded width E and state size N, the current state is [b, E, N]. The scan input and output are [b, T, E]. A has shape [E, N] and does not change with the token. Delta has shape [b, T, E]. In the pinned mamba_simple block, B and C have shape [b, T, N] and are shared across the expanded channels. The projection weights are learned and shared over positions; their outputs depend on the current locally mixed input. Algorithm 2 and §3.6, pinned implementation

Because A is negative and delta is positive, each retention lies between zero and one. A larger delta advances further along these decaying dynamics. It also scales the write. Delta is a time scale, not a probability gate, and B and C can contain signed values. The update is therefore not generally a convex average of old state and current input.

There is a discretization detail worth keeping explicit. The paper introduces exact zero-order hold for the continuous system, whose input map includes an integral:

The official readable selective_scan_ref instead uses delta × B × x for input injection while retaining exp(delta × A) for the transition. The runnable reference and every numeric figure on this page follow that code. Calling this write term the exact zero-order-hold integral would describe a different computation. Paper Equation 4, official reference recurrence

With fixed discrete coefficients, unrolling a state-space model gives a fixed global convolution kernel. Once the coefficients change with the token, the influence of an earlier input depends on the intervening inputs too. The same stationary kernel no longer describes every position. The recurrence still works, but computing it efficiently now needs a different organization.

Compose the updates instead of waiting for every state

If every update waits for the preceding state, training appears sequential across T positions. But once the inputs and their coefficients are available, each update is a known affine map. Combining two consecutive maps gives another affine map:

Composition is associative, so neighboring intervals can combine in parallel while preserving chronological order. It is not commutative: swapping updates generally changes which information is retained. Select a coordinate and a prefix in the figure to inspect the composed pair and compare its state with sequential evaluation, including the incoming memory.

The browser and CPU prefix example expose the algebra. Their snapshot-based Hillis–Steele rounds have logarithmic dependency depth and O(T log T) pair work. They also materialize intermediate arrays for inspection. That is different from claiming they reproduce a work-efficient GPU scan.

The paper’s hardware contribution combines parallel scan with fusion and recomputation. Rather than write a full [b, T, E, N] expanded state tensor through high-bandwidth device memory between operations, the fused computation constructs and consumes states in faster on-chip storage and recomputes needed intermediates during backward. This reduces memory traffic; it does not make all training activations constant-size. Paper §3.3 and Appendix D

Training can regroup known updates this way. Autoregressive decoding still waits for the newly generated token before it can form that token’s coefficients. Before following that decode path, we need to place the selective operator inside the complete neural-network block.

The selective operator is one part of the Mamba block

The block first projects each model-width input into two expanded branches. One branch provides x; a short causal depthwise convolution and SiLU activation add local context. From that convolved x, smaller projections produce B, C and a low-rank representation that expands into delta. The selective scan then mixes information over the sequence. The other branch provides z, which multiplies the scan result through SiLU before an output projection restores the model width. Paper Figure 3

Enlarge figure

This complete trace uses a tiny mixer with model width 4, expanded width 8, state size 3 and convolution kernel width 3. Its six-position inputs and sinusoidal untrained weights are authored; the displayed tensors were recorded on CPU in float64. Select an operation to see its actual tensor shape and a labeled sample of its values. Samples are flattened excerpts, not a display of every element.

There are three distinct ways information travels here. The depthwise convolution mixes neighboring sequence positions within a channel. Linear projections mix channels. The selective recurrence carries a longer-range summary. The z gate changes the mixed output before it leaves the mixer; it is separate from the input-dependent state update itself.

There are also two distinct additions. D × x is a direct path inside the scan output, at expanded width. The enclosing residual adds the mixer’s projected result back to its model-width input, typically after a prenormalization. The figure shows that outside wrapper separately so the direct skip is not mistaken for the block residual. The original paper image and pinned code correspondence appear below with the executable reference.

The whole mixer is causal. Its complete-sequence result should therefore agree with processing one token at a time, provided the implementation preserves every piece of history that the next token needs.

What persists when the next token arrives

The SSM array is only one persistent buffer. The causal convolution also needs recent projected inputs. In the pinned Mamba implementation, the convolution allocation stores K samples per expanded channel, including the latest sample. Each decode step shifts the buffer and appends the arriving projected input before computing the convolution. Separately, the SSM updates its [b, E, N] state. Pinned step and cache allocation

The small streaming trace starts both buffers at zero. This differs from the nonzero incoming memory supplied to the earlier core recurrence demo. Inspecting the convolution buffer before and after arrival makes the short local memory visible beside the channel-by-state array.

The storage comparison is a calculation under a separate, explicitly stated configuration. If every stored element uses s bytes, the two Mamba buffers across L layers require:

A conventional full-context Transformer KV cache with T positions, H KV heads and head width P requires:

The factor two stores both keys and values. The figure specifies batch, layers, expanded width, state size, stored convolution width, KV heads, head width and element size. Move the sequence-length control to isolate the growth with T. Grouped-query attention changes the KV-head count; a sliding-window cache changes how many positions persist. Mixed state dtypes require accounting for the buffers separately.

These formulas count persistent inference state, excluding weights, temporary workspaces, allocator overhead and training tensors. They compare storage assumptions, not equally accurate models. Compressing all previous tokens into fixed-size memory creates a capacity constraint: a bounded state cannot promise exact retrieval of arbitrary histories. Making storage independent of context length does not remove that tradeoff.

The buffers now close the implementation loop: the same operations must agree as a recurrence, a prefix computation and a streamed block. The reference below makes that check runnable.

Original paper, equations and runnable CPU code

The six rows connect the coefficient projections, state update, scan algebra, local convolution, full mixer and persistent buffers to actual synchronized Python source. Displayed snippets and downloadable files come from the canonical example, rather than separate handwritten pseudocode. The self-check compares sequential and prefix states, checks causality and full-sequence versus token-step behavior, and verifies finite gradients. Float64 comparisons use absolute and relative tolerances of 1e−9; float32 checks use 1e−5.

This authored CPU adaptation follows official v2.2.4 (commit 95d8aba8), with upstream attribution retained. Its deterministic mechanics and random-weight smoke runs do not use a pretrained model.

Original-paper images: Albert Gu and Tri Dao, Mamba (arXiv v2), under CC BY 4.0. PDF excerpts are cropped; Figure 3 is unchanged. Select any source image to enlarge it.

Download Python reference · Pinned requirements

python -m pip install -r requirements-mamba.txt
python mamba.py --self-check
python mamba.py --smoke --backward

1. Produce coefficients from the locally mixed input

Original paper

Enlarge figure

Math and tensor shapes

[r_t,B_t,C_t]=W_{x}x_t

\Delta_t=\operatorname{softplus}(W_{\Delta}r_t+b_{\Delta}),\quad A=-\exp(A_{\log})

Here E is expanded channel width, not the expansion factor. B and C are shared across channels in this selected Mamba implementation. The low-rank delta projection has rank ceil(model_width/16). Softplus produces a positive time scale, which affects retention and writing together.

x: [b,T,E]; B/C: [b,T,N]; delta: [b,T,E]; A: [E,N]

Pinned official code

Python implementation

python

def coefficients(self, x):
    """Produce token-dependent delta/B/C after the local convolution."""
    low_rank, b, c = self.x_proj(x).split(
        (self.rank, self.state_size, self.state_size), -1
    )
    delta = F.softplus(self.dt_proj(low_rank))
    a = -self.A_log.exp()  # Learned, input-independent continuous dynamics.
    return delta, a, b, c

2. Retain, write and read each channel’s state

Original paper

Enlarge figure

Math and tensor shapes

a_{t,c,n}=\exp(\Delta_{t,c}A_{c,n})

h_{t,c,n}=a_{t,c,n}h_{t-1,c,n}+\Delta_{t,c}B_{t,n}x_{t,c}

y_{t,c}=\sum_n C_{t,n}h_{t,c,n}+D_cx_{t,c}

The readable reference writes delta*B*x; this differs from the exact zero-order-hold input integral in paper Equation 4. Read after updating the state, then add the direct D*x path. An optional incoming state is preserved, so the loop also describes continuation from an earlier prefix.

Current state: [b,E,N]; sequence output: [b,T,E]

Pinned official code

Python implementation

python

def selective_scan(x, delta, a, b, c, d=None, initial_state=None):
    """Return outputs and all states of a diagonal selective recurrence.

    Args:
        x: Content [batch, length, channels].
        delta: Already positive step sizes, same shape as x.
        a: Negative continuous dynamics [channels, state].
        b: Input-dependent write coefficients [batch, length, state].
        c: Input-dependent read coefficients, same shape as b.
        d: Optional direct skip [channels], separate from the outside residual.
        initial_state: Optional state [batch, channels, state].

    Returns:
        Outputs [batch,length,channels] and states [batch,length,channels,state].
        Input injection is delta*B*x, not the exact zero-order-hold integral.
    """
    _validate_scan(x, delta, a, b, c, initial_state)
    if d is not None:
        _finite_shape(d, (x.shape[-1],), "D")
    state = (
        x.new_zeros(x.shape[0], x.shape[2], a.shape[1])
        if initial_state is None
        else initial_state
    )
    outputs, states = [], []
    for index in range(x.shape[1]):
        step = delta[:, index, :, None]
        retention = torch.exp(step * a)
        write = step * b[:, index, None, :] * x[:, index, :, None]
        state = retention * state + write
        readout = (state * c[:, index, None, :]).sum(-1)
        outputs.append(readout if d is None else readout + d * x[:, index])
        states.append(state)
    return torch.stack(outputs, 1), torch.stack(states, 1)

3. Compose chronological affine updates

Original paper

Enlarge figure

Math and tensor shapes

f_t(h)=a_t\odot h+w_t

(a_2,w_2)\circ(a_1,w_1)=(a_2\odot a_1,\ a_2\odot w_1+w_2)

Each round reads a snapshot of the previous round; updating in place would change the parallel dependency graph. This teaching Hillis–Steele scan has logarithmic depth but O(T log T) pair work and materializes intermediates. It establishes prefix agreement, not the work or memory efficiency of the production fused kernel.

Affine pairs: [b,T,E,N]; prefix states: [b,T,E,N]

Pinned official code

Python implementation

python

def affine_prefix(x, delta, a, b, initial_state=None):
    """Hillis-Steele scan; each round reads a snapshot, not its own writes."""
    _validate_scan(x, delta, a, b, b, initial_state)
    retention = torch.exp(delta[..., None] * a)
    write = delta[..., None] * b[:, :, None, :] * x[..., None]
    pairs = [(retention, write)]
    stride = 1
    while stride < x.shape[1]:
        previous_a, previous_b = pairs[-1]
        # later(earlier(h)) = later.a*earlier.a*h + later.a*earlier.b+later.b
        combined_a = previous_a[:, stride:] * previous_a[:, :-stride]
        combined_b = (
            previous_a[:, stride:] * previous_b[:, :-stride]
            + previous_b[:, stride:]
        )
        pairs.append(
            (
                torch.cat((previous_a[:, :stride], combined_a), 1),
                torch.cat((previous_b[:, :stride], combined_b), 1),
            )
        )
        stride *= 2
    final_a, final_b = pairs[-1]
    initial = (
        x.new_zeros(x.shape[0], x.shape[2], a.shape[1])
        if initial_state is None
        else initial_state
    )
    return final_a * initial[:, None] + final_b, pairs

4. Add causal local context before selecting coefficients

Original paper

Enlarge figure

Math and tensor shapes

x_{t,c}=\operatorname{SiLU}\!\left(b_c+\sum_{k=0}^{K-1}q_{c,k}v_{t-K+1+k,c}\right)

Left padding supplies zeros at the start. The convolution reads only the present and previous positions; channels interact in the surrounding linear projections. This short local convolution remains valid even though input-dependent SSM coefficients prevent a single fixed global SSM convolution kernel.

Projected branch [b,T,E] → causal depthwise Conv1d → [b,T,E]

Pinned official code

Python implementation

python

def causal_convolution(x, convolution):
    """Convolve [batch,length,channels] with left context only, then SiLU."""
    width = convolution.kernel_size[0]
    padded = F.pad(x.transpose(1, 2), (width - 1, 0))
    return F.silu(
        F.conv1d(
            padded, convolution.weight, convolution.bias, groups=x.shape[-1]
        )
    ).transpose(1, 2)

5. Put the selective operator inside the complete mixer

Original paper

Enlarge figure

Math and tensor shapes

[v,z]=W_{in}u,\quad x=\operatorname{SiLU}(\operatorname{CausalDWConv}(v))

o=W_{out}\!\left(y\odot\operatorname{SiLU}(z)\right)

u_{next}=u+\operatorname{Mixer}(\operatorname{RMSNorm}(u))

The scan already includes D*x. The z branch gates that scan result before output projection. The final equation describes an enclosing prenorm/residual wrapper, implemented separately from this forward method. It is a different addition from D*x inside the operator.

Tiny reference: [b,6,4] → two width-8 branches → [b,6,4]

Pinned official code

Python implementation

python

def forward(self, u, return_trace=False):
    """Trace the mixer; an enclosing block supplies prenorm and residual."""
    projected = self.in_proj(u)
    x, z = projected.chunk(2, -1)
    convolved = causal_convolution(x, self.conv)
    delta, a, b, c = self.coefficients(convolved)
    scanned, states = selective_scan(convolved, delta, a, b, c, self.D)
    gated = scanned * F.silu(z)
    output = self.out_proj(gated)
    if return_trace:
        return output, {
            "projected": projected,
            "x": x,
            "z": z,
            "convolved": convolved,
            "delta": delta,
            "A": a,
            "B": b,
            "C": c,
            "states": states,
            "scanned": scanned,
            "gated": gated,
            "output": output,
        }
    return output

6. Carry both buffers to the next token

Original paper

Enlarge figure

Math and tensor shapes

\mathrm{conv}: [b,E,K],\quad\mathrm{ssm}: [b,E,N]

\mathrm{bytes}=s\,b\,L_{layers}\,E\,(K+N)

Shift the convolution history, append the projected current input, convolve, produce coefficients and update the SSM. K counts the stored samples including the latest, matching the pinned allocation rather than assuming K−1 slots. s is bytes per element, with a common dtype assumed for this formula; weights and temporary tensors are excluded.

Tiny reference buffers: [b,8,3] and [b,8,3]

Pinned official code

Python implementation

python

def initial_cache(self, batch):
    """Store K projected samples, including the latest, as the source does."""
    reference = self.in_proj.weight
    return (
        reference.new_zeros(batch, self.channels, self.conv_width),
        reference.new_zeros(batch, self.channels, self.state_size),
    )

def step(self, token, cache):
    """Process one [batch,model_width] token with functional cache updates."""
    conv_state, ssm_state = cache
    x, z = self.in_proj(token).chunk(2, -1)
    conv_state = torch.cat((conv_state[..., 1:], x[..., None]), -1)
    x = F.silu(
        (conv_state * self.conv.weight[:, 0]).sum(-1) + self.conv.bias
    )
    delta, a, b, c = self.coefficients(x)
    ssm_state = (
        torch.exp(delta[..., None] * a) * ssm_state
        + delta[..., None] * b[:, None, :] * x[..., None]
    )
    readout = (ssm_state * c[:, None, :]).sum(-1) + self.D * x
    output = self.out_proj(readout * F.silu(z))
    return output, (conv_state, ssm_state)

The reference establishes these mechanics. Its random-weight forward/backward smoke output is not text generation, a model-quality evaluation or a throughput measurement. To assess the learned model, we need the paper’s trained-model experiments and their protocol.

Published results and the compression tradeoff

The following values are published Pile validation perplexities, copied from the original paper’s Table 1. Lower is better. The comparison uses the GPT-NeoX tokenizer and 300 billion training tokens; Mamba and Pythia use context length 2,048. This is a selected model pair, not a claim about every Transformer training recipe.

Source: Mamba v2, Table 1 and §4.2.2. Table 1 also reports downstream zero-shot tasks, which measure different behavior from validation perplexity. RWKV’s context length in that comparison is 1,024, so sharing the dataset and training-token count does not make every experimental setting identical.

The paper separately measures kernel efficiency and end-to-end generation. Those boundaries matter: its generation protocol uses a prompt of 2,048 tokens, generates 128 tokens and varies batch size, while the scan benchmark isolates a different workload. Neither the storage formula nor the browser animation predicts those runtimes. Paper §4.5 and Appendix E.5

Mamba’s central result is that input-dependent compression can support useful sequence modeling while preserving a recurrent decode state. Its limitations follow the same mechanism: state capacity is finite, repeated retention can attenuate older signals, and operation count alone does not determine latency. Long-context success on the paper’s synthetic, audio or genomic tasks does not establish lossless recall of arbitrary language histories.

This leaves a computational question. The scan’s arithmetic is mostly elementwise, while modern GPUs devote substantial capability to matrix multiplication. Can we keep a recurrent state and reorganize much more of the sequence computation into matrix products? Mamba-2 answers by restricting the transition structure, deriving a structured mixing matrix and computing it in chunks.

Key takeaways

• Selection makes delta, B and C depend on the input; the learned continuous A stays input-independent in the implementation shown here.
• The state update combines retained memory and a new write, then reads the updated state. Delta changes both terms and is not a probability.
• Affine composition permits parallel prefix computation during training. Fusion and memory movement determine whether that algebra executes efficiently.
• A complete Mamba block includes local convolution, projections, a multiplicative gate and an outside residual, beyond the selective recurrence.
• Both convolution history and SSM state persist during decoding. Their fixed size trades growing storage for compressed context.
• Authored arithmetic, untrained CPU traces and published model evaluations provide different kinds of evidence.

Related reading

• Mamba-2: structured state space duality: turn a restricted state update into a structured mixing matrix and a chunked algorithm.
• VMamba: follow four spatial traversal orders, aligned cross-merge and a hierarchical image backbone.
• Vision Mamba: follow forward and backward patch sequences through a bidirectional visual mixer.
• Convolution operation: distinguish the short local convolution from the selective running state.
• Linear attention approximations: compare compressed sequence summaries with other alternatives to full pairwise attention.
• Original Mamba paper, arXiv v2 and official v2.2.4 source: the full derivations, trained-model experiments and optimized implementations.