Mamba Selective State-Space Module

A sequence-mixing module that updates a compact recurrent state as each token arrives, using input-dependent parameters to decide what the state stores, forgets, and emits.

A selective state-space module (SSM) processes tokens in order while carrying a small hidden state forward through time. Mamba makes that state input-dependent so each step can selectively store, forget, and emit information instead of using fixed recurrence rules.

At a glance

Released

December 2023

Authors

Albert Gu, Tri Dao

Optimizes

  • Sequence Scaling
  • Long Context Inference
  • Memory Bandwidth

Example models

What It Is

A state-space module is a sequence-mixing block that reads tokens one at a time and keeps a compact hidden state that summarizes what came before. At each position, the module combines the current token with that carried state to produce an output and an updated state for the next step. A selective state-space model (SSM) changes the update rules based on the current input rather than applying the same fixed transition at every position.

Why It Exists

Long sequences strain attention because dense attention compares every token to every other token. State-space modules offer a different path: they propagate information through a fixed-size state as the sequence unfolds, which can scale more gently with length when the recurrence is implemented efficiently.

How It Works

The module walks the sequence in order. For each token, it projects the input into input-dependent parameters that control the state transition and output. The previous hidden state is updated into a new state, and the module emits an output vector for that position. In Mamba, those transition parameters depend on the current token, so the block can decide what to remember, what to discard, and what to expose at each step. That selective behavior is the main change from earlier fixed-coefficient state-space designs.
Sequence mixing over time
State carried across steps
Input-dependent update path
Current token input
Selective gate parameters
Dense attention compares the current query against every past key position. Mamba instead walks tokens in order and updates a compact hidden state through input-dependent gates at each step.

Math Or Compute Schema

Dense attention forms outputs from pairwise token scores. A selective state-space step instead updates a hidden state and projects an output from that state. The formulas below contrast multi-head attention against a selective state update and output projection at one time step.
Multi-head attention (MHA)
Attention(Qi,Ki,Vi)=softmax ⁣(QiKi⊤dk)Vi\text{Attention}(Q_i, K_i, V_i) = \mathrm{softmax}\!\left(\frac{Q_i K_i^{\top}}{\sqrt{d_k}}\right) V_i
QQ
Query vectors for head i.
KK
Key vectors for head i.
VV
Value vectors for head i.
HH
Number of query heads.
dkd_k
Key dimension per head.
ii
Query head index.
Selective state-space step
ht=Aˉtht−1+Bˉtxt,yt=Cthth_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t, \quad y_t = C_t h_t
hth_t
Hidden state after processing token at position t.
ht−1h_{t-1}
Hidden state carried from the previous position.
xtx_t
Input vector for the token at position t.
yty_t
Output vector emitted at position t.
Aˉt\bar{A}_t
Input-dependent state transition applied at step t.
Bˉt\bar{B}_t
Input-dependent input injection at step t.
CtC_t
Input-dependent output projection at step t.
tt
Sequence position index.

Compared To Nearby Modules

Multi-head attention mixes tokens by scoring every query against every key position. Linear attention keeps a running summary but still routes through attention-style feature maps. A Mamba-style state-space module does not build an explicit all-pairs score map; it compresses history into a recurrent state and lets input-dependent gates shape each update.
Comparison dimensionMamba Selective State-SpaceMulti-Head AttentionLinear Attention
Sequence mixing mechanismOrdered recurrent state updates with input-dependent transition parametersSoftmax-weighted combination of value vectors from all key positionsFeature-map accumulation over keys without materializing full score matrix
Direct all-pairs token comparisonNo explicit n-by-n attention map; history flows through a compact stateEvery query token scores against every key tokenApproximates dense attention through kernel-style summaries
History representationFixed-size hidden state carried across positionsPast tokens remain addressable through keys and values at each stepRunning key-value summaries rather than a single recurrent state vector

Example Architectures

Hybrid stacks such as Nemotron 3 Super combine Mamba-style state-space layers with attention layers so the model can use efficient recurrence for much of the depth while keeping attention where direct token-to-token comparison is still useful.

Limitations And Tradeoffs

State-space mixing is inherently sequential: each step waits on the previous state, which affects parallel training patterns and can complicate implementation compared with batched attention kernels. The compact state also means less direct all-pairs token comparison than dense attention provides in a single layer, which is one reason hybrid Mamba-attention models still include attention blocks. None of this replaces benchmark evaluation; it describes architectural tradeoffs readers should expect when comparing sequence modules.

Why It Still Matters

Selective state-space modules remain a practical reference for long-context and hybrid designs because they show how recurrent state updates can carry sequence information with different cost and inductive bias than attention-only stacks.

Tags

References

  1. Gu, Albert, and Tri Dao. "Mamba: Linear-Time Sequence Modeling with Selective State Spaces." arXiv, 2023, https://arxiv.org/abs/2312.00752.

Looped Transformers

Previous Page

Multi-Token Prediction

Next Page

On this page

What It IsWhy It ExistsHow It WorksMath Or Compute SchemaCompared To Nearby ModulesExample ArchitecturesLimitations And TradeoffsWhy It Still MattersRelatedTagsReferences