State Space Models (Mamba, S4): A Frontier-Lab Interview Deep Dive

Why this exists. SSMs are the most credible challenger to transformers. They give $O(N)$ sequence complexity with constant memory at inference — properties transformers don’t have. Frontier-lab interviews increasingly ask about Mamba, the selectivity mechanism, and hybrid models. This document covers the math without the dense academic notation.

1. The big picture

A state space model maintains a hidden state $h_t$ and applies a recurrence:

$$h_t=A{h}_{t-1}+B{x}_t\text{(state update)}$$

$$y_t=C{h}_t+D{x}_t\text{(output)}$$

This is exactly an RNN with linear (no nonlinearity in the recurrence) dynamics. The trick is what you do with this.

Linear recurrences have two equivalent computational forms:

1. Recurrent: compute $h_t$ from $h_{t-1}$ step by step. $O(N)$ time, $O(d)$ memory.
2. Convolutional: unroll into a convolution $y=x\ast K$, where $K$ is a learned kernel. $O(N\log N)$ with FFT.

The duality is the key idea: SSMs train via convolution (parallel) but generate via recurrence (constant memory per step). Best of both worlds.

Saying it out loud. If someone asks what a state space model is, I’d say it’s a machine that keeps one running summary of everything it has seen. Each step it fades the old summary a bit, mixes in the new token, and reads an answer off of it. The clever part isn’t the equation, it’s that because the update is linear you can compute it two different ways — step by step when generating, or all at once as a convolution when training. So you get an RNN’s constant-memory decoding without an RNN’s sequential training bottleneck. That duality is the entire reason anyone builds these.

2. The classical state space ODE

In plain language. This section is borrowed straight from control theory, where the same equations describe things like a cruise-control system or a robot arm. The idea is that there’s a hidden internal state you never see directly, an input pushing on it, and an output you read off it. The math below just writes down how the state changes over time and how the output is read out.

The continuous version:

$$\frac{dh}{dt}=Ah(t)+Bx(t)$$

$$y(t)=Ch(t)+Dx(t)$$

For input signal $x(t)$, the state $h(t)$ evolves linearly via $A$; output $y(t)$ is a linear readout. Same equations as in control theory and signal processing.

Saying it out loud. The continuous form is the original: it says the rate of change of the internal state is some mix of the current state plus whatever the input is pushing in, and the output is a linear readout of that state. Control engineers have used exactly this to model physical systems for decades — think of the state as position and velocity and the input as the throttle. Machine learning borrows the form and just learns the matrices instead of deriving them from physics. The useful mental hook is that $A$ controls how memory decays on its own, $B$ controls what gets written in, and $C$ controls what gets read out.

Discretization

For machine learning, we work with discrete sequences. Use zero-order hold to discretize:

$$A_d=\exp (A\Delta t)$$

$$B_d=(A_d-I)A^{-1}B$$

$\Delta t$ is a step size (often learned). The discretized recurrence:

$$h_t=A_d{h}_{t-1}+B_d{x}_t$$

Same form as before but with $A_d$, $B_d$ as discrete-time matrices.

Saying it out loud. Discretization is how you get from smooth continuous time to one-token-at-a-time. You pick a step size $\Delta t$ and ask what the continuous system would have done over that interval, assuming the input stays constant across the step — that assumption is what ‘zero-order hold’ means. Out of that falls a matrix exponential, $A_d=\exp (A\Delta t)$, which is the discrete-step version of the decay. The thing to actually remember is what $\Delta t$ means intuitively: large $\Delta t$ means the state moves a lot per token, small $\Delta t$ means the token barely registers, and Mamba’s whole contribution is making that dial input-dependent.

3. The convolutional view

In plain language. This section shows that if you write out the recurrence for a few steps, a pattern appears: the output at any time is just a weighted sum of all past inputs, with weights that follow a fixed geometric-ish pattern. That means the whole recurrence can be rewritten as a convolution with a single filter. The algebra below is just that unrolling done carefully.

Unrolling the recurrence:

$$\begin{array}{rl}{h}_0& =B{x}_0\\ h_1& =AB{x}_0+B{x}_1\\ h_2& =A^{2}B{x}_0+AB{x}_1+B{x}_2\\ & ⋮\\ y_t& =\sum _{k=0}^{t}C{A}^{k}B{x}_{t-k}\end{array}$$

This is a convolution with kernel $K_k=C{A}^{k}B$ (and $D{x}_t$ if $D$ is included). For a length-$N$ sequence: $y=K\ast x$.

The kernel $K$ has length $N$ (or up to $N$). Computing this convolution:

• Direct: $O(N^2)$ — same as attention.
• FFT: $O(N\log N)$ — better but requires structured $A$.

The breakthrough: choose $A$ such that $A^k$ can be efficiently computed.

Saying it out loud. If you unroll the recurrence by hand for a few steps, you notice the output is just a weighted sum of every past input, and the weight on something $k$ steps back is always $C{A}^{k}B$. That’s exactly what a convolution is — one fixed filter slid over the sequence. And convolutions are perfect for GPUs, because every output position can be computed in parallel instead of waiting for the previous one. The catch is the filter is as long as your sequence, so you need an FFT to do it in $O(N\log N)$, and you need a special structure on $A$ to build those powers cheaply in the first place. That last requirement is what the next two sections are about.

4. HiPPO: the theoretical foundation

In plain language. HiPPO is about how you pick the $A$ matrix at initialization. The problem it solves is that a random $A$ either forgets everything within a few steps or blows up. HiPPO derives a specific $A$ whose state is a compressed sketch of the entire history so far.

HiPPO (High-order Polynomial Projection Operators) — Gu et al. 2020. Provides a principled choice for $A$:

The HiPPO matrix is constructed so that the hidden state $h_t$ is a compressed representation of the history of $x$ up to time $t$. Specifically, the columns of $h_t$ represent coefficients of a polynomial approximation of $x_{0:t}$ in some basis (Legendre, Fourier, etc.).

This gives the SSM a principled inductive bias: the model can in principle “remember” all of history, weighted by an interpretable polynomial basis.

Saying it out loud. HiPPO answers a very concrete question: what should $A$ be so the state actually remembers? The trick is to stop thinking of the state as raw memory and start thinking of it as the coefficients of a curve fit — the state stores the best polynomial approximation of everything the model has seen. Once you frame it that way you can solve for the matrix that keeps that approximation optimal as new data arrives, and that’s the HiPPO matrix. So instead of hoping the optimizer discovers long-range memory, you build it in at step zero — and empirically that’s worth tens of points on long-range benchmarks, not a few.

Why this matters for ML

A randomly-initialized SSM doesn’t have any reason to remember long-range patterns. HiPPO initialization guarantees that, at init, the model can capture history with bounded error. Empirically: HiPPO-initialized SSMs train much better than randomly initialized ones.

Saying it out loud. The blunt version is that random initialization basically never gives you long-range memory. The eigenvalues of a random $A$ are either slightly inside the unit circle, in which case memory decays exponentially and you’ve forgotten anything past a few dozen steps, or slightly outside, in which case the state explodes. HiPPO lands you in the narrow useful regime by construction, with a proven bound on how much history the state preserves. It’s a great example of an initialization choice mattering more than an architecture choice — same model, same parameter count, and the ablation swings Long Range Arena results dramatically.

5. S4: Structured State Spaces

S4 (Gu, Goel, Re, 2022). Practical SSM that combines HiPPO with computational efficiency.

Key contributions

1. Diagonal Plus Low-Rank (DPLR) parameterization of $A$.

HiPPO matrices are dense. S4 reparameterizes as $A=\Lambda +p{q}^{\top }$ (diagonal + rank-1 update). This makes computing $A^k$ tractable.

2. Efficient kernel computation.

The convolution kernel $K_k=C{A}^{k}B$ can be computed in $O(N\log N)$ time using a Cauchy-style structured matrix multiplication (instead of $O(N^2)$ for general $A$).

3. Stable parameterization.

Use HiPPO-LegS initialization for theoretical guarantees, plus tricks to ensure $A$’s eigenvalues stay stable.

Saying it out loud. S4’s problem was purely computational: to build the convolution kernel you need every power of $A$ up to sequence length, and for a dense matrix that’s hopeless inside a training loop. So S4 constrains $A$ to a diagonal matrix plus a rank-one correction, which is expressive enough to represent HiPPO but structured enough that the powers collapse into fast Cauchy-kernel math. Add a careful parameterization that keeps eigenvalues stable and you get an $O(N\log N)$ kernel. The tradeoff is implementation pain — DPLR is genuinely nasty to write, which is why later work found that plain diagonal $A$ with good init gets most of the benefit.

Result

S4 was the first SSM to match transformers on long-range tasks (Long Range Arena benchmark) while having $O(N\log N)$ complexity. It established SSMs as a credible architecture.

Saying it out loud. The headline is that S4 was the first state space model that actually beat transformers at something people cared about. On Long Range Arena, which has sequences in the thousands of steps, S4 cleared the whole benchmark including Path-X, which every transformer variant had been failing at chance level. That result is what turned SSMs from a control-theory curiosity into a serious architecture line. The honest caveat is that Long Range Arena is not language — S4’s language modeling was still noticeably behind, and closing that gap took Mamba.

6. Mamba: Selective State Spaces

Mamba (Gu & Dao, 2023). The breakthrough that made SSMs competitive at LLM scale.

The selectivity insight

In plain language. This is the core of Mamba. The equations below say that instead of using the same fixed matrices at every position, the model computes fresh ones from the current token. That single change lets each token control how strongly it writes into memory and how fast old memory fades.

In S4, the matrices $A$, $B$, $C$ are shared across all positions — a single linear time-invariant (LTI) system. This is fast (the convolution view works) but inflexible: the model cannot decide that some inputs are more “important” or change its dynamics based on input.

Mamba makes $B$, $C$, and $\Delta t$ input-dependent:

$$B(x_t)={\mathrm{Linear}}_B(x_t)$$

$$C(x_t)={\mathrm{Linear}}_C(x_t)$$

$$\Delta t(x_t)=\mathrm{softplus}({\mathrm{Linear}}_{\Delta t}(x_t))$$

The state update becomes:

$$h_t=A(\Delta t(x_t))h_{t-1}+B(x_t)x_t$$

$$y_t=C(x_t)h_t$$

In Mamba 1, $A$ is initialized as a diagonal real-valued matrix (S4D-Real / HiPPO-LegS-diagonal), a simplification from the full HiPPO-DPLR structure of S4. Its discretization now depends on input via $\Delta t$. Each token can choose how much to remember (large $\Delta t$) vs forget (small $\Delta t$). This is the “selective” mechanism.

Saying it out loud. The insight is that S4 treats every token identically, which is a strange thing for a language model to do. A comma and a person’s name got the exact same dynamics, because the matrices were fixed across the sequence — that’s what ‘time-invariant’ means, and it’s what made the convolution trick legal. Mamba makes $B$, $C$ and $\Delta t$ into cheap linear functions of the current token, so each position sets its own remember-versus-forget dial. Concretely, a big $\Delta t$ means this token writes hard into the state and the state persists; a small one means it slides by without disturbing anything. That’s what fixed the selective-copying tasks that S4 simply could not do.

Cost of selectivity

The convolutional view no longer works: $K$ depends on input, so it’s no longer a single shared kernel. Mamba reverts to the recurrent view but uses a parallel scan algorithm (Blelloch scan) to compute the recurrence in parallel.

Parallel scan: compute $h_1,h_2,\dots ,h_N$ in $O(N\log N)$ parallel ops with $O(N)$ work.

Hardware-aware implementation in CUDA. Throughput comparable to (or better than) attention on modern GPUs.

Saying it out loud. There’s no free lunch here — the moment the matrices depend on the input, the convolution view dies. The kernel was one fixed filter for the whole sequence, and now every position has a different one, so there’s nothing to precompute and FFT. Mamba goes back to the recurrence but computes it with a parallel scan, the same associativity trick behind fast prefix sums, which turns $N$ sequential steps into $\log N$ parallel ones. The other half is engineering: the CUDA kernel keeps state in SRAM instead of round-tripping through HBM, exactly like FlashAttention, and without that memory-traffic work the whole thing would be slower than the attention it replaces.

Result

Mamba matches transformer quality at the same parameter count and compute on language modeling, with $O(N)$ sequence complexity and constant-memory inference. The first credible drop-in replacement for transformer attention.

Saying it out loud. The bottom line is that Mamba was the first SSM to actually match transformers on language modeling at matched parameters and compute, not just on synthetic long-range benchmarks. On top of that it kept the two structural advantages: linear scaling in sequence length and a decode state that doesn’t grow with context. That combination is what made people take it seriously as a drop-in replacement rather than a niche architecture. The asterisk is scale — the results were convincing up to a few billion parameters, and the recall gap reappears as you go bigger.

7. Why SSMs are interesting for LLMs

Linear sequence complexity

Attention is $O(N^2)$ in compute and memory. SSMs are $O(N)$ (with $\log N$ factors for parallel scan). At long context (32K+), this is a huge advantage.

Saying it out loud. Attention compares every token to every other token, so cost scales with the square of context length — double the context and you quadruple the work. An SSM just walks the sequence once updating a fixed-size state, so it’s linear. At 32K tokens that difference is enormous on paper. The honest caveat is that attention kernels are extremely well optimized, so the crossover point where SSMs actually win in wall-clock is further out than the asymptotics suggest — at 2K context the transformer is usually still faster.

Constant-memory inference

For autoregressive decoding: each step is $O(d)$ work and $O(d)$ memory. KV cache doesn’t grow. Massive memory savings for long-context inference.

Saying it out loud. This is the advantage that’s hardest to argue with. When a transformer generates, it keeps a key and value vector for every token it has ever seen, so the cache grows linearly with context and every new token has to read all of it. An SSM keeps one fixed-size state — same size at token ten as at token a hundred thousand. In serving terms that’s the difference between a KV cache measured in tens of gigabytes at long context and a state measured in megabytes, which directly translates to how many concurrent users fit on one GPU.

Empirical quality

Mamba matches transformer quality on many language tasks at small-to-medium scale (≤7B). At larger scale, an “in-context recall” gap re-emerges (transformers’ attention is naturally good at copying). Mamba-2 (Dao & Gu 2024) reformulates the selective SSM as structured state space duality (SSD) — using semiseparable matrices, the SSM operation becomes a structured matmul that maps onto tensor cores efficiently. This is a substantial speedup, not a minor improvement, and enables much larger state dimensions. Hybrid models (Jamba, Falcon Mamba 7B, Codestral Mamba) interleave SSM and attention layers to recover transformer-level recall while keeping Mamba’s long-context efficiency.

Saying it out loud. The picture is: competitive at small and medium scale, with a specific weakness that grows with scale. Up to around 7B, Mamba matches transformers on most language benchmarks. What re-emerges as you scale is in-context recall — the ability to copy something verbatim from earlier in the prompt, which attention does natively via induction heads and a compressed state has to approximate. Mamba-2 helped on the efficiency side by reformulating the operation as structured matmuls that hit tensor cores, but the fix for recall is architectural: interleave a few attention layers, which is what Jamba and the other hybrids do.

Why they haven’t replaced transformers (yet)

• In-context learning weaker. Transformers’ attention is naturally good at copying from earlier in context (induction heads). SSMs have weaker copy-and-recall behavior empirically.
• Calibration / uncertainty. Transformers’ attention provides interpretable patterns; SSMs less so.
• Ecosystem. Transformers have years of optimization (FlashAttention, vLLM, paged attention). SSM tooling is younger.
• Scaling laws. Whether SSMs match transformers at frontier scale (100B+) is still being established.

Saying it out loud. Three things, and only one is really about quality. The genuine quality issue is recall — a fixed-size state has to compress everything, so exact lookup from far back in the context is fundamentally harder than for attention, which can just point at the token. The other two are ecosystem and evidence: every serving stack in the world is built around attention, and nobody has published a clean frontier-scale head-to-head. That’s why labs are hedging with hybrids instead of betting a full pretraining run on a pure SSM.

8. Hybrid architectures

Recent research suggests mixing attention and SSM layers gives the best of both:

• Attention layers for in-context learning, copy, and exact recall.
• SSM layers for long-range mixing with $O(N)$ cost.

Examples

Jamba (AI21, 2024). 7-to-1 SSM-to-attention ratio. 256K+ context. Combines Mamba blocks with transformer attention blocks and MoE.

Zamba (Zyphra). Hybrid SSM-attention with MoE.

Bamba, Samba, Hymba — various hybrid designs. Active research area.

The frontier-lab interview question: “Are pure SSMs going to replace transformers?” Most likely answer: hybrids are the practical compromise; pure SSMs may not catch up at frontier scale, but mixed-block architectures will be increasingly common.

Saying it out loud. The pattern across all of these is the same: mostly SSM layers with a few attention layers sprinkled in, usually around one attention layer per six or seven SSM layers. Jamba was the first production one, from AI21 in 2024, and it added mixture-of-experts on top to get more capacity without more compute per token. The reason it works is that you don’t need attention everywhere — you need it in a handful of places where exact retrieval matters, and the SSM majority keeps the cache small. The number that made people pay attention was Jamba fitting 140K tokens of context on a single 80GB GPU, which a dense transformer that size simply cannot do.

9. Mamba vs LSTM vs RNN

People sometimes ask “isn’t this just an RNN?”

Yes, mathematically

Mamba is a linear RNN (no nonlinearity in the recurrence — selectivity is in the input-dependent matrices, not in a nonlinear gate).

But practically very different

• LSTM: nonlinear gates, hard to parallelize, vanishing gradients with depth. Mamba: linear recurrence, parallel scan, stable gradients via structured $A$.
• Vanilla RNN: unstable, can’t be trained at scale. Mamba: HiPPO-initialized, stable, scales.
• Linear attention: also $O(N)$, but Mamba’s selectivity gives more expressiveness.

So “linear RNN” is technically right but misleading. Mamba is what RNNs always wanted to be.

Saying it out loud. Saying Mamba is just an RNN is technically true and completely unhelpful. Every property that made RNNs unusable comes from details Mamba changed: the nonlinearity in the recurrence, which forces sequential training and vanishing gradients, and the dense random recurrent matrix, which makes it unstable. Mamba keeps the recurrence linear so it parallelizes with a scan, uses a structured diagonal $A$ with principled init so gradients stay well-behaved over thousands of steps, and puts its gating in the input-dependent $\Delta t$ instead. The lesson is that in sequence modeling the equation tells you almost nothing — initialization and hardware mapping decide whether something scales.

10. The mathematical machinery (briefly)

In plain language. This section collects the three formulas you’d actually be asked to reproduce on a whiteboard: how continuous time becomes discrete steps, what the convolution filter looks like written out, and the associativity trick that makes the recurrence parallelizable. None of it is deep — it’s the notation you need to have at your fingertips.

For interview-grade understanding:

Discretization:

$$A_d=\exp (A\Delta t)\approx I+A\Delta t+\frac{(A\Delta t{)}^2}{2}+\cdots \text{(Taylor)}$$

or the zero-order hold formula above.

Convolution kernel:

$$K=(CB,CAB,C{A}^{2}B,\dots ,C{A}^{N-1}B),y=K\ast x$$

Parallel scan for selective SSM:

The associativity trick: define $(a_1,b_1)\oplus (a_2,b_2)=(a_2{a}_1,a_2{b}_1+b_2)$. Then $h_t=a_t{h}_{t-1}+b_t$ can be computed via prefix scan over $(a,b)$ pairs. Parallelizable; runs in $O(\log N)$ parallel time.

For $A$ diagonal in Mamba (each state dim is independent), this scan is straightforward. The hardware-aware Mamba kernel does this efficiently.

Saying it out loud. If I had to whiteboard this, three things carry the whole story. First, discretization is a matrix exponential — you can quote the Taylor expansion if pressed, and the first-order term $I+A\Delta t$ is usually enough to make the point. Second, the convolution kernel is literally the sequence $CB,CAB,C{A}^{2}B$ and so on, which is why structure on $A$ matters so much. Third, the parallel scan works because composing two linear updates gives you another linear update, and that composition is associative — associativity is exactly the property that lets you use a tree instead of a chain, so $N$ steps become $\log N$ depth. In Mamba, $A$ is diagonal, so each state dimension scans independently and the whole thing is embarrassingly parallel.

11. Common interview gotchas

12. The 8 most-asked SSM interview questions

1. What’s an SSM? Linear recurrence $h_t=A{h}_{t-1}+B{x}_t,y_t=C{h}_t$. Trains via convolution; generates via recurrence.
2. What’s Mamba? SSM with input-dependent $B,C,\Delta t$ (selectivity). Matches transformer quality at $O(N)$ complexity.
3. What’s HiPPO? Theoretical initialization for $A$ that makes the state a polynomial approximation of past inputs.
4. Why is the convolutional view useful? Parallelizes training (FFT, $O(N\log N)$). The recurrence runs sequentially.
5. Why does Mamba use parallel scan, not convolution? Selectivity makes $K$ input-dependent; can’t use a fixed kernel.
6. Memory advantage of SSMs? Constant memory at decode ($O(d)$), vs growing KV cache for transformers.
7. What’s the in-context learning gap? Transformers excel at copy-and-recall via induction heads. SSMs are weaker; hybrid models compensate.
8. What’s a hybrid SSM-attention model? Mix attention and SSM layers. Jamba, Zamba. Combines both architectures’ strengths.

13. Drill plan

1. Memorize the $h_t=A{h}_{t-1}+B{x}_t$ recurrence and its convolutional unrolling.
2. Know HiPPO’s role in initializing $A$.
3. Explain Mamba’s selectivity: input-dependent $B,C,\Delta t$.
4. Know parallel scan = how Mamba parallelizes training.
5. Cite hybrid models (Jamba) as the practical compromise.
6. Drill INTERVIEW_GRILL.md.

14. Further reading

• Gu et al., “HiPPO: Recurrent Memory with Optimal Polynomial Projections” (2020).
• Gu, Goel, Re, “Efficiently Modeling Long Sequences with Structured State Spaces” (S4, 2022).
• Gu & Dao, “Mamba: Linear-Time Sequence Modeling with Selective State Spaces” (2023).
• Dao & Gu, “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality” (Mamba-2, 2024).
• AI21, “Jamba” (2024) — hybrid SSM-Transformer.