Mamba and Mamba-2: Selective State Spaces and Structured State Space Duality

A state space model (SSM) is a sequence model rooted in control theory: a hidden state evolves through time according to a linear differential equation, driven by the input. Discretize the equation, plug in trainable matrices, and you get a sequence layer that runs as a recurrence at inference, scales linearly in sequence length, and looks nothing like attention. The catch is that classical SSMs apply the same transformation to every token regardless of what the token contains — a property that makes them excellent at signal processing and hopeless at selective recall like “Harry … Harry Potter”.

This post derives Mamba and Mamba-2, the two papers that fixed this. The first paper makes the SSM parameters input-dependent so the model can decide per token whether to read or skip. The second reveals that this selective SSM is, under a natural restriction, the same function as a structured form of masked attention computed by a different algorithm. We will work through both papers using a 4-token running example, and arrive at the structured state space duality that ties SSMs and attention together. Concretely, we draw from:

1. Gu and Dao (2023), “Mamba: Linear-Time Sequence Modeling with Selective State Spaces”: identifies the fundamental limitation of prior SSMs (fixed dynamics cannot reason about content), introduces the selection mechanism that makes SSM parameters input-dependent, and proposes a hardware-aware architecture that matches Transformer quality for the first time.

2. Dao and Gu (2024), “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality”: reveals that SSMs and attention are not separate ideas but two algorithms for computing the same function on semiseparable matrices, uses this duality to design a faster algorithm (SSD) that leverages matrix multiplication units, and proposes the Mamba-2 architecture that is 2–8× faster than Mamba while improving quality.

The first paper says: give SSMs the ability to select. The second paper says: selection makes SSMs equivalent to a form of attention. Together, they close the loop between the two paradigms.

---

The Running Example

We use a single-channel input sequence of 4 tokens:

$$x = (x_1, x_2, x_3, x_4) = (1.0, \; 0.5, \; -0.3, \; 0.8)$$

This represents one channel ($D = 1$) of a sequence after the input projection. For the SSM, we use a state dimension of $N = 1$ (the simplest possible latent state — a single scalar). For the model-scale analysis, we use the same parameters as the series: $d_\text{model} = 512$, $L = 12$ layers.

For the selective SSM and the Mamba-2 duality sections, we extend to a multi-channel example with $D = 4$ and head dimension $P = 2$.

---

1. State Space Models

1.1 The continuous system

We will build the state space model one concept at a time. Each piece answers one question.

What is a “state”?

The state is a small running summary of everything the model has seen so far. We denote it $h(t) \in \mathbb{R}^N$, where $t$ is time and $N$ is how many numbers the summary contains. Imagine a bucket sitting under a tap. The tap drips water in (the input), the bucket has a small hole at the bottom that drips water out (forgetting), and at any moment the only thing that matters about the past is the current water level. That water level is the state. Everything the model “remembers” must fit inside it.

In our running example we will pick the smallest possible state: $N = 1$, a single scalar. So the bucket holds just one number.

How does the state change?

The state evolves continuously over time. We need an equation that says how it changes from one instant to the next. Two things can change the state: the state itself (the bucket leaks proportional to how full it is), and the current input (water arriving through the tap). The simplest equation that captures both is:

$$h'(t) = Ah(t) + Bx(t) \qquad \text{(state equation)}$$

Read this slowly. $h'(t)$ is the rate of change of the state — how fast the bucket level is rising or falling at this exact moment. $Ah(t)$ is the part of that change that depends on the current state itself (the leak). $Bx(t)$ is the part driven by the current input (the tap). Add the two contributions and you get the total rate of change.

There is a separate equation for what we read out of the state at each moment:

$$y(t) = Ch(t) \qquad \text{(output equation)}$$

This says the output $y(t)$ is just a linear measurement of the current state. We do not get to look directly at the input $x(t)$ when producing the output — everything must flow through the state.

What do $A$, $B$, and $C$ do?

Each matrix has one job. $A \in \mathbb{R}^{N \times N}$ is the state matrix: it controls how the state evolves on its own (how fast the bucket leaks, and in our scalar case, whether the leak is even leaking — if $A$ were positive instead of negative, the bucket would refill itself). $B \in \mathbb{R}^{N \times 1}$ is the input matrix: it controls how the new input enters the state (how wide the tap is). $C \in \mathbb{R}^{1 \times N}$ is the output matrix: it controls how we read out of the state (which dipstick we use to measure the level). Together, $A$, $B$, and $C$ fully specify the system’s behavior. This setup is the linear time-invariant (LTI) system from classical control theory (Kalman, 1960), and the entire space of possible state trajectories — what could ever happen to $h(t)$ — is called the state space.

What does “time-invariant” mean?

The matrices $A$, $B$, $C$ do not depend on $t$. The bucket has the same hole size and the same tap width at every moment. Pour the same water in at $t = 3$ or at $t = 300$ and the bucket level reacts identically. This is what time-invariant means: the rule that governs the system is the same at every instant.

This sounds innocent. It will turn out to be the entire problem we eventually need to fix in Section 2 — because in language modeling, we want the rule to depend on what token just arrived.

Plugging in the running example

For our running example with $N = 1$, we pick $A = -1$, $B = 1$, $C = 1$. The state equation becomes:

$$h'(t) = -h(t) + x(t)$$

This is the canonical leaky integrator: the state grows by whatever you put in (the $+x(t)$ term, which is $Bx(t)$ with $B=1$) and shrinks in proportion to how much is currently there (the $-h(t)$ term, which is $Ah(t)$ with $A=-1$). The decay rate is $|A| = 1$, meaning that if you stop pouring anything in, the bucket loses about 63% of its level per unit of time. This is the same dynamics as an RC circuit discharging through a resistor, or a hot cup of coffee cooling toward room temperature — equations from physics that share the same first-order structure.

1.2 Discretization

Why we need to discretize

The state equation we just wrote down is a continuous-time equation: it talks about $h'(t)$, the rate of change of the state at every real-numbered moment in time. But neural networks do not see continuous signals. They see a discrete sequence of tokens $x_1, x_2, x_3, \ldots$ arriving one at a time. We need a way to convert “the continuous bucket equation” into “a recurrence that takes one token in and produces one new state out.” That conversion is called discretization.

A discretization rule is a recipe with one job: turn the continuous parameters $(\Delta, A, B)$ into discrete parameters $(\bar{A}, \bar{B})$ such that the discrete recurrence $h_t = \bar{A} h_{t-1} + \bar{B} x_t$ produces the same trajectory at sample points $t = 0, \Delta, 2\Delta, \ldots$ that the continuous bucket would have produced.

What is $\Delta$?

$\Delta \in \mathbb{R}_{>0}$ is the step size: how much continuous time elapses between two adjacent tokens. It is a knob. Crank it small, and adjacent tokens are very close together in continuous time — the bucket barely changes between them. Crank it large, and adjacent tokens are far apart in continuous time — the bucket has long stretches to leak and fill. So $\Delta$ controls how aggressively each new token affects the state. We will see in Section 2.3 that this knob is exactly what Mamba turns into a learnable, input-dependent quantity.

What does “zero-order hold” mean?

To do the conversion, we have to make some assumption about what the input $x$ is doing between samples — between, say, $t$ and $t + \Delta$. We only have a sample at $t$ and a sample at $t + \Delta$. What happens in between? The simplest possible answer is: it stays at the value it had at $t$. The input is a flat plateau between samples. That assumption is called the zero-order hold (ZOH): we hold the input constant (“zero-order” means “constant,” as opposed to “first-order” which would be linear interpolation, etc.). Under this pretense, the continuous equation can be integrated exactly between samples, and we get the closed-form discrete update:

$$\bar{A} = \exp(\Delta A)$$ $$\bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I) \cdot \Delta B$$

Two formulas. The first says: in $\Delta$ seconds, the bucket leaks by a factor of $\exp(\Delta A)$ on its own (with $A < 0$, this is between 0 and 1). The second says: how much new input gets pumped into the bucket during that same interval, accounting for the fact that it leaks while it fills.

Where these formulas come from

The two boxed expressions for $\bar{A}$ and $\bar{B}$ are not assumed — they are derived. The continuous equation $h'(t) = Ah(t) + Bx(t)$ is a first-order linear ODE; solving it with the integrating factor method gives a closed-form expression for $h(t + \Delta)$ in terms of $h(t)$ and an integral of the input over $[t, t+\Delta]$. Under ZOH the input is constant on that interval, so it factors out of the integral, and what remains evaluates to $A^{-1}(\exp(\Delta A) - I)$ via a matrix-exponential integral identity. Reading off the coefficients of $h(t)$ and $x(t)$ in the resulting expression gives the two formulas above. The full step-by-step derivation — integrating factor, matrix exponential, the integral identity, and the ZOH substitution — is in Mathematical Prerequisites for Mamba, Sections 2 through 5.

Plugging in the running example

For $A = -1$, $B = 1$, $\Delta = 1.0$ (one second between samples):

$$\bar{A} = \exp(-1 \cdot 1.0) = \exp(-1) = 0.368$$

So one second of leak shrinks the bucket level to 36.8% of what it was — equivalently, 63.2% leaks out. That matches the leaky-integrator interpretation from Section 1.1.

$$\bar{B} = (-1)^{-1}(\exp(-1) - 1) \cdot 1.0 = (-1)(0.368 - 1) = (-1)(-0.632) = 0.632$$

So 63.2% of the new input gets written into the bucket per step. Notice that $\bar{A} + \bar{B} = 1$ in this scalar case — the fraction that leaks out and the fraction that gets written in sum to 1. The bucket update is a convex combination of the old level and the new input. We will see in Section 2 that this identity is exactly what makes the discretized SSM equivalent to a classical RNN gate.

1.3 The discrete recurrence

With the discretized parameters, the SSM becomes a linear recurrence:

$$h_t = \bar{A} h_{t-1} + \bar{B} x_t$$ $$y_t = C h_t$$

For our running example with $\bar{A} = 0.368$, $\bar{B} = 0.632$, $C = 1$, $h_0 = 0$:

Step $t = 1$: $x_1 = 1.0$.

$$h_1 = 0.368 \cdot 0 + 0.632 \cdot 1.0 = 0.632$$ $$y_1 = 1 \cdot 0.632 = 0.632$$

Step $t = 2$: $x_2 = 0.5$.

$$h_2 = 0.368 \cdot 0.632 + 0.632 \cdot 0.5 = 0.233 + 0.316 = 0.549$$ $$y_2 = 0.549$$

Step $t = 3$: $x_3 = -0.3$.

$$h_3 = 0.368 \cdot 0.549 + 0.632 \cdot (-0.3) = 0.202 - 0.190 = 0.012$$ $$y_3 = 0.012$$

Step $t = 4$: $x_4 = 0.8$.

$$h_4 = 0.368 \cdot 0.012 + 0.632 \cdot 0.8 = 0.004 + 0.506 = 0.510$$ $$y_4 = 0.510$$

The output is $y = (0.632, 0.549, 0.012, 0.510)$.

1.4 The convolutional form

Since the parameters $(\bar{A}, \bar{B}, C)$ are constant across time (the LTI property), the recurrence can be unrolled into a global convolution. Let us derive this.

Expanding the recurrence:

$$h_1 = \bar{B} x_1$$ $$h_2 = \bar{A} h_1 + \bar{B} x_2 = \bar{A}\bar{B} x_1 + \bar{B} x_2$$ $$h_3 = \bar{A} h_2 + \bar{B} x_3 = \bar{A}^2 \bar{B} x_1 + \bar{A}\bar{B} x_2 + \bar{B} x_3$$

The pattern: $h_t = \sum_{s=0}^{t-1} \bar{A}^s \bar{B} \, x_{t-s}$.

Multiplying by $C$:

$$y_t = C h_t = \sum_{s=0}^{t-1} C \bar{A}^s \bar{B} \, x_{t-s}$$

Define the SSM convolution kernel $\bar{K} = (C\bar{B}, \; C\bar{A}\bar{B}, \; C\bar{A}^2\bar{B}, \; \ldots)$. Then:

$$\boxed{y = x * \bar{K} \qquad \text{(causal convolution)}}$$

For our running example:

$$\bar{K}_0 = C\bar{B} = 1 \cdot 0.632 = 0.632$$ $$\bar{K}_1 = C\bar{A}\bar{B} = 1 \cdot 0.368 \cdot 0.632 = 0.233$$ $$\bar{K}_2 = C\bar{A}^2\bar{B} = 1 \cdot 0.368^2 \cdot 0.632 = 1 \cdot 0.135 \cdot 0.632 = 0.086$$ $$\bar{K}_3 = C\bar{A}^3\bar{B} = 1 \cdot 0.368^3 \cdot 0.632 = 1 \cdot 0.050 \cdot 0.632 = 0.031$$

Numerical check for $y_2$:

$$y_2 = \bar{K}_0 x_2 + \bar{K}_1 x_1 = 0.632 \cdot 0.5 + 0.233 \cdot 1.0 = 0.316 + 0.233 = 0.549 \quad \checkmark$$

Numerical check for $y_3$:

$$y_3 = \bar{K}_0 x_3 + \bar{K}_1 x_2 + \bar{K}_2 x_1 = 0.632 \cdot (-0.3) + 0.233 \cdot 0.5 + 0.086 \cdot 1.0$$ $$= -0.190 + 0.117 + 0.086 = 0.013$$

This matches $y_3 = 0.012$ from the recurrence (the difference is rounding in the kernel coefficients). $\checkmark$

1.5 The dual computation modes

This is the crucial property. The SSM has two equivalent computation modes:

1. Recurrent mode (equation in Section 1.3): processes one token at a time with $O(1)$ memory per step. Total cost: $O(TN)$ for $T$ tokens. Ideal for autoregressive inference.

2. Convolutional mode (equation in Section 1.4): processes the entire sequence at once via an FFT-based convolution. Total cost: $O(T \log T)$ for $T$ tokens. Ideal for parallel training.

Prior SSMs (S4, DSS, S4D, S5, H3, Hyena) exploit this duality: train with convolutions, infer with recurrence. The convolutional mode gives training parallelism. The recurrent mode gives constant-time inference per step.

1.6 The LTI property — and its limitation

The duality between recurrence and convolution only exists because $(\bar{A}, \bar{B}, C)$ are constant across time. This is the linear time-invariance (LTI) property. The bucket has the same hole and the same tap width at every step, regardless of which token just walked in the door.

LTI is a double-edged sword. On one hand it enables convolutions, because a single fixed kernel can be applied everywhere. On the other hand it prevents content-based reasoning — the model cannot decide to pay attention to one token and ignore another based on what those tokens actually contain. To see why this matters concretely: imagine a sequence “Harry … (lots of irrelevant words) … Harry Potter,” and you want to predict “Potter” after seeing the second “Harry.” The model has to recognize that the current token matches a token it saw earlier and route information based on that match. But an LTI SSM applies the same $\bar{A}, \bar{B}$ at every step — it cannot make any decision conditional on what the input looks like. The dynamics are frozen. Mamba’s whole job is to unfreeze them.

---

2. Why Selection Matters

2.1 The failure mode

Gu and Dao (2023) identify two tasks that reveal the LTI limitation:

The Selective Copying task modifies the standard Copying task by randomizing the spacing between tokens that need to be memorized. The standard Copying task has constant spacing, so a fixed convolution kernel can solve it by simply counting positions. The Selective Copying task has random spacing, so the model must look at the content of each token to decide whether to memorize it.

The Induction Heads task requires the model to perform associative recall: given a pattern like “Harry … Harry Potter”, the model must predict “Potter” when it sees the second “Harry”. This requires recognizing that the current token matches a previously seen token — a content-dependent operation.

LTI SSMs fail on both tasks. From the recurrent view, the $(\bar{A}, \bar{B})$ transitions are constant, so the model cannot selectively focus on or ignore tokens based on their content. From the convolutional view, a fixed convolution kernel is inherently position-aware but not content-aware — it cannot vary the spacing dynamically.

2.2 What selection means

The solution is to make the SSM parameters functions of the input. Instead of fixed $(\Delta, B, C)$, the selective SSM uses:

$$B_t = s_B(x_t) = \text{Linear}_N(x_t)$$ $$C_t = s_C(x_t) = \text{Linear}_N(x_t)$$ $$\Delta_t = \tau_\Delta(\text{Parameter} + s_\Delta(x_t)) = \text{softplus}(\text{Parameter} + \text{Linear}_1(x_t))$$

where $\text{Linear}_d$ denotes a learned linear projection to dimension $d$, and $\tau_\Delta = \text{softplus}$ ensures $\Delta_t > 0$.

The softplus function is $\text{softplus}(x) = \log(1 + \exp(x))$, a smooth approximation to $\text{ReLU}$. It is always positive: $\text{softplus}(x) > 0$ for all $x$.

The matrix $A$ remains fixed (not input-dependent). Gu and Dao hypothesize that making $A$ selective in addition to $\Delta$ would have similar performance, since $\Delta$ already controls $A$ through the discretization $\bar{A}_t = \exp(\Delta_t A)$.

The critical consequence: once $(\Delta, B, C)$ vary with time, the parameters $(\bar{A}_t, \bar{B}_t)$ are no longer constant. The LTI property breaks. The convolution kernel $\bar{K}$ is no longer well-defined (it would need to be different at every position). The convolutional computation mode is lost.

This is the fundamental tradeoff. Selection gives the model content-dependent dynamics. But it removes the fast convolutional training path. The model must be computed recurrently — or with a new algorithm.

2.3 Interpretation of $\Delta$

Of the three input-dependent parameters $(\Delta, B, C)$, the most important is $\Delta$ — and the easiest way to see why is to think of it as a knob between two extremes. $\Delta$ controls the balance between focusing on the current input $x_t$ and persisting the state $h_{t-1}$. Mechanistically:

• A large $\Delta_t$ means $\bar{A}_t = \exp(\Delta_t A) \to 0$ (since $A$ has negative entries) and $\bar{B}_t \to$ large. The state is reset, and the current input $x_t$ is written strongly. The system is “selecting” $x_t$.

• A small $\Delta_t$ means $\bar{A}_t \to 1$ and $\bar{B}_t \to 0$. The state is preserved, and the current input is ignored. The system is “skipping” $x_t$.

This is exactly the behavior needed for Selective Copying: the model should produce large $\Delta_t$ for content tokens and small $\Delta_t$ for noise tokens.

2.4 Numerical example: selective vs. non-selective

Let us trace the selective SSM for our running example. We keep $A = -1$, $C = 1$, but now $\Delta_t$ varies.

Suppose the selection mechanism produces:

$$\Delta_1 = 2.0, \quad \Delta_2 = 0.1, \quad \Delta_3 = 0.1, \quad \Delta_4 = 2.0$$

This says: focus on tokens 1 and 4, ignore tokens 2 and 3.

Step $t = 1$: $\Delta_1 = 2.0$.

$$\bar{A}_1 = \exp(-2.0) = 0.135, \quad \bar{B}_1 = (-1)^{-1}(\exp(-2.0) - 1) \cdot 2.0 = (-1)(-0.865) \cdot 2.0 = 1.729$$

Wait — let us be more careful with the ZOH formula. For the scalar case with $A = -1$:

$$\bar{B}_t = (\Delta_t A)^{-1}(\exp(\Delta_t A) - 1) \cdot \Delta_t B = \frac{\exp(-\Delta_t) - 1}{-\Delta_t} \cdot \Delta_t \cdot 1 = 1 - \exp(-\Delta_t)$$

This uses the simplification: $\frac{\exp(-\Delta_t) - 1}{-\Delta_t} \cdot \Delta_t = -((\exp(-\Delta_t) - 1)) = 1 - \exp(-\Delta_t)$. The $\Delta_t$ in the numerator and denominator cancel (by algebraic cancellation).

So:

$$\bar{A}_t = \exp(-\Delta_t), \qquad \bar{B}_t = 1 - \exp(-\Delta_t)$$

Note that $\bar{A}_t + \bar{B}_t = 1$. This is a convex combination of the previous state and the current input. This identity holds specifically for the scalar case $A = -1$, $B = 1$ with ZOH discretization.

Step $t = 1$: $\Delta_1 = 2.0$.

$$\bar{A}_1 = \exp(-2.0) = 0.135, \quad \bar{B}_1 = 1 - 0.135 = 0.865$$ $$h_1 = 0.135 \cdot 0 + 0.865 \cdot 1.0 = 0.865, \quad y_1 = 0.865$$

Strong focus on $x_1 = 1.0$: the state captures 86.5% of the input.

Step $t = 2$: $\Delta_2 = 0.1$.

$$\bar{A}_2 = \exp(-0.1) = 0.905, \quad \bar{B}_2 = 1 - 0.905 = 0.095$$ $$h_2 = 0.905 \cdot 0.865 + 0.095 \cdot 0.5 = 0.783 + 0.048 = 0.831, \quad y_2 = 0.831$$

The state barely changed: 90.5% of the previous state is retained, and only 9.5% of the new input enters. Token 2 is effectively ignored.

Step $t = 3$: $\Delta_3 = 0.1$.

$$\bar{A}_3 = 0.905, \quad \bar{B}_3 = 0.095$$ $$h_3 = 0.905 \cdot 0.831 + 0.095 \cdot (-0.3) = 0.752 - 0.029 = 0.723, \quad y_3 = 0.723$$

Again, the state is mostly preserved. Token 3 is ignored.

Step $t = 4$: $\Delta_4 = 2.0$.

$$\bar{A}_4 = 0.135, \quad \bar{B}_4 = 0.865$$ $$h_4 = 0.135 \cdot 0.723 + 0.865 \cdot 0.8 = 0.098 + 0.692 = 0.790, \quad y_4 = 0.790$$

Strong focus on $x_4 = 0.8$: the state is mostly overwritten.

Compare the outputs:

• Non-selective (Section 1.3): $y = (0.632, 0.549, 0.012, 0.510)$
• Selective: $y = (0.865, 0.831, 0.723, 0.790)$

The selective model retains the signal from token 1 through tokens 2 and 3 (where $y$ stays near 0.8), then switches to capture token 4. The non-selective model treats all tokens equally, leading to the state collapsing near zero at token 3 (because $x_3 = -0.3$ partially cancels the accumulated positive state).

---

3. The Connection to Gating

3.1 Theorem 1: selective SSMs are gated RNNs

Gu and Dao (2023) prove that the selective SSM, under specific parameter choices, reduces exactly to a classical gated RNN. This is Theorem 1 of the Mamba paper.

Theorem 1. When $N = 1$, $A = -1$, $B = 1$, $s_\Delta = \text{Linear}(x)$, and $\tau_\Delta = \text{softplus}$, the selective SSM recurrence takes the form:

$$g_t = \sigma(\text{Linear}(x_t))$$ $$h_t = (1 - g_t) h_{t-1} + g_t x_t$$

where $\sigma$ is the sigmoid function $\sigma(z) = 1 / (1 + \exp(-z))$.

3.2 Proof

The proof is in Appendix C of the Mamba paper. Let us re-derive it step by step.

The continuous system with $N = 1$, $A = -1$, $B = 1$ is:

$$h'(t) = -h(t) + x(t)$$

The discretization step size is:

$$\Delta_t = \text{softplus}(\text{Parameter} + \text{Linear}(x_t))$$

We observe that the Parameter can be absorbed as a bias term in the linear projection. So we write $\Delta_t = \text{softplus}(\text{Linear}(x_t))$ where $\text{Linear}$ includes the bias.

Applying ZOH with $A = -1$:

$$\bar{A}_t = \exp(\Delta_t \cdot (-1)) = \exp(-\Delta_t) = \frac{1}{\exp(\Delta_t)}$$

Now we use the identity $\text{softplus}(z) = \log(1 + \exp(z))$, so $\Delta_t = \log(1 + \exp(\text{Linear}(x_t)))$. Exponentiating:

$$\exp(\Delta_t) = 1 + \exp(\text{Linear}(x_t))$$

Therefore:

$$\bar{A}_t = \frac{1}{1 + \exp(\text{Linear}(x_t))} = \sigma(-\text{Linear}(x_t))$$

By the sigmoid reflection identity $\sigma(-z) = 1 - \sigma(z)$:

$$\bar{A}_t = 1 - \sigma(\text{Linear}(x_t))$$

Define $g_t = \sigma(\text{Linear}(x_t))$. Then $\bar{A}_t = 1 - g_t$.

For $\bar{B}_t$, we showed in Section 2.4 that $\bar{B}_t = 1 - \exp(-\Delta_t) = 1 - \bar{A}_t = g_t$.

The recurrence becomes:

$$h_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t = (1 - g_t) h_{t-1} + g_t x_t$$ $$\boxed{h_t = (1 - g_t) h_{t-1} + g_t x_t}$$

This is exactly the gated recurrence. $\square$

3.3 Numerical verification

For $x_1 = 1.0$ with $\text{Linear}(x) = 2x - 1$ (a concrete linear projection):

$$g_1 = \sigma(2 \cdot 1.0 - 1) = \sigma(1.0) = \frac{1}{1 + \exp(-1)} = \frac{1}{1 + 0.368} = \frac{1}{1.368} = 0.731$$ $$h_1 = (1 - 0.731) \cdot 0 + 0.731 \cdot 1.0 = 0.731$$

Cross-check via $\Delta_1 = \text{softplus}(1.0) = \log(1 + \exp(1.0)) = \log(1 + 2.718) = \log(3.718) = 1.313$:

$$\bar{A}_1 = \exp(-1.313) = 0.269, \quad \bar{B}_1 = 1 - 0.269 = 0.731$$ $$h_1 = 0.269 \cdot 0 + 0.731 \cdot 1.0 = 0.731 \quad \checkmark$$

For $x_3 = -0.3$:

$$g_3 = \sigma(2 \cdot (-0.3) - 1) = \sigma(-1.6) = \frac{1}{1 + \exp(1.6)} = \frac{1}{1 + 4.953} = \frac{1}{5.953} = 0.168$$

The gate is small ($0.168$), so the model retains most of the state and largely ignores $x_3 = -0.3$. This is the selection mechanism in action.

3.4 Interpretation

The connection to gating is not merely formal. It means that discretization of SSMs is the principled foundation of heuristic gating mechanisms. The gate $g_t = \sigma(\text{Linear}(x_t))$ in an LSTM or GRU was introduced as a heuristic to control information flow. The SSM perspective derives the same gate from first principles: start with a continuous dynamical system, discretize with ZOH, and the gate emerges naturally from the interaction between $\Delta$ and $A$.

This also explains why making $\Delta$ input-dependent is the most important selective parameter. Table 7 of the Mamba paper ablates the three selective parameters ($\Delta$, $B$, $C$). Making $\Delta$ alone selective reduces perplexity from 10.93 (no selection) to 10.15. Making $B$ or $C$ alone selective gives smaller improvements (10.93 → 10.15 for $\Delta$ vs 10.93 → 9.98 for $\Delta + B$ vs 10.93 → 9.81 for all three). $\Delta$ is the most important because it directly controls the gate — it determines whether to read or skip each token.

---

4. The Mamba Architecture

4.1 The architecture

Prior SSM architectures (H3, Hyena) interleave an SSM layer with an MLP block, following the Transformer’s pattern of alternating attention and MLP. Mamba simplifies this by merging the two blocks into one.

The Mamba block consists of:

1. Input projection: $x \in \mathbb{R}^D \to (x', z) \in \mathbb{R}^{ED} \times \mathbb{R}^{ED}$, where $E$ is the expansion factor (typically $E = 2$).

2. Short convolution: a 1D depthwise convolution with kernel size $d$ (typically $d = 4$) applied to $x'$. This provides local context before the SSM.

3. SSM: the selective state space model applied to the convolved $x'$. The parameters $(\Delta, B, C)$ are computed from the post-convolution activation.

4. Gating: the SSM output is multiplied element-wise by $\sigma(z)$, where $\sigma$ is the SiLU (Sigmoid Linear Unit) activation $\sigma(z) = z \cdot \text{sigmoid}(z)$.

5. Output projection: the gated result is projected back to $\mathbb{R}^D$.

The SiLU gating makes the Mamba block analogous to a SwiGLU MLP (Shazeer, 2020) — a variant of the standard two-layer MLP in which one of the linear branches is multiplied element-wise by a Swish-activated copy of the input, used in PaLM and LLaMA. Compared to the MLP block, Mamba simply adds a convolution and SSM to the main branch.

4.2 Parameter count

For each Mamba block with model dimension $D$ and expansion factor $E$:

• Input projections: $2ED^2$ (for $x'$ and $z$)
• Output projection: $ED^2$
• Total from projections: $3ED^2$

With $E = 2$, this is $6D^2$ per block. The SSM parameters ($A$, projections for $\Delta$, $B$, $C$) are much smaller in comparison. Two Mamba blocks (stacked homogeneously) match the $12D^2$ parameters of a Transformer layer (one attention + one MLP):

$$\text{Transformer layer}: \underbrace{4D^2}_\text{Q,K,V,O projections} + \underbrace{8D^2}_\text{MLP (up + down)} = 12D^2$$ $$\text{Two Mamba blocks}: 2 \times 6D^2 = 12D^2$$

4.3 Key results

Mamba achieves several firsts:

1. Selective Copying: Mamba solves the task with 99.8% accuracy, compared to 97.0% for S4 (no gate) and 18.3% for S4 without the selection mechanism (Table 1).

2. Induction Heads: Mamba extrapolates perfectly to sequences 4000× longer than training length ($2^8 = 256$ training → $2^{20} = 1{,}048{,}576$ test). No other method exceeds 2× extrapolation (Table 2).

3. Language modeling: Mamba is the first linear-time model to match the quality of a strong Transformer++ recipe (PaLM/LLaMA-style) on scaling laws from 125M to 1.3B parameters (Figure 4). Mamba-3B matches Transformers at twice the size on downstream tasks.

4. Inference throughput: Mamba achieves 4–5× higher generation throughput than a Transformer of similar size, because it does not require a KV cache that grows with sequence length.

---

5. The Hardware-Aware Algorithm

5.1 The problem

The selective SSM loses the convolutional computation mode. The naive recurrence requires materializing the expanded state $h \in \mathbb{R}^{B \times L \times D \times N}$ in GPU HBM (high-bandwidth memory), which is prohibitively large. For batch size $B = 16$, sequence length $L = 2048$, $D = 2048$, $N = 16$: the state requires $16 \times 2048 \times 2048 \times 16 \times 2 = 2$ GB in fp16.

5.2 The solution: kernel fusion

The key insight is that the SSM can be computed entirely in fast SRAM (on-chip memory), without materializing the full state in HBM.

1. Load the SSM parameters $(\Delta, A, B, C)$ from HBM to SRAM. Size: $O(BLD + DN)$.
2. Compute discretization ($\bar{A}_t, \bar{B}_t$) in SRAM.
3. Perform the selective scan (recurrence) in SRAM.
4. Multiply by $C$ and write the output $(B, L, D)$ back to HBM.

The intermediate states of size $(B, L, D, N)$ never leave SRAM. This reduces memory IOs by a factor of $O(N)$ (the state dimension), which in practice gives 20–40× speedup over a naive implementation.

5.3 Parallel scan

Despite being sequential in nature, the recurrence can be parallelized with a parallel associative scan (Blelloch, 1990). The key observation: the recurrence $h_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t$ can be written as an associative binary operation on pairs $(\bar{A}_t, \bar{B}_t x_t)$. The scan computes all prefix products in $O(\log L)$ parallel steps, using $O(L)$ work.

5.4 Recomputation

To reduce memory during training, Mamba does not save intermediate states for backpropagation. Instead, it recomputes them in the backward pass by reloading the inputs from HBM and re-running the scan in SRAM. This is the same technique as gradient checkpointing in Transformers (e.g., FlashAttention). The result: the selective SSM layer uses the same activation memory as a FlashAttention layer.

---

6. SSMs Are Structured Matrices

We now turn to the second paper (Dao and Gu, 2024), which reveals a deep connection between SSMs and attention. The starting point is a simple observation: every SSM can be written as a matrix multiplication.

6.1 The matrix transformation form

Recall the SSM recurrence:

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

Unrolling this from $h_0 = 0$ — multiply out $h_1, h_2, h_3, \ldots$ in turn and collect the contribution of each input $x_s$ — produces a closed-form expression for $h_t$ as a sum over the inputs, each weighted by a cumulative product of the $A$ matrices that came after it. The full unrolling is derived in Mathematical Prerequisites for Mamba, Section 6. The result is

$$h_t = \sum_{s=0}^{t} A_{t:s}^\times B_s x_s$$

where $A_{t:s}^\times = A_t A_{t-1} \cdots A_{s+1}$ denotes the cumulative product of the $A$ matrices from time $s+1$ to $t$, with the convention $A_{t:t}^\times = I$.

Multiplying by $C_t^\top$:

$$y_t = C_t^\top h_t = \sum_{s=0}^{t} C_t^\top A_{t:s}^\times B_s x_s$$

This is a matrix-vector product $y = Mx$ where:

$$\boxed{M_{ts} = C_t^\top A_{t:s}^\times B_s \qquad \text{for } t \geq s}$$

and $M_{ts} = 0$ for $t < s$ (causality). The matrix $M \in \mathbb{R}^{T \times T}$ is lower-triangular.

6.2 Semiseparable matrices

Dao and Gu (2024) identify the matrix $M$ as belonging to a well-studied class called semiseparable matrices.

Definition 3.1. A lower-triangular matrix $M$ is N-semiseparable if every submatrix contained in the lower-triangular portion has rank at most N.

The SSM matrix $M$ satisfies this with $N$ equal to the state dimension: the formula $M_{ts} = C_t^\top A_{t:s}^\times B_s$ factors any lower-triangular submatrix as an outer product of a chain of $C$‘s and a chain of $B$‘s, capped at rank $N$. This is Theorem 3.5 of the paper: the SSM transformation $y = \text{SSM}(A, B, C)(x)$ is identical to matrix multiplication by an N-semiseparable matrix $M = \text{SSS}(A, B, C)$.

6.3 The scalar case: 1-semiseparable matrices

The most important special case is when $A_t$ is a scalar times the identity: $A_t = a_t I$ for some scalar $a_t \in [0, 1]$. The cumulative product collapses to a scalar, $A_{t:s}^\times = a_{t:s}^\times \cdot I$ where $a_{t:s}^\times = a_t a_{t-1} \cdots a_{s+1}$, and the matrix factors as

$$M_{ts} = a_{t:s}^\times \cdot (C_t^\top B_s), \qquad M = L \circ (CB^\top)$$

where $L_{ts} = a_{t:s}^\times$ is a 1-semiseparable matrix (also called a 1-SS matrix) and $\circ$ denotes the Hadamard (element-wise) product. The 1-SS structure of $L$ — a lower-triangular matrix whose every entry below the diagonal is a cumulative product, equivalently whose every lower-triangular submatrix has rank at most 1 — is derived directly from the unrolling in Mathematical Prerequisites for Mamba, Section 7, where the explicit rank-1 factorization $L_{ts} = c_t / c_s$ is built up from first principles.

Concretely, $L$ has the form:

$$L = \text{1SS}(a) = \begin{pmatrix} 1 \\ a_1 & 1 \\ a_2 a_1 & a_2 & 1 \\ a_3 a_2 a_1 & a_3 a_2 & a_3 & 1 \end{pmatrix}$$

Each entry is a cumulative product of consecutive $a_t$ values. The diagonal is all 1’s. This is the structured mask — it replaces the causal mask of standard attention.

6.4 Numerical example

Let us trace through a 4-token example with scalar $A$. We use $a = (a_1, a_2, a_3, a_4) = (0.9, 0.8, 0.5, 0.7)$ and:

$$B = \begin{pmatrix} 1.0 \\ 0.5 \\ -0.3 \\ 0.8 \end{pmatrix}, \quad C = \begin{pmatrix} 0.5 \\ 1.0 \\ 0.7 \\ 0.3 \end{pmatrix}, \quad X = \begin{pmatrix} 1.0 \\ 2.0 \\ 1.5 \\ 0.5 \end{pmatrix}$$

(Here $B$, $C$, $X$ are all 1-dimensional per token since $N = P = 1$.)

The 1-SS mask $L$:

$$L = \begin{pmatrix} 1 & 0 & 0 & 0 \\ 0.8 & 1 & 0 & 0 \\ 0.8 \cdot 0.5 & 0.5 & 1 & 0 \\ 0.8 \cdot 0.5 \cdot 0.7 & 0.5 \cdot 0.7 & 0.7 & 1 \end{pmatrix} = \begin{pmatrix} 1 & 0 & 0 & 0 \\ 0.8 & 1 & 0 & 0 \\ 0.4 & 0.5 & 1 & 0 \\ 0.28 & 0.35 & 0.7 & 1 \end{pmatrix}$$

Note: $a_1$ is not used in $L$ because $L_{ts}$ involves the product $a_{s+1} \cdots a_t$, and the first row only has $L_{11} = 1$.

The Gram matrix $G = CB^\top$:

$$G_{ts} = C_t \cdot B_s = \begin{pmatrix} 0.5 \cdot 1.0 & 0.5 \cdot 0.5 & 0.5 \cdot (-0.3) & 0.5 \cdot 0.8 \\ 1.0 \cdot 1.0 & 1.0 \cdot 0.5 & 1.0 \cdot (-0.3) & 1.0 \cdot 0.8 \\ 0.7 \cdot 1.0 & 0.7 \cdot 0.5 & 0.7 \cdot (-0.3) & 0.7 \cdot 0.8 \\ 0.3 \cdot 1.0 & 0.3 \cdot 0.5 & 0.3 \cdot (-0.3) & 0.3 \cdot 0.8 \end{pmatrix}$$ $$= \begin{pmatrix} 0.5 & 0.25 & -0.15 & 0.4 \\ 1.0 & 0.5 & -0.3 & 0.8 \\ 0.7 & 0.35 & -0.21 & 0.56 \\ 0.3 & 0.15 & -0.09 & 0.24 \end{pmatrix}$$

The full matrix $M = L \circ G$ (element-wise product, lower-triangular):

$$M = \begin{pmatrix} 0.5 & 0 & 0 & 0 \\ 0.8 & 0.5 & 0 & 0 \\ 0.28 & 0.175 & -0.21 & 0 \\ 0.084 & 0.0525 & -0.063 & 0.24 \end{pmatrix}$$

Output $Y = MX$:

$$y_1 = 0.5 \cdot 1.0 = 0.5$$ $$y_2 = 0.8 \cdot 1.0 + 0.5 \cdot 2.0 = 0.8 + 1.0 = 1.8$$ $$y_3 = 0.28 \cdot 1.0 + 0.175 \cdot 2.0 + (-0.21) \cdot 1.5 = 0.28 + 0.35 - 0.315 = 0.315$$ $$y_4 = 0.084 \cdot 1.0 + 0.0525 \cdot 2.0 + (-0.063) \cdot 1.5 + 0.24 \cdot 0.5 = 0.084 + 0.105 - 0.0945 + 0.12 = 0.2145$$

We can verify $y_2$ via the recurrence. With scalar $A$, the recurrence is:

$$h_t = a_t h_{t-1} + B_t x_t, \quad y_t = C_t h_t$$ $$h_1 = a_1 \cdot 0 + B_1 x_1 = 1.0 \cdot 1.0 = 1.0, \quad y_1 = C_1 h_1 = 0.5 \cdot 1.0 = 0.5 \quad \checkmark$$ $$h_2 = a_2 h_1 + B_2 x_2 = 0.8 \cdot 1.0 + 0.5 \cdot 2.0 = 0.8 + 1.0 = 1.8, \quad y_2 = C_2 h_2 = 1.0 \cdot 1.8 = 1.8 \quad \checkmark$$

The matrix form and recurrence produce the same output.

6.5 Interpretation

The matrix $M$ encodes the full input-output map of the SSM. The recurrent mode computes $y = Mx$ by exploiting the sequential structure (each row depends on the previous state). The quadratic mode computes $y = Mx$ by materializing $M$ and doing direct matrix multiplication. These are two algorithms for the same computation — one is $O(TN)$ time, the other is $O(T^2N)$ time but more hardware-friendly because it uses matrix multiplications.

This is the core insight of the Mamba-2 paper: different methods of computing SSMs can be reframed as different algorithms for multiplying by semiseparable matrices.

---

7. Structured State Space Duality

7.1 From SSMs to attention

The matrix form $M = L \circ (CB^\top)$ with $L = \text{1SS}(a)$ looks strikingly similar to masked attention:

$$Y = (L \circ QK^\top) \cdot V$$

In standard causal attention, $L$ is the lower-triangular matrix of all 1’s (the causal mask), $Q$ and $K$ are queries and keys, and $V$ is the value matrix. In the SSM, $L$ is the 1-semiseparable mask of cumulative decay products, $C$ plays the role of queries, $B$ plays the role of keys, and $X$ plays the role of values.

The correspondence is exact:

SSM	Attention
$C$ (output matrix)	$Q$ (queries)
$B$ (input matrix)	$K$ (keys)
$X$ (input sequence)	$V$ (values)
$a_{t:s}^\times$ (cumulative product)	$L_{ts}$ (mask entry)
$N$ (state dimension)	$N$ (feature dimension)

7.2 The duality

Dao and Gu (2024) make this precise with structured state space duality (SSD).

State space models are usually defined through a recurrence (Definition 2.2 of the paper) and computed with a linear-time algorithm (the scan). Attention is usually defined through pairwise comparisons (equation 9) and computed with a quadratic-time algorithm (materializing $QK^\top$).

But both have dual forms:

• An SSM can be computed quadratically by materializing $M = \text{SSS}(A, B, C)$ and multiplying $Y = MX$. This is the quadratic (attention-like) mode.

• Attention can be computed linearly by using the cumsum trick from the linear attention framework (Katharopoulos et al., 2020). This is the linear (recurrent) mode.

The duality says: for scalar-identity $A$ matrices, these are the same function computed by different algorithms. The SSM recurrence is the linear mode. Materializing $M$ is the quadratic mode. They produce identical outputs.

7.3 The SSD layer

The state space dual (SSD) layer is the specific SSM that the duality applies to. Compared to Mamba’s selective SSM (S6), SSD makes two simplifications:

1. $A$ is restricted from diagonal to scalar times identity: $A_t = a_t I$. Each $a_t$ is a single scalar shared across all state dimensions.

2. The head dimension $P$ is increased from $P = 1$ (in Mamba) to $P = 64$ or $128$ (matching Transformer conventions).

The first restriction slightly decreases expressivity but enables the quadratic mode. The second compensates by using attention-like multi-head structure.

7.4 Why scalar $A$ matters

With diagonal $A_t = \text{diag}(a_t^{(1)}, \ldots, a_t^{(N)})$, each state dimension has an independent decay rate. The matrix $M$ becomes:

$$M_{ts} = C_t^\top \text{diag}(a_{t:s}^{(1)\times}, \ldots, a_{t:s}^{(N)\times}) B_s$$

This involves $N$ different 1-SS masks, one per state dimension. The quadratic mode requires materializing all $N$ masks and their Hadamard products, which is expensive.

With scalar $A_t = a_t I$, all state dimensions share the same decay: $a_{t:s}^{(i)\times} = a_{t:s}^\times$ for all $i$. The matrix simplifies to $M = L \circ (CB^\top)$ with a single mask $L$, and the quadratic mode becomes a single masked matrix multiplication — just like attention.

---

8. The SSD Algorithm

8.1 The idea

The linear (recurrent) mode takes $O(TN)$ time. The quadratic (attention-like) mode takes $O(T^2N)$ time. Neither is optimal in practice:

• The recurrent mode is sequential and cannot exploit matrix multiplication units (tensor cores on GPUs).
• The quadratic mode is parallelizable but scales poorly with sequence length.

The SSD algorithm combines both: split the sequence into chunks, compute within each chunk quadratically (using matmul), and connect chunks recurrently (using a scan). This is a block decomposition of the semiseparable matrix $M$.

8.2 Block decomposition

Partition the $T$-length sequence into $T/Q$ chunks of size $Q$. The matrix $M$ decomposes into blocks:

$$M = \begin{pmatrix} M^{(0,0)} & & \\ M^{(1,0)} & M^{(1,1)} & \\ \vdots & & \ddots \end{pmatrix}$$

The diagonal blocks $M^{(j,j)}$ represent intra-chunk interactions. They are small ($Q \times Q$) and can be computed quadratically using matrix multiplication.

The off-diagonal blocks $M^{(j,i)}$ for $j > i$ represent inter-chunk interactions. By the semiseparable property, these blocks are low-rank (rank at most $N$). They factor as:

$$M^{(j,i)} = \underbrace{C_\text{block}}_\text{left factor} \cdot \underbrace{A_\text{chain}}_\text{center factor} \cdot \underbrace{B_\text{block}^\top}_\text{right factor}$$

The center factors are connected by a scalar recurrence (a 1-SS multiplication of length $T/Q$), which is $Q$ times shorter than the original sequence.

8.3 The four steps

The SSD algorithm has four steps:

Step 1: Diagonal blocks (intra-chunk). For each chunk $j$, compute the output from tokens within the chunk using the quadratic form:

$$Y_\text{diag}^{(j)} = (L^{(j)} \circ C^{(j)} B^{(j)\top}) X^{(j)}$$

This is a batched matrix multiplication. Cost: $\text{BMM}(T/Q, Q, Q, P)$.

Step 2: Right factors (chunk → state). For each chunk, compute the final state assuming the initial state is zero:

$$h_\text{local}^{(j)} = B^{(j)\top} \text{decay} \cdot X^{(j)}$$

This is a matrix multiplication. Cost: $\text{BMM}(T/Q, N, P, Q)$.

Step 3: Center factors (state → state). Connect the chunks by propagating states through a scalar SSM scan of length $T/Q$:

$$h_\text{true}^{(j)} = a_\text{chunk}^{(j)} h_\text{true}^{(j-1)} + h_\text{local}^{(j)}$$

This is a 1-SS multiplication on $(N, P)$ independent channels. Cost: $O(T/Q \cdot NP)$ — negligible.

Step 4: Left factors (state → output). For each chunk, compute the output contribution from prior chunks:

$$Y_\text{off}^{(j)} = C^{(j)} \text{decay}_\text{out} \cdot h_\text{true}^{(j-1)}$$

This is a matrix multiplication. Cost: $\text{BMM}(T/Q, Q, P, N)$.

Final output: $Y = Y_\text{diag} + Y_\text{off}$.

8.4 Complexity

Setting $N = P = Q$ (state dimension = head dimension = chunk length):

• Total FLOPs: $O(TN^2)$ — same as attention but linear in $T$ for the dominant terms.
• Total memory: $O(TN)$ — linear in both sequence length and state size.
• The work is dominated by matrix multiplications on $(N, N)$ matrices.

This is the key advantage over Mamba’s selective scan: SSD uses matrix multiplication as its core primitive, which tensor cores are optimized for. Mamba’s scan is a custom CUDA kernel that cannot leverage these hardware units.

8.5 Speed comparison

The SSD algorithm is 2–8× faster than Mamba’s fused selective scan (Figure 10 of the Mamba-2 paper). For large state expansion ($N = 256$), SSD is 6× faster. For the default $N = 64$, SSD is 2× faster. SSD is also faster than FlashAttention-2 at sequence lengths beyond 2K and 6× faster at 16K.

---

9. The Mamba-2 Architecture

9.1 Block design changes

The Mamba-2 block modifies the Mamba block in two ways motivated by the attention connection:

Parallel parameter projections. In Mamba, the SSM parameters $(A, B, C)$ are computed from the post-convolution activation $x_c$, which depends on the initial linear projection. The projections are sequential.

In Mamba-2, $(A, X, B, C)$ are all produced from the input $u$ in parallel — analogous to how $(Q, K, V)$ are produced in parallel in a Transformer. This slightly reduces parameters and, more importantly, enables tensor parallelism for larger models by reducing the number of synchronization points per block from two to one.

Extra normalization. Mamba-2 adds a normalization layer (GroupNorm or RMSNorm) after the gating multiplication and before the output projection. GroupNorm (Wu and He, 2018) splits a feature vector into groups of channels and normalizes each group independently to zero mean and unit variance, sitting between LayerNorm (one group, the whole vector) and InstanceNorm (one group per channel). This improves training stability at larger scales and is analogous to the NormFormer architecture (Shleifer, Weston, and Ott, 2021) that adds normalization at the end of MLP and attention blocks.

9.2 Multi-head patterns

The state space duality allows transferring multi-head design choices from attention to SSMs.

Multi-head SSM (MHS) / Multi-head attention (MHA). The classic pattern: $H$ independent heads, each with its own $(A, B, C, X)$. The state size per head is $N$, and the head dimension is $P = D/H$.

Multi-input SSM (MIS) / Multi-value attention (MVA). The original Mamba architecture uses this pattern: $X$ has $H$ heads (one per channel), but $B$ and $C$ are shared across all heads. In attention terms, the keys and queries are shared while the values have independent heads. This is the natural choice from the SSM perspective because $X$ is the main input to the SSM, while $B$ and $C$ are auxiliary parameters.

Grouped-input SSM (GIS) / Grouped-value attention (GVA). Analogous to grouped-query attention (GQA), this creates $G$ groups of $B$ and $C$ projections, each shared across $H/G$ input heads. Mamba-2 uses this pattern with $G$ set to be a multiple of the tensor parallelism degree for efficient sharding.

The paper ablates these patterns (Table 5) and finds that the MVA/MIS pattern performs best, matching the choice naturally derived from the SSM perspective.

9.3 Kernel feature maps

The SSD framework allows incorporating kernel feature maps from the linear attention literature. In Mamba-2, the feature map $\psi$ is applied to the $B$ and $C$ branches (corresponding to $K$ and $Q$ in attention). By default, $\psi(x) = \text{Swish}(x) = x \cdot \sigma(x)$, following Mamba’s SiLU activation.

The paper ablates various kernel approximations (Table 6): cosFormer, Random Feature Attention, and Positive Random Features (Performer). None significantly improve over simple pointwise nonlinearities. This is expected because SSD differs from vanilla linear attention by the inclusion of the 1-semiseparable mask $L$, which already captures positional structure that kernel approximations were designed to provide.

9.4 Hybrid architectures

A striking finding from the Mamba-2 paper is that mixing SSD layers with attention layers improves over either alone. Table 2 shows that adding approximately 10% attention layers (6 out of 48 layers in a 350M model) reduces perplexity from 8.60 (pure SSD) to 8.26, with the best configuration using 7 attention layers.

At the 2.7B scale (Table 3), a Mamba-2 + MLP + Attention hybrid (28 SSD + 4 attention + 32 MLP layers) achieves an average downstream accuracy of 60.7%, compared to 60.2% for pure Transformer++ and 60.2% for pure Mamba-2.

The hypothesis: SSM layers function as general sequence-to-sequence mappings that compress context into their recurrent state, while attention layers act as a retrieval mechanism that can refer directly to previous tokens without compression. A small number of attention layers provides an “escape hatch” for tasks that require exact token lookup, while SSD handles the bulk of the computation more efficiently.

---

10. Scaling Results

10.1 Mamba scaling laws

On the Pile dataset (Figure 4 of the Mamba paper), Mamba matches Transformer++ scaling from 125M to 1.3B parameters at context length 2048. At context length 8192, Mamba further improves relative to Transformers (which are limited by the quadratic cost of longer sequences).

On downstream zero-shot evaluations (Table 3), Mamba at each model size matches baselines at twice the size:

• Mamba-130M matches Pythia-160M
• Mamba-370M matches Pythia-410M
• Mamba-1.4B matches Pythia-2.8B

10.2 Mamba-2 scaling laws

Mamba-2 is Pareto-dominant over both Mamba and Transformer++ (Figure 9 of the Mamba-2 paper): it achieves lower perplexity at every FLOP budget from 125M to 1.3B parameters. This is because SSD is both faster (enabling more training tokens per wall-clock hour) and slightly more expressive (due to larger state sizes enabled by the efficient algorithm).

On downstream evaluations at 2.7B scale (Table 1 of the Mamba-2 paper), Mamba-2 matches Mamba’s quality while being 2–8× faster to train.

10.3 State expansion

One of SSD’s most important practical benefits is efficient state expansion. In Mamba, increasing the state dimension $N$ from 16 to 64 provides a significant perplexity improvement (from 9.82 to 8.71 for a 350M model — Table 10 of the Mamba paper) but at the cost of proportionally slower selective scan.

In Mamba-2, the SSD algorithm’s speed is nearly independent of $N$ up to $N = 256$ (Figure 10, right panel). This allows Mamba-2 to use much larger state sizes without slowdown, effectively making the capacity-efficiency tradeoff from the Kernel Zoo blog far less severe.

---

11. The Full Picture: From Selection to Duality

11.1 The progression

The two papers together tell a coherent story:

1. Prior SSMs (S4, S5, H3, Hyena) are LTI systems. They process sequences through convolutions during training and recurrence during inference. They are fast but cannot do content-based reasoning.

2. Mamba makes the SSM parameters input-dependent (selective). This enables content-based reasoning and matches Transformer quality. But the convolutional mode is lost, and the model relies on a custom scan kernel.

3. Mamba-2 reveals that selective SSMs with scalar $A$ are equivalent to a form of structured attention. This equivalence — structured state space duality — exposes both a quadratic (attention-like) and a linear (recurrence-like) algorithm. The SSD algorithm combines both by chunking the sequence: quadratic within chunks, linear between chunks. The result is faster than both pure attention and pure recurrence.

11.2 Connection to the blog series

The linear attention framework from the Why Replace Attention blog replaces the softmax with a kernel: $Y = (L \circ \phi(Q)\phi(K)^\top) V$, where $L$ is the all-1’s causal mask. RetNet replaces $L$ with a decay mask $L_{ts} = \gamma^{t-s}$. The Gated DeltaNet blog modifies the recurrent update rule.

SSD generalizes all of these. The mask $L$ is a 1-semiseparable matrix with input-dependent entries $a_t$, not fixed scalars. This means:

• Linear attention is SSD with $a_t = 1$ for all $t$ (no decay, causal mask of 1’s).
• RetNet is SSD with $a_t = \gamma$ for all $t$ (constant decay).
• Mamba-2 is SSD with $a_t = \exp(\Delta_t \cdot A_\text{scalar})$ varying per token (input-dependent decay).

The feature map $\phi$ and the mask $L$ are orthogonal design choices. The Kernel Zoo blog explored $\phi$. The SSD framework shows that $L$ — the structured mask — is equally important, and that making $L$ input-dependent is what gives SSMs their selectivity.

---

Summary

State space models process sequences through a latent state governed by a linear recurrence, with dual convolutional and recurrent computation modes. Prior SSMs kept their parameters fixed (LTI), which enabled convolutions but prevented content-based reasoning. Mamba introduces selection — making $\Delta$, $B$, $C$ functions of the input — which breaks the LTI property but gives the model a learnable gate ($g_t = \sigma(\text{Linear}(x_t))$, equivalent to the ZOH discretization of a leaky integrator) that decides per token whether to read or skip. Mamba-2 then reveals the deeper structure: every SSM is a multiplication by a semiseparable matrix, and when $A$ is scalar, this matrix factors as a 1-semiseparable mask times a Gram matrix — exactly the structure of masked attention with an input-dependent decay mask. The SSD algorithm exploits this duality through block decomposition: quadratic attention within chunks (leveraging matrix multiplication hardware), linear recurrence between chunks (keeping the cost linear in sequence length), producing an architecture that is 2–8× faster than Mamba’s custom scan while matching or exceeding Transformer quality at scales up to 2.7B parameters.

---

Previous: Mathematical Prerequisites for Mamba
Next: The Efficient Transformer Design Space: Comparing All Variants and the Three Futures of Attention