State Space Models and Mamba: Is There Life Beyond Attention?
How state space models and Mamba process sequences with a fixed-size memory instead of attention, what they gain in long-context efficiency, and where transformers still win.
Ask a transformer to read a book and it does something slightly absurd: for every new token it looks back at every previous token, and it keeps a record of all of them in memory. That brute-force lookup is the reason transformers are so good at recalling a detail from page three, and it is also the reason long contexts are expensive. State space models (SSMs), and Mamba in particular, make the opposite bet. They read a sequence once, left to right, and squeeze everything they have seen into a fixed-size summary. This article explains how that works, why it is cheap, what it gives up, and why many recent designs mix both ideas.
The Cost of Attention
In the transformer, self-attention compares each token with every other token. Compute for a full sequence grows quadratically with length. During generation, the model avoids recomputation by keeping a KV cache, but that cache grows linearly with the number of tokens and must be read at every step.
So two things scale badly with context length:
- Training and prompt processing cost grows quadratically with sequence length.
- Generation needs memory that grows with every token, and memory bandwidth to read it back.
Engineers have attacked this with sparse patterns, sliding windows, grouped-query attention and clever kernels. State space models attack the root instead: what if the model did not keep the whole history at all?
Recurrence, Revisited
Before transformers, recurrent neural networks (RNNs and LSTMs) handled sequences by carrying a hidden state forward. At each step the network takes the new input and the old state and produces a new state:
state_t = f(state_{t-1}, input_t)
output_t = g(state_t)
The appeal is obvious. The state has a fixed size, so each step costs the same no matter how long the sequence is, and generation needs no growing cache. The problems are equally well known. Classic RNNs are hard to train because gradients vanish or explode across many steps, and they are sequential by nature: you cannot compute step 1000 until you have computed step 999, which wastes the parallelism of modern GPUs during training. That is exactly the weakness the transformer was designed to remove.
State space models are best understood as recurrence rebuilt on better mathematics, with training made parallel.
What a State Space Model Is
The idea comes from control theory. A continuous linear system keeps a hidden state h that evolves under the input x and emits an output y:
h'(t) = A h(t) + B x(t)
y(t) = C h(t)
Here A, B and C are learned matrices. To apply this to text or audio, which are discrete, the system is discretized with a step size, giving a recurrence that looks very much like an RNN:
h_t = A_bar h_{t-1} + B_bar x_t
y_t = C h_t
Two properties make this interesting.
First, the recurrence is linear (there is no nonlinearity between steps). A linear recurrence can be unrolled into a convolution: the output is the input convolved with a long kernel built from A, B and C. During training you can use the convolutional view and compute the whole sequence in parallel. During inference you switch to the recurrent view and take constant-time, constant-memory steps. One model, two modes.
Second, the choice and initialization of A matters enormously. Earlier work on structured SSMs (the S4 family) showed that carefully structured A matrices let the state remember information over very long ranges, which is where plain RNNs failed. The state is effectively a compressed, continually updated representation of the history, with some dimensions forgetting quickly and others slowly.
Where Classic SSMs Fell Short
Early SSMs were time-invariant: A, B and C were the same for every token. That is what makes the convolution trick possible, but it has a real cost. A fixed filter treats every input the same way. It cannot decide “this token is irrelevant, skip it” or “this token is important, store it firmly.”
Language needs exactly that kind of content-based decision. Consider the sentence “The key to the cabinet is in the drawer.” A model has to hold onto the specific noun that matters and let filler words wash past. Time-invariant SSMs did well on continuous signals such as audio but lagged behind transformers on text, in particular on tasks that require looking up something specific from earlier in the context.
Mamba: Making the State Selective
Mamba’s central idea is selectivity. Instead of fixed parameters, the step size and the B and C matrices become functions of the current input token. In pseudo-code:
# conceptual sketch, not an optimized implementation
def selective_ssm(x, A, W_B, W_C, W_delta):
h = zeros(state_size)
outputs = []
for x_t in x:
delta_t = softplus(W_delta @ x_t) # how much to move the state
B_t = W_B @ x_t # what to write into the state
C_t = W_C @ x_t # what to read out of the state
A_bar = exp(delta_t * A) # how much of the old state to keep
h = A_bar * h + delta_t * B_t * x_t
outputs.append(C_t @ h)
return outputs
The input now controls the gate. A large step size effectively resets the state toward the current token; a tiny one lets the token pass without disturbing memory. The model can therefore ignore noise, hold on to key facts, and reset at boundaries such as the start of a new document in a packed training batch.
The price is that the convolution trick no longer applies, because the parameters change at each step. Mamba’s authors solved this with a hardware-aware parallel scan: the linear recurrence is still associative, so it can be computed with a parallel prefix algorithm, and the implementation is carefully fused so that the large expanded state lives in fast on-chip GPU memory rather than being written out to slow memory. The engineering here is as important as the math, which is a recurring theme in efficient architectures.
A Mamba block also looks different from a transformer block. It typically contains a linear projection that expands the width, a short causal convolution, the selective SSM, a gating branch, and a projection back. There is no attention and no separate feed-forward layer in the classic design. Blocks are stacked with normalization and residual connections, just as in a transformer.
What You Gain
Constant-size state at inference. Generating the ten-thousandth token costs about the same as generating the tenth. There is no KV cache growing with context, which frees memory and makes serving very long sequences far more practical.
Linear scaling in sequence length. Processing time grows linearly rather than quadratically. For domains with extremely long sequences, such as DNA, long audio and some time series, this is a major advantage.
Simple, predictable memory. Capacity planning gets easier when each active request holds a fixed-size state instead of a cache that depends on how long the conversation becomes.
What You Give Up
The fixed-size state is both the feature and the limitation. Whatever the model wants to remember must fit in that state, so information is compressed lossily. A transformer, by contrast, keeps every token’s keys and values and can retrieve any of them exactly.
This shows up in practice in several ways:
- Exact recall. Tasks like copying a long string, finding one specific entry in a large table, or retrieving an arbitrary earlier fact are where pure SSMs tend to struggle relative to attention, because attention can address memory directly while an SSM must have decided in advance to store it.
- In-context learning. Transformers are strong at picking up patterns from examples in the prompt. Pure SSMs have been reported to be weaker on some of these behaviors, particularly as tasks demand more precise retrieval.
- Ecosystem maturity. Transformers have years of tooling behind them: optimized kernels, quantization recipes, serving engines, fine-tuning libraries. SSM support is improving but is thinner, so the theoretical efficiency does not always translate into easy deployment.
None of these is a flaw in the math so much as a consequence of the design: a bounded summary cannot be a perfect index of an unbounded history.
Hybrids: Using Both
Because attention and SSM layers have complementary strengths, a popular direction is to interleave them. A hybrid model might use mostly SSM layers for cheap long-range processing and sprinkle in a smaller number of attention layers to provide exact lookup. Some designs also combine them with mixture-of-experts layers for capacity.
The reasoning is practical. Most layers in a network do not need random access to the whole history; a few well-placed attention layers can supply it. The KV cache then shrinks, because only the attention layers need one, while retrieval quality stays close to a full transformer. Several research and open-weight models have taken this hybrid route, which suggests the future is likely a blend rather than a replacement.
Using SSMs in Practice
If you want to experiment, the open-source ecosystem gives you a straightforward path. A minimal example with the Hugging Face transformers library, assuming a Mamba-style checkpoint you have chosen and that the required packages are installed:
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
model_id = "your-org/your-mamba-checkpoint" # substitute a real checkpoint
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id, torch_dtype=torch.bfloat16, device_map="auto"
)
inputs = tok("State space models work by", return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=60)
print(tok.decode(out[0], skip_special_tokens=True))
Things worth checking before you commit to an SSM or hybrid for a real project:
- Your workload’s retrieval needs. If you must reliably pull exact facts from long documents, test on that task directly. Pair the model with retrieval rather than assuming the state will hold everything.
- Kernel and hardware support. The fast path relies on custom fused kernels. Without them, throughput can be disappointing.
- Serving stack compatibility. Check that your inference engine supports the architecture, including batching and state management.
- Evaluation on your own data. Public leaderboards rarely reflect long-context behavior on your documents.
When Not to Use Them
Do not reach for an SSM just because it is novel. A standard transformer is the safer choice when your context is short to moderate (the quadratic cost does not hurt there), when you depend on mature tooling, when your task is retrieval-heavy, or when you need the widest choice of pretrained checkpoints. SSMs and hybrids earn their keep when sequences are very long, memory per request is the bottleneck, or the data is naturally continuous.
Common Mistakes
- Assuming “linear scaling” means “always faster.” For short sequences, a well-optimized attention kernel can win, and constant factors matter.
- Expecting an SSM to behave like a transformer on retrieval tests and being surprised by failures.
- Ignoring the hybrid option and treating the choice as binary.
- Comparing models at different sizes or training budgets and attributing the difference to architecture alone.
- Forgetting that the fixed state bounds how much can be remembered, regardless of how long the input is.
Key Takeaways
- State space models carry a fixed-size hidden state through a sequence, so inference cost per token and memory stay constant instead of growing with context like a KV cache.
- They descend from recurrent networks but use structured linear recurrences that can be trained in parallel, fixing the old RNN training problems.
- Mamba’s selectivity makes the state parameters depend on the input, letting the model decide what to remember and what to ignore.
- The bounded state is lossy, so pure SSMs tend to be weaker than attention at exact recall and some in-context learning behaviors.
- Hybrids that mix SSM layers with a few attention layers are a practical compromise and a strong current direction.
- Choose based on your workload and tooling: transformers for mature, retrieval-heavy, short-to-medium contexts; SSMs or hybrids for very long sequences and tight memory budgets.