What are State Space Models?

If you've been keeping up with our blog, then you probably know that the biggest scaling liability in modern Transformers is the KV cache. The reason for this is the attention mechanism retaining every prior token. Keeping track of each token in generation creates a massive mountain of data that needs to be moved back and forth repeatedly. Even with the fastest hardware, a memory-bandwidth bottleneck can happen long before the compute cores hit saturation.

State Space Models (SSMs) replace the ever-growing token cache with a fixed-dimension hidden state. By compressing context recurrently rather than caching raw token representations, SSMs keep the inference memory footprint static, no matter how long the sequence runs.

The Core Mental Model: Continuous Systems to Digital Tokens

At its core, a state space model functions like a physical system tracking changes over time. Imagine an audio filter processing a continuous sound wave:  
 
Incoming signal enters the system -> system updates an internal hidden state based on the new signal and its existing conditions -> the system projects that hidden state into an observable output.

To apply this concept to digital text, continuous-time equations undergo a process called discretization. Discretization samples the continuous mathematical framework into discrete steps, aligning it with individual token intervals.

This transformation produces a system governed by clear, intuitive roles:

  • A state transition rule dictates how previous context decays or persists over time.

  • An input mapping determines how strongly the new token alters the internal state.

  • An output mapping projects the internal state into the next generated representation.

Because the internal state summarizes the sequence into a fixed-size vector, the model processes language continuously without keeping historical data in active memory.

The Dual Operational Modes: Training vs. Inference

What makes state space models unique is that they adapt their execution style depending on the job.

The Recurrent Mode for Inference

During live inference, the model behaves like a streamlined recurrent loop: it takes in one token at a time, updates its fixed internal state, and generates the next word. Because it only looks at the current token and its existing state, memory usage never grows, whether a conversation is five words or 50,000 words long.

The Convolutional Mode for Training

Running step-by-step updates like inference mode during training would prevent GPUs from parallelizing work across long documents. When the internal transformation rules remain constant across time, the entire sequence of steps can be mathematically unrolled into a single, global one-dimensional convolution.

Using Fast Fourier Transforms, the training engine applies this convolution across the entire sequence at once. This enables full GPU hardware saturation during training runs without the sequential lag that is typical of older recurrent networks.

Selective State Spaces and Hardware-Aware Execution

Early state space models had a major blind spot: they treated every token identically. A comma or period decayed at the exact same rate as a critical entity or instruction. Because the transformation math was static across the entire document, the model couldn't selectively filter out conversational noise or reset its attention for a new topic.

Selective State Spaces (like Mamba) solved this by making the update rules dynamic. The model inspects the current token and decides on the fly what to write into memory, what to discard, and how fast past context should fade.

The catch? Making updates dependent on the input token breaks the global convolution trick that made training fast in the first place.

Modern system architectures recover that speed through clever hardware engineering. Instead of repeatedly moving intermediate states back and forth across off-chip VRAM, the runtime fuses the entire scan loop into ultra-fast, on-chip SRAM. Paired with parallel prefix scans, which are tree-based algorithms that calculate sequential updates across parallel GPU thread blocks, these models retain linear scaling during training without ever needing static convolutions.

Architectural Trade-offs and Production Realities

While state space models solve the memory bloat of Transformers, they don’t render standard attention obsolete. In production, the trade-off comes down to lossless memory versus predictable scaling.

  • Flat Inference Latency: Eliminating the dynamic KV cache removes memory-bandwidth bottlenecks. Systems can serve significantly larger batch sizes on the same hardware, while generation speeds remain consistent across deep context windows.

  • Efficient Context Scaling: Linear computational complexity enables processing input documents of several hundred thousand tokens without triggering memory-allocation crashes.

  • The Compression Bottleneck: Squeezing an entire document into a fixed-size internal state comes with a cost. Unlike Transformers, which preserve an exact lossless record of every prior token, an SSM continually summarizes context. When an application requires pinpoint recall, like extracting a specific invoice total, serial number, or buried legal clause from deep within a long document, pure SSMs can fail where uncompressed attention caches excel.

Because of these trade-offs, production systems are increasingly converging on hybrid architectures. Interleaving selective state space layers with occasional standard attention layers delivers the best of both worlds: the low memory footprint and blazing decode speeds of recurrent models, paired with the recall precision of full self-attention.

Back to Main   |  Share