Core ML

Mamba

Mamba makes state space models input-dependent: instead of using the same fixed transition for every token, it learns separate selection gates that let each token control how much it updates the state and how much of the output to reveal.

Why Does This Exist?

State space models achieve linear-time inference by compressing the past into a fixed-size state updated per token. The problem is their selection is "dumb": the matrices $A$ and $B$ that control how each input modifies the state are fixed — the same for every token, regardless of content. That's the same limitation RNNs have, and it's why information that entered the state many tokens ago tends to decay regardless of whether it was important.

Attention solves this by having every query look back at every key directly — expensive but precise. Gu and Dao (2023) asked: can you keep the linear-time budget of SSMs while adding content-aware selection? Mamba is their answer. The key change is to make the discretisation step size $\Delta$ , which controls how strongly a new input modifies the state, a function of the input itself. A "boring" token produces a small $\Delta$ — the state barely changes, old information is preserved. A "salient" token produces a large $\Delta$ — it rewrites the state strongly. The model learns which tokens matter.

This one change closes much of the quality gap between SSMs and attention, particularly on recall-intensive tasks, while preserving linear inference cost.

Think of It Like This

A highlighter you only press hard on important sentences

Imagine reading a document and updating your notes as you go. A vanilla SSM is like pressing the highlighter with exactly the same pressure on every sentence, so every sentence fades at the same rate. Mamba is like having a smart grip that tightens when the sentence is important — important sentences get pushed hard into the notes, unimportant ones barely leave a mark. The notebook (state) ends up with the salient sentences prominent and the filler already faded.

How It Actually Works

The selection mechanism

In a standard SSM, the step size $\Delta$ is a learned scalar — fixed across all inputs. Mamba makes $\Delta$ input-dependent:

\Delta_t = \text{softplus}(W_\Delta \, x_t + b_\Delta)

where $W_\Delta$ is a small learned linear projection and $x_t$ is the current input. Then:

\bar{A}_t = \exp(\Delta_t A), \quad \bar{B}_t \approx \Delta_t B

When $\Delta_t$ is large, $\bar{A}_t \approx 0$ (old state forgotten, new input dominates). When $\Delta_t$ is small, $\bar{A}_t \approx I$ (state preserved, new input barely absorbed). This is exactly attention's selectivity, but implemented through state transitions rather than a score matrix.

Why this breaks the parallel form

Standard SSMs can be parallelised with a fast convolution during training because the kernel is fixed — the same for every position. Mamba's input-dependent $\Delta$ means the kernel changes at every position, so the standard FFT trick doesn't apply. Gu and Dao wrote a hardware-efficient parallel scan algorithm that exploits the associativity of the state transitions, achieving $O(N \log N)$ on GPU with careful use of SRAM — a different but similarly IO-aware trick to FlashAttention.

What the Mamba block looks like

A Mamba block wraps the SSM with a gated MLP: the input is split into two branches, one goes through the SSM, the other through a pointwise gate (SiLU), and the two are multiplied. This is a gated linear unit pattern, similar to SwiGLU in transformers.

Show Me the Code

Implementing the selective recurrence with input-dependent $\Delta$ and comparing to a fixed- $\Delta$ SSM.

1import numpy as np2
3rng = np.random.default_rng(13)4n, d, d_in = 12, 4, 8   # seq len, state dim, input dim5
6A = -0.5 * np.ones(d)           # stable diagonal A (continuous)7B = rng.standard_normal((d, d_in)) * 0.18C = rng.standard_normal((d_in, d)) * 0.19W_delta = rng.standard_normal((d_in,)) * 0.01  # simplified scalar projection10x_seq = rng.standard_normal((n, d_in))11
12def softplus(x):13    return np.log1p(np.exp(x))14
15# Mamba: input-dependent delta16h = np.zeros(d)17y_mamba = np.zeros((n, d_in))18for t in range(n):19    delta = softplus(W_delta @ x_seq[t] + 0.5)   # varies per token20    A_bar = np.exp(delta * A)21    B_bar = (1 - A_bar) / (-A) * (B @ x_seq[t])  # simplified ZOH22    h = A_bar * h + B_bar23    y_mamba[t] = C @ h24
25# Fixed SSM: constant delta26delta_fixed = 0.327A_bar_f = np.exp(delta_fixed * A)28h_f = np.zeros(d)29y_fixed = np.zeros((n, d_in))30for t in range(n):31    B_bar_f = (1 - A_bar_f) / (-A) * (B @ x_seq[t])32    h_f = A_bar_f * h_f + B_bar_f33    y_fixed[t] = C @ h_f34
35# The variance of delta shows how much it varies with input36deltas = [softplus(W_delta @ x_seq[t] + 0.5) for t in range(n)]37print(f"Delta range: {min(deltas):.3f} – {max(deltas):.3f}")38# Varies meaningfully across tokens — the selection is working39print(f"Output MSE (Mamba vs fixed): {((y_mamba - y_fixed)**2).mean():.4f}")40# Non-zero — the two models differ, confirming selection changes the result

The varying delta is the model's content-aware gate — some tokens are absorbed strongly, others nearly pass through.

Watch Out For

Expecting Mamba to fully replace attention on all tasks

Mamba's selectivity is better than standard SSMs but still operates through a fixed-size compressed state. Tasks requiring precise recall of a single specific token from thousands of steps ago — needle-in-a-haystack retrieval, exact memorisation — can still exceed what the state can faithfully store. Attention handles these by direct lookup; Mamba must store the needle in the state, which is harder when the state has been modified by many subsequent tokens. For these tasks, hybrid architectures that pair Mamba with periodic attention layers tend to do better.

Assuming Mamba's parallel scan is equivalent to FFT convolution

Standard SSMs parallelize training with FFT convolution, $O(N \log N)$ . Mamba's input-dependent transitions break that — the kernel is different at every position. The parallel scan algorithm Mamba uses is also $O(N \log N)$ on a GPU but has different memory access patterns and implementation requirements. If you naively try to adapt standard SSM training code to Mamba's selective transitions, you won't get the parallel speedup.

The Quick Version

Mamba makes the SSM step size $\Delta$ a learned function of the input, so each token controls how strongly it modifies the state.
Large $\Delta$ : the old state is largely forgotten, the new input dominates. Small $\Delta$ : state is preserved, input barely registers.
This single change closes much of the quality gap between SSMs and attention, especially on recall-heavy tasks.
The input-dependent transitions break the FFT parallelism trick; Mamba replaces it with a hardware-efficient parallel scan.
For tasks that need precise long-range retrieval, hybrid attention+Mamba architectures outperform Mamba alone.