# Mamba and State Space Models, Simply Explained

#58 Jul 7, 2025 Architecture

https://www.aibutsimple.com/p/mamba-and-state-space-models-simply-explained

Transformer-based LLMs have become the dominant architecture for text processing and NLP in recent years. They are behind some of the most widely used LLMs, like ChatGPT, Claude, and Gemini.

While the transformer brought on a huge wave of development and growth, researchers have started developing new architectures that might outperform it.

- Don’t understand transformers? Learn about transformers or brush up your knowledge in our previous issue, “[Transformers and the Attention Mechanism](https://www.aibutsimple.com/p/transformers-and-the-attention-mechanism?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=mamba-and-state-space-models-simply-explained).”

One of these new architectures is called Mamba, and it is what we call a “state-space” model (SSM).

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/3d30e68e-c7e2-4c5d-9178-4394f1b134e5/Mamba_Cover__1_.png)

Mamba was introduced by researchers at Carnegie Mellon University in their 2023 paper: “[Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://arxiv.org/abs/2312.00752?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=mamba-and-state-space-models-simply-explained).”

State space models have been around for quite a while (introduced in the 1960s), but why specifically were state space models introduced to NLP? The reason is that state space models like Mamba solve key problems with transformers and RNNs.

Essentially, transformers have great attention mechanisms that model long-range dependencies well, but their computation and memory complexity scales quadratically with sequence length—they are too computationally expensive.

RNNs scale better, but they have small, fixed-size hidden states that do not allow for good comprehension of long sequences.

Due to these constraints, there is a need for an architecture that could be both memory-efficient and manage long-range dependencies well.

State Space Models (SSMs) for NLP were the researchers' answer to this problem. Mamba is one of the most standout SSMs that have shown fantastic performance.

### The Issue with Transformers

State space models like Mamba solve a key problem with transformers—they are quite memory intensive and, although fairly scalable, scale quadratically with input sequence length.

Let’s first do a short recap of transformers, then explore the scaling issue. Transformers take sequences of tokens as input.

Through attention, they have the ability to look back at any of the earlier tokens in the sequence to derive its representation. This gives it a fantastic ability to understand context and longer sequences.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/2b25f072-cc33-4add-b59d-3cd88ef2b815/Transformer_Tokens__1_.png)

But how do transformers scale quadratically? If you recall, transformers have an attention pattern that is in the shape of a 2D grid.

It creates a matrix comparing each token with every token that came before. The weights in the matrix are determined by how relevant the token pairs are to one another.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/3925a048-eb13-41dc-a69f-1fead73fb1c9/Transformer_Attention_Grid__1_.png)

For self-attention, we calculate attention scores between each token and every other token, giving us an attention pattern of size n×n, where n is the size of the input sequence in tokens.

This means the grid will have n2 entries, scaling quadratically with the input sequence length for both training and inference. At least training is easily parallelizable, so it can be finished quicker with more GPUs.

Even though transformers have this key issue, they are still used, since they still provide a fantastic understanding of context and long-range sequences.

### The Issue with RNNs

As a quick recap, Recurrent Neural Networks (RNN) are neural networks made for sequence modelling.

RNNs take two inputs at every timestep: the current input at time t and the previous hidden state at time t-1. With these inputs, they are able to produce an output and the next hidden state (at time t).

At every timestep, they iteratively pass information from previous steps to the next.

Since the RNN only needs the previous hidden state and current input for output generation, we don’t need to recalculate all previous hidden states like the Transformer does.

With this in mind, RNNs have inference that scales linearly—much better than quadratic!

So, RNNs do solve the transformer’s scaling problem, but the issue with them is that they lack the ability to model longer-range dependencies with their fixed-size hidden states. They tend to “forget” information over time.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/b94e09a1-390f-41a4-b24a-670b9798a4a3/RNN_sequence_forgetting__2_.png)

With the focus shifting to longer and longer sequences and context windows, it’s essential that our text models can handle them. And so, SSMs were applied to NLP, which has long dependency handling and good scalability. But before we get into state space models, what even is a state space?

### What is a State Space?

A state space is the set of all possible states a system can be in. But what is a state, and what is a system?

Imagine we are navigating through a 2D grid world—this grid navigation problem is our system we want to analyze. The system is essentially the environment and the rules governing the process we’re analyzing.

Locations of this grid are the states, and the state space is made up of all possible states, or locations.

For a more complicated example, maybe our state has variables representing our x and y coordinates in the grid but also contains distances to enemies.

With all this information, we can pack it into a vector called a state vector. State vectors are useful since vector embeddings in language are also used to describe the state of an input sequence.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/89c0c14d-08c4-447a-98ab-386b81b9dc34/State_Vector.png)

In terms of neural networks, the “state” of a system is typically its hidden state. For instance, in RNNs, the state space is the set of all possible hidden state ht vectors.

State space models use equations and matrices to track the state space. SSMs track inputs, outputs, and states in mathematical equations—we call this the state space representation.

### State Space Models

SSMs mathematically represent these states and produce predictions of what the next state could be depending on some input.

At some timestep t, SSMs map an input sequence x(t) to a latent state/hidden state representation h(t) and produce a predicted output sequence y(t).

SSMs are represented by two core mathematical equations: the state equation and the output equation.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/75d38a48-8ee2-432d-ba4d-7dacbcacf98c/State_Vector__1080_x_530_px___1_.png)

The state equation describes how the state h(t) changes through matrix A and how the input x(t) influences the state using matrix B.

The output equation describes how the state transitions to the output using matrix C and how the input directly affects the output through matrix D.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/bd84d742-d71c-4b73-967c-e7558ded3b4c/Explaining_SSM_Equations.png)

We can visualize these two equations with the following diagram. This is commonly how state space models are visualized as an architecture.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/95038ff3-d743-4a8a-8e27-dcd336757d04/Copy_of_State_Vector__1080_x_530_px_.png)

Above, all four matrices, A, B, C, and D, are parameters since they are learnable and change their values during training.

If we look closely, the input can flow directly to the output without going through the rest of the computations through matrix D—this is like a modified skip connection.

The only difference between this and an actual skip connection is that the input has to be transformed through a matrix first before reaching the output.

### Continuous to Discrete

Since SSMs take in continuous inputs and produce continuous sequences as output, the model needs to be discretized for discrete applications (like processing sequences of tokens).

To do so, we perform something called the “zero-order hold technique.” When we receive a discrete signal or input, we repeat its numerical value until we get a new discrete signal. This creates a continuous “staircase” signal the SSM can use:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/e6741d10-1ee0-4b9f-a2cb-13f18eb3768b/Discretize.png)

How long we hold the value is determined by a learnable parameter called the step size, and we represent it using the uppercase delta symbol ∆.

Using our newly continuous input, we generate a continuous output, and then we sample values with interval ∆ to obtain a discretized output.

Mathematically, we can apply the zero-order hold technique using the exp function or ex along with the identity matrix I and the step size ∆.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/8b74b64c-5707-48c4-850f-dcdb4d203d06/image.png)

This discretizes matrix A and matrix B for use in discrete applications.

The discretized SSM now maps input sequences to output sequences instead of input functions to output functions. The discrete SSM’s state and output equations are shown below with our new discretized matrices.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/873c30dc-541f-480f-8376-41b4fe5b37d3/Discrete_vs_Continuous.png)

A quick note here: in simple SSM implementations, oftentimes we will set D to zero to ignore the Dxt term. If we disable the Dxt term, we prevent the skip connection from bypassing the model’s state or memory.

### SSMs as RNNs

Representing SSMs discretely translates to also representing them as RNNs! Their underlying structures are very similar.

In equation form, we use the hidden state of the previous timestep and the input of the current timestep to obtain the output and state of the current timestep. It’s clearer with this image:

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/faa2ab92-5e21-4df2-9cda-28e010242e43/SSMs_as_RNNs__3_.png)

Discrete SSMs can kind of be viewed as a generalization of RNNs but with better stability and long-range memory. In most modern comparisons, SSMs outperform traditional RNNs like LSTMs and GRUs in both performance and scalability.

Discrete SSMs can also be visualized using convolutions, and both representations have different benefits. Recurrent SSMs have efficient inference but slower training (no parallelization), while convolutional SSMs have slow inference (fixed kernel size) but faster training with parallelization.

We go into convolutional SSMs in detail for [supporters only](https://www.aibutsimple.com/upgrade?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=mamba-and-state-space-models-simply-explained), so if you want to learn more, consider [upgrading](https://www.aibutsimple.com/upgrade?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=mamba-and-state-space-models-simply-explained)!

---

---

### Mamba—The Selective SSM (S6)

Mamba is a popular SSM that has shown exceptional performance, matching and even sometimes exceeding the performance of transformers of the same size. They also provide fast inference, training, and unbounded context (issue with the convolutional representation).

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/35597b3b-f7f4-4ae9-8b44-aa1ec2dc8463/image.png)

Sample Mamba Architecture. The SSM is shown in the middle by “Selective SSM”

Mamba is also referred to as the selective SSM or S6. The popularity of Mamba boils down to two main contributions that solve a set of disadvantages with SSMs.

The first contribution is a selective scan algorithm, which allows the model to filter information.

The second is a hardware algorithm that stores intermediate computations efficiently through fancy techniques called parallel scan and kernel fusion.

These contributions solve the issue that previous SSMs had difficulty focusing on or ignoring certain inputs.

They struggle since recurrent or convolutional representations of SSMs are linear time invariant. This is just a fancy way of saying how the matrices A, B, and C are the same for every token the SSM generates.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/3adf4c16-ea77-43e7-9745-2601389bdac1/Problem_Const_Matrices__1_.png)

As a result, SSMs have problems with making decisions based on the actual content of the input sequence, since they process inputs uniformly—without adapting to the meaning of every token. This is a problem, as we want the SSM to properly interpret the input.

Reasoning using the different meanings of tokens is relatively easy for transformers to dynamically change their attention based on the input sequence. They can selectively “attend” to different tokens in a sequence.

#### Selective Scan

Mamba aims to place itself in between the recurrent representation of an SSM (small state, efficient) and a transformer (huge state, attention matrix, powerful).

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/8fbadd7e-1d78-4d29-ab01-1de15dba49b5/Problem_Const_Matrices__1080_x_500_px_.png)

It does this by compressing information selectively into the state. Inputs often have unnecessary information like [stop words](https://en.wikipedia.org/wiki/Stop_word?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=mamba-and-state-space-models-simply-explained), which we can filter out.

To selectively compress information, we need the parameters to be dependent on the input.

Mamba makes matrices B and C dependent on the input by incorporating the sequence length and batch size of the input into the dimensions. It also does this for the step size (∆).

You might be wondering why we aren’t making matrix A dependent on the input. Matrix A stays static since we want the state itself to be unchanged, but the way it is influenced through B and C to be dynamic.

These matrices selectively choose what to keep in the hidden state and what to ignore since they are now dependent on the input.

- For example, a smaller step size (∆) focuses more on the previous context, while a larger step size focuses on the input tokens more than the context.

Since these matrices are now dynamic, they can’t be computed using the convolution representation since it has a fixed kernel.

However, the recurrent representation does not have parallelization capabilities. To solve this, Mamba uses the [parallel scan (prefix sum)](https://developer.nvidia.com/gpugems/gpugems3/part-vi-gpu-computing/chapter-39-parallel-prefix-sum-scan-cuda?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=mamba-and-state-space-models-simply-explained) algorithm to parallelize.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/1fe534ca-7002-4261-ada8-8be4f6127adf/image.png)

High-Level Diagram of an Efficient Parallel Scan Algorithm

Explained simply, parallel scan uses the associative property of an operator like addition to calculate the sequences in parts and iteratively combine them. As it’s associative, the order of the sequence doesn’t matter.

The usage of the parallel scan algorithm along with dynamic matrices B and C is the core of Mamba’s “selective scan” algorithm.

#### Hardware Algorithm

Mamba also performs a hardware algorithm to increase the speed at which the model can be processed through the GPU. GPUs have a large amount of DRAM, which is a slower memory, and a small amount of SRAM, which is a quick and efficient memory.

When we copy information between DRAM and SRAM, this takes quite some time and is slow. Mamba limits the number of times we need to go from DRAM to SRAM and from SRAM to DRAM.

![](https://media.beehiiv.com/cdn-cgi/image/fit=scale-down,width=1200,format=auto,onerror=redirect,quality=80/uploads/asset/file/8e30bb7e-bdd1-412a-97df-bb8d63cb4bfb/Copy_of_Problem_Const_Matrices__1080_x_500_px_.png)

It does this through a process called “kernel fusion,” making the GPU perform computations without writing intermediate results, shortening the number of steps to writing the final results.

Mamba can be implemented as a block, similar to how we represent a transformer decoder block. Like the decoder, multiple Mamba blocks can be stacked, where the output of one serves as the input to the next.

Mamba can then be combined with other techniques or architectures to achieve even better performance.

---

Here’s a special thanks to our [biggest supporters](https://www.aibutsimple.com/upgrade?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=mamba-and-state-space-models-simply-explained):

Sushant Waidande

---

If you enjoy our content, consider supporting us so we can keep doing what we do. Please share this with a friend!

Want to reach 50000+ ML engineers? Let’s work together:

[WORK WITH US 📈](https://www.passionfroot.me/aibutsimple?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=mamba-and-state-space-models-simply-explained)

Feedback, inquiries? Send us an email at [edwin@aibutsimple.com](mailto:edwin@aibutsimple.com).

If you like, you can also donate to our team to push out better newsletters every week!

[Support us ☕](https://ko-fi.com/aibutsimple?utm_source=newsletter.aibutsimple.com&utm_medium=referral&utm_campaign=mamba-and-state-space-models-simply-explained)

---

That’s it for this week’s issue of AI, but simple. See you next week!

—AI, but simple team