Mamba

The Problem

Transformers suffer from a fundamental flaw: their attention mechanism scales quadratically with sequence length. If you double the context window, the compute required quadruples, and the memory required to store the KV cache explodes.

For years, researchers tried to build "Sub-quadratic Transformers" or use Recurrent Neural Networks (RNNs) to achieve linear scaling. However, RNNs cannot be trained efficiently in parallel, and previous State Space Models (SSMs) like S4 struggled heavily with tasks requiring discrete information retrieval (like copying text or recalling specific facts).

The Idea

The authors hypothesized that previous SSMs failed because their state transitions were time-invariant—the model applied the exact same mathematical transformation regardless of what the input token actually was.

They introduced Selective State Spaces. In Mamba, the parameters that dictate how information flows into the hidden state are no longer fixed; they are a function of the input itself. The model can choose to "select" important information to remember and actively "forget" irrelevant information, much like a highly advanced LSTM, but designed for modern hardware.

How It Works

1. Selective Mechanism: The matrices $B$, $C$, and $\Delta$ (step size) in the state space equations are made data-dependent. When encountering a filler word (like "um"), the model effectively zeroes out the step size, ignoring it. When encountering a critical noun, it maximizes the step size to engrave it into the hidden state.
2. Hardware-Aware Algorithm: The primary challenge with making SSMs data-dependent is that they can no longer be computed efficiently using fast Fourier transforms (FFTs). The authors (including Tri Dao, creator of FlashAttention) solved this by designing a hardware-aware algorithm (a parallel scan) that computes the selective SSM entirely within the ultra-fast SRAM of the GPU, avoiding the HBM memory bottleneck.
3. Architecture: They simplified the standard Transformer block by combining the SSM with an MLP, creating a streamlined, uniform block architecture.

Why It Mattered

Mamba was a watershed moment. It was the first non-Transformer architecture that matched or exceeded the performance of Transformers of the same size, while boasting linear scaling in both training and inference.

Because Mamba compresses all context into a fixed-size hidden state (like an RNN), it does not require a KV cache. This means generating a token takes the exact same amount of time and memory regardless of whether the context is 10 tokens or 1,000,000 tokens.

What Came After

Mamba sparked immense interest in SSMs. It was quickly integrated into hybrid architectures (e.g., Jamba by AI21 Labs), which interleave Mamba layers with Transformer layers to get the best of both worlds (the infinite context of Mamba with the sharp recall of attention). Mamba-2 (2024) further refined the underlying math to bridge the theoretical gap between SSMs and attention. While Transformers still dominate, Mamba represents the most credible threat to their monopoly.