Mamba: Linear-Time Sequence Modeling with Selective State Spaces

[Paper] introduces an algorithm called Mamba — a non-attention-based approach to reasoning over sequential data.

How did we end up with an algorithm like this?

Let's take a brief look at the history.

It is fair to say that virtually every promising modern LLM traces its roots back to Google's Transformer.

The Transformer looks impressive at first glance, but when you strip it down, it is ultimately just a prediction model that predicts the embedding of the next token.

In other words, it performs Next Sequence Prediction — looking at the embeddings of previous states to predict the embedding of the next state — and in doing so, boosts performance through Attention, which captures the relationships among those previous embeddings.

So why did Attention become necessary in the first place?

Algorithms that take previous states into account have existed for a long time.

RNN is the most basic incarnation of this idea.

It uses the output of the previous state alongside the current input in order to inform the current prediction.

RNN succeeds in capturing temporal structure through this simple mechanism.

The problem is that its sequential nature gives rise to several issues.

The most prominent one is its inability to handle long sequences properly.

As the structure implies, because the output of each state feeds into the input of the next, as sequences grow longer,

earlier information fails to propagate effectively all the way to the final state.

LSTM was introduced to address this problem directly.

Put simply, it keeps the core RNN mechanism intact while adding a Cell Container that decides what to pass forward to the next state.

LSTM was a remarkably effective algorithm and appeared to solve many of the sequence problems that existed at the time.

However, a fundamental limitation remained in these sequence models:

to infer the current state X_t, you must already have the result of state X_{t-1}.

So while LSTM does have the ability to handle long sequences effectively …

doing so requires sequentially computing every state and passing that result forward — for every single sequence, every single time.

This means memory can never be parallelized,
and as sequences grow longer … and models grow larger … training becomes extremely slow.

It was against this backdrop that the Attention Mechanism emerged in 2015 [link],

followed shortly by Google's Transformer in 2017, which resolved the core challenges of processing sequential data.

The key to this solution is the Key, Query, and Value matrices at the heart of the Attention Mechanism.

The parameters that the Transformer learns live in these three matrices.

Matrix multiplication against the incoming embedded vectors enables inference at remarkably high speed.

Because relationships to prior information (Attention Scores) are expressed through these matrix operations,

one of the core problems with RNN-based algorithms — the need to compute state t-1 explicitly — is eliminated.

That is, sequential inference is no longer required to express temporal or sequential relationships!

This property made the weight matrices amenable to parallel training,
gave the model strong performance on long sequences,
and because it fundamentally relies on matrix operations, it is considerably fast relative to its parameter count.

On the strength of these capabilities, post-Transformer models pursued

bigger architectures … more parameters … more attention heads …

and threw ever-larger amounts of GPU resources at the problem.

Then an event occurred that poured fuel on the "Bigger is Better" trend.

A 2022 paper (again from Google) reported that models crossing the 10^22 parameter threshold

exhibited sudden, dramatic performance improvements.

This phenomenon, dubbed Emergent Ability, describes how next-word-prediction models, as they scale up,

acquire capabilities that go far beyond what we expected — surpassing mere next-token prediction and developing abilities we cannot fully explain.

This can also be seen in a 2017 OpenAI blog post:

even though the model was trained on an entirely different task from sentiment analysis,

certain neurons produced outputs that took the form of sentiment analysis results.

After that, the Transformer came to dominate the landscape of sequential problem domains … but was it without flaws?

Not quite.

The parameters encoding the relationships among all tokens (more precisely, their embedded vectors) became an enormously heavy tool.

To compute the Attention Score between the word "person" and the word "animal," the Transformer must load all of its parameters …

meaning even the simplest tasks carry a tremendous computational cost.

And as the model grows larger, just loading it onto a GPU becomes a problem in itself.

In short, what made it great for training becomes a serious liability at inference time.

It is in this context that State Space Model-based research leveraging non-attention approaches began to emerge.

---

Memorizing what to remember (pass forward) vs. memorizing the rule for reconstructing memory

What RNN parameters learn is

"how to bring memories from the past into the present."

As noted, RNN processes an entire sequence and relies on an unstructured, sequential information-passing mechanism.

This means to retrieve information from state t-1 while at the current state t, you must trace all the way back to the first state t_0.

As we saw, the Transformer addresses this by encoding Attention information over the full sequence,
which handles long sequences effectively — but at the cost of being very large and heavy.

If we could define a structured way to recover a memory,

then whenever we need state t-1, we would only need a suitable rule to reconstruct it — nothing more.

New approaches emerged to tackle this problem.

State Space Models had already been widely validated across fields such as control theory and computational neuroscience, but they had not made a strong showing in the deep learning domain for long sequences that we care about here.

The core idea of SSMs is that any state x we wish to find can be expressed as a discretized differential equation of a consistent form,

and that matrix A — the central component of this differential equation — exhibits Linear Time-Invariant (LTI) properties throughout.

In other words, every state that makes up a sequence can be represented by a single linear system.

The fact that a sequence's states are expressed in such an explicitly structured form is the core value proposition of SSMs.

Accordingly, in traditional models, matrix A holds the same fixed value throughout the model.

However, this style of modeling fails to map states that exhibit high-dimensional, nonlinear characteristics adequately.

So why not let matrix A vary with time t?

But if it varies and is also learnable, it becomes essentially equivalent to an RNN (the two share nearly the same properties),

and building a structured latent space over all states via a well-defined equation becomes difficult.

HiPPO demonstrates that structured matrices can effectively resolve this problem.

In traditional approaches, matrix A is a "fixed matrix";
in unconstrained deep learning methods like RNNs, it is an "unstructured matrix."

HiPPO constructs matrix A using a "structured matrix",

and shows that this allows sequence states to be represented effectively.

As discussed, we want to express some state as a discretized, time-invariant differential equation,

while being able to describe the original state faithfully,

and having the matrix A required for this take a highly structured form.

This is precisely what HiPPO is — one of the building blocks in the knowledge structure that underlies Mamba.

As the figure shows, HiPPO can be understood as a memory storage scheme consisting of three steps: [Projection] → [Coefficient Estimation] → [Discretization].

Once projection and coefficient estimation are performed, the approximation above holds for all time T.

Here φ(T) denotes an orthogonal polynomial basis (let's focus on the Legendre polynomials).

The original state can be approximated using these orthogonal polynomials and their corresponding coefficients.

Since the orthogonal polynomial basis itself is fixed,

the information that actually needs to be stored for reconstruction is the coefficient function c(T).

HiPPO (through some fairly involved mathematics) pre-computes this as well via a discretized differential equation.

Since directly integrating to obtain c(T) is intractable …

the change in c(T) can instead be expressed as a differential equation involving matrices A and B.

(Again, through considerable mathematical machinery …) this differential equation can be written as a discrete state equation,

and all the information C needed to estimate the original function F can be obtained in continuous form using the highly structured matrix A.

In other words, the target form S_{t+1} = Matrix A × S_t is preserved entirely within an LTI system, and through a structured matrix A,
the ability to handle long sequences effectively is demonstrated.

---

That covers a surface-level tour of SSMs.

To summarize in one line:

SSMs are a methodology for expressing the hidden state function — the kind learned by models like RNNs — through a state equation within a Linear Time-Invariant system.

In the next section, we will take a closer look at HiPPO and S4.