State-Space Models (S4 / Mamba)

Every time a transformer reads a sentence, it lets each word look at every other word. That flexibility is powerful — but it costs $O(n^2)$ in time and memory. Double the sequence and you pay four times the price.

State-space models (SSMs) take a different bet. Instead of full attention, they compress the entire past into a compact hidden state and update it one step at a time. The update is a linear recurrence — fast, parallelisable, and $O(n)$ in the sequence length.

The key insight, pioneered by S4 (Gu et al., 2021) and sharpened by Mamba (Gu & Dao, 2023), is that the right choice of structured matrices makes linear recurrence just as expressive as attention on most real tasks — language, audio, DNA — while being dramatically cheaper on long sequences.

This is an open research frontier. SSMs are not a settled theory; they are a fast-moving race to find the architecture that beats transformers and attention without paying the quadratic price.

The demo below runs the same sequence through two algorithms — a linear recurrent scan (SSM style) and a simulated quadratic attention step count — as you change the sequence length. Watch the operation count diverge.

<!-- {{c_html_intro}} -->
<p class="hint">{{hint_para}}</p>
<div class="controls">
  <label for="seqLen">{{label_seq_len}} <span id="lenDisplay">16</span></label>
  <input type="range" id="seqLen" min="4" max="128" value="16" step="4">
</div>
<div class="chart-wrap">
  <canvas id="chart" width="440" height="200"></canvas>
</div>
<div class="stats-grid">
  <div class="stat ssm-stat">
    <div class="stat-label">{{label_ssm}}</div>
    <div class="stat-value" id="ssmOps">—</div>
    <div class="stat-note">O(n)</div>
  </div>
  <div class="stat attn-stat">
    <div class="stat-label">{{label_attn}}</div>
    <div class="stat-value" id="attnOps">—</div>
    <div class="stat-note">O(n²)</div>
  </div>
</div>
<div class="status" id="status"></div>
<div class="btns">
  <button id="stepBtn" type="button">{{btn_step}}</button>
  <button id="resetBtn" type="button" class="ghost">{{btn_reset}}</button>
</div>
<!-- {{c_html_seq}} -->
<div id="seqViz" class="seq-viz"></div>

/* {{c_css_intro}} */
* { box-sizing: border-box; }
body { font-family: system-ui, sans-serif; color: #222; margin: 0; padding: 14px 2px; }
.hint { font-size: .88rem; color: #444; margin: 0 0 .7rem; line-height: 1.45; }
.controls { display: flex; align-items: center; gap: .6rem; margin-bottom: .6rem; font-size: .9rem; }
input[type=range] { flex: 1; accent-color: #1d3557; }
.chart-wrap { background: #f4f7fa; border-radius: 10px; padding: 8px; margin-bottom: .7rem; }
canvas { display: block; width: 100%; height: auto; }
.stats-grid { display: grid; grid-template-columns: 1fr 1fr; gap: .5rem; margin-bottom: .7rem; }
.stat { border-radius: 10px; padding: .5rem .7rem; text-align: center; }
.ssm-stat { background: #e0f0e9; }
.attn-stat { background: #fde8e8; }
.stat-label { font-size: .78rem; color: #555; font-weight: 600; text-transform: uppercase; letter-spacing: .04em; }
.stat-value { font-size: 1.35rem; font-weight: 700; margin: .1rem 0; }
.ssm-stat .stat-value { color: #0a7d33; }
.attn-stat .stat-value { color: #c92f3c; }
.stat-note { font-size: .75rem; color: #888; font-family: ui-monospace, monospace; }
.status { font-size: .95rem; font-weight: 600; min-height: 1.4em; margin-bottom: .4rem; color: #1d3557; }
.btns { display: flex; gap: .5rem; flex-wrap: wrap; margin-bottom: .7rem; }
button { font: 600 14px system-ui, sans-serif; padding: .45rem .9rem; border: 1px solid #1d3557;
         background: #1d3557; color: #fff; border-radius: 8px; cursor: pointer; }
button.ghost { background: #fff; color: #1d3557; }
/* {{c_css_seq}} */
.seq-viz { display: flex; flex-wrap: wrap; gap: 4px; margin-top: .3rem; }
.tok { width: 32px; height: 32px; border-radius: 6px; display: flex; align-items: center;
       justify-content: center; font: 700 11px ui-monospace, monospace;
       background: #e8eef3; color: #1d3557; border: 1px solid #cdd9e3; transition: background .2s; }
.tok.processed { background: #0a7d33; color: #fff; border-color: #0a7d33; }
.tok.current { background: #1d3557; color: #fff; border-color: #1d3557; box-shadow: 0 0 0 2px #a8c4e0; }

// Code not found

Notice the gap. For short sequences both are manageable. As the slider climbs, attention's cost rockets while the linear scan stays flat by comparison. This is exactly why SSMs matter at the scale of genomics, audio, and very long documents.

The core of every state-space model is deceptively simple:

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

where $h_t$ is the hidden state, $x_t$ the input token, and $A, B, C$ are learnable matrices. One step: one matrix–vector multiply. No attention mask, no softmax.

Three ideas make this competitive with transformers:

• Parallel prefix scan: unrolling the recurrence into a convolution means the whole sequence can be processed in $O(n \log n)$ on GPUs — not sequentially one step at a time.
• HiPPO initialisation (Gu et al., 2020): choosing $A$ as a specific polynomial-projection matrix gives the model a structured way to remember the distant past, solving the gradient-vanishing problem that plagued older RNNs.
• Selective state spaces (Mamba): input-dependent $B$ and $C$ let the model decide which tokens are worth remembering — achieving the selective attention of transformers without the quadratic cost.

The computational status: SSMs are not NP-hard in any classical sense — they are a machine-learning architecture question, not a complexity one. The open question is whether any fixed-size state can match the expressivity of unbounded attention, which connects to circuit complexity and learning theory.

SSMs shine whenever sequences are long enough to make attention prohibitively expensive:

• Genomics: DNA and protein sequences stretch to millions of bases. Mamba-based models like HyenaDNA process full chromosomes that would be impossible for standard transformers.
• Audio: raw waveforms at 44 kHz are far too long for attention. SSMs were state-of-the-art on audio classification benchmarks years before they attracted mainstream attention.
• Time-series forecasting: sensor logs, EHRs and financial tick data all benefit from linear-time processing over very long horizons.
• Language at scale: Mamba-3 and hybrid SSM-attention architectures now approach GPT-4 quality at a fraction of the inference cost, suggesting the two paradigms may converge.

The same $O(n)$ advantage that makes SSMs interesting to researchers also makes them practical to deploy: inference on a trained SSM is a simple recurrence — one hidden state, one matrix multiply per token, constant memory.

Transformers taught us that letting every token talk to every other token is extraordinarily powerful. State-space models ask: what if you could get that power by remembering the right summary instead?

S4 proved the idea works. Mamba made it selective. Hybrid architectures are now blending both. The race is live, the theory is incomplete, and the winner — if there is one — will reshape how we think about sequence modeling from proteins to paragraphs.

The deeper lesson is familiar: the hardest part of a computational problem is often not the algorithm but the representation. Choose the right structure for your state matrix, and a recurrence that looks trivial can match the most powerful architecture we have built.