Mamba remembers by forgetting, read closely

Imagine reading a long receipt to remember one line marked important. Attention keeps every line available and looks back when asked. A recurrent model carries one running note instead. That note is cheaper, but every new line risks smudging the number you meant to keep.

Mamba's central move is to make the smudging decision depend on the current input. A useful token can overwrite the state; a distractor can barely touch it. This is selection: deciding what enters a compact memory and what should be forgotten.

A state-space model, or SSM, passes a hidden state h from one sequence position to the next. The state is a fixed-size summary of everything seen so far. In the scalar version we will rebuild, each input x produces:

hₜ = exp(−Δₜ) hₜ₋₁ + (1 − exp(−Δₜ)) xₜ

The timestep Δ controls memory. When it is near zero, exp(−Δ) is near one: keep the old state and nearly ignore the input. When it is large, retention falls toward zero: overwrite the state with the new input. Mamba predicts this control from the token itself, alongside input and output projections called B and C.

def selective_scan(inputs, deltas, state=0.0):
    out = []
    for x, delta in zip(inputs, deltas):
        retention = exp(-delta)
        state = retention * state + (1 - retention) * x
        out.append(state)
    return out

This is a one-dimensional zero-order-hold SSM with A = −1 and B = 1. The full Mamba block has many channels and state dimensions plus learned projections, convolution, normalization, and gating. The tiny recurrence preserves the paper's mechanism without pretending to be its language model.

The paper motivates Mamba with selective copying. The model sees marked values at unpredictable positions, surrounded by distractors, and must reproduce the marked values later. A time-invariant model applies the same update everywhere; it cannot know whether the current token deserves memory.

Our seeded version marks the first scalar as the target and appends Gaussian distractors. The fixed SSM uses Δ = 0.25 for every item. The selective SSM uses Δ = 6 for the target and Δ = 0.002 for noise. We average absolute recall error over 400 independently generated sequences at each gap.

The paper's language-model ablation tells the same story at scale. Perplexity measures how surprised a model is by held-out text; lower is better. Making any one parameter selective helps, but the three controls work best together.

Here is the part a skim misses. A conventional SSM uses the same dynamics at every position. That linear time-invariance means its entire recurrence can be rewritten as one convolution—a fixed filter sliding across the sequence. GPUs are excellent at that operation.

Once Δ, B, and C depend on each token, the filter changes at every position. The convolution equivalence disappears. Selection creates the content awareness, but it also removes the old fast training path. The paper therefore needs a second contribution: a hardware-aware selective scan.

Write one scalar update as the affine function h ↦ ah + b. Two consecutive updates compose into another affine function:

(a₂, b₂) ∘ (a₁, b₁) = (a₂a₁, a₂b₁ + b₂)

Composition is associative: grouping changes the work schedule, not the answer. A tree-shaped prefix scan can therefore combine distant chunks concurrently. The TypeScript core uses the simple sequential form for clarity; the three-language code below shows the associative operator that makes the hardware algorithm possible.

def compose(left, right):
    a1, b1 = left
    a2, b2 = right
    return a2 * a1, a2 * b1 + b2

During autoregressive generation, a Transformer stores keys and values for previous tokens in a KV cache. That cache grows with context length. A recurrent SSM carries forward a fixed state, so the next token costs constant memory with respect to prior sequence length.

The bargain has a price: every past detail must fit through that finite state. Attention delays compression by keeping the original token representations available. Mamba compresses immediately and learns what to retain. That is why selective copying is more than a toy—it isolates the central bet that learned filtering can preserve the useful history.

• Stress retrieval when many marked facts compete for a fixed state, rather than when one target only needs protection from noise.
• Compare wall-clock speed and peak memory across sequence lengths on several GPUs; fused kernels can move the crossover point substantially.
• Inspect which tokens cause large timesteps and state resets, then test whether those choices remain stable under paraphrase or tokenization changes.

Want the alternative computation in front of you? Use the Transformer page to inspect attention scores and the causal mask. Attention chooses which stored token to read; Mamba chooses which information is allowed into a compressed state.