Mamba/SSM Basics

State Space Models offer linear-time sequence modeling with content-aware selective filtering, challenging Transformers for long-context inference.

Why This Matters

State Space Models (SSMs) provide a principled alternative to Transformers for long-sequence modeling. In production systems handling long contexts (e.g., code generation, genomic analysis), Transformer attention's quadratic cost becomes a bottleneck. Mamba achieves linear-time inference with constant-memory state, making it viable for million-token contexts where attention-based models are prohibitively expensive.

Core Idea

SSMs originate from continuous-time control theory: a latent state evolves over time driven by input, and observations are linear projections of that state. Mamba's key innovation is making the SSM parameters input-selective — the model learns to gate which information enters and exits the state, mimicking attention's ability to focus on relevant tokens without the


O(n2)O(n^2)O(n2)

cost.

Technical Details

The continuous-time SSM is defined as:

where

x(t)∈RNx(t) \in \mathbb{R}^Nx(t)∈RN

is latent state,

$u(t)$

is input, and

A∈RN×NA \in \mathbb{R}^{N \times N}A∈RN×N

,

B∈RN×1B \in \mathbb{R}^{N \times 1}B∈RN×1

,

C∈R1×NC \in \mathbb{R}^{1 \times N}C∈R1×N

. Using zero-order hold discretization with step

$\Delta$

:

The recurrent update becomes:

Mamba's selective mechanism makes

$B$

,

$C$

, and

$\Delta$

input-dependent:

The parallel scan algorithm computes this recurrence in

$O(n\log n)$

during training. Inference is

$O(1)$

O(1) per token with fixed state size N

$N$

, yielding constant-memory decoding regardless of sequence length.

How It Works

1. 
   Project input: Map token
   
   uku_kuk​
   
   to expanded dimension
   
   $D\cdot N$
   
   .


2. 
   Generate selective parameters: Compute input-dependent
   
   BkB_kBk​
   
   ,
   
   CkC_kCk​
   
   ,
   
   Δk\Delta_kΔk​
   
   from
   
   uku_kuk​
   
   .


3. 
   Discretize: Convert continuous
   
   $(A,B)$
   
   to discrete
   
   (Aˉk,Bˉk)(\bar{A}_k, \bar{B}_k)(Aˉk​,Bˉk​)
   
   using
   
   Δk\Delta_kΔk​
   
   .


4. 
   Recurrent scan: Apply parallel scan (training) or sequential update (inference) to compute hidden states
   
   xkx_kxk​
   
   .


5. 
   Output projection: Compute
   
   yk=Ckxky_k = C_k x_kyk​=Ck​xk​
   
   , then project through gating (SiLU) to output dimension.

Key Insights

• 
  Selectivity is essential: Non-selective SSMs (S4) cannot do in-context retrieval; making
  
  $B,C,\Delta$
  
  input-dependent enables content-aware filtering.


• 
  Diagonal + low-rank structure on
  
  $A$
  
  enables
  
  $O(n)$
  
  recurrence; Mamba uses diagonal
  
  $A$
  
  matrices exclusively.


• 
  Hardware-aware design: The scan kernel is IO-bound, not compute-bound — Mamba's CUDA kernel fuses discretization, scan, and output projection to minimize memory reads.


• 
  Linear decoding cost: Unlike KV-cache which grows linearly, SSM state is fixed-size
  
  $O(ND)$
  
  , making generation memory-constant.

Sources

• Gu, A. & Dao, T. "Mamba: Linear-Time Sequence Modeling with Selective State Spaces." arXiv:2312.00752 (2023). https://arxiv.org/abs/2312.00752
  


• Gu, A. & Dao, T. "Mamba-3: Improved Sequence Modeling using State Space Principles." arXiv:2603.15569 (2026). https://arxiv.org/abs/2603.15569