Research Briefings: Mamba - 3

To understand Mamba, you need to understand the core tradeoff in sequence modeling. Transformers use self-attention, which stores all past tokens in a key-value (KV) cache. This gives them perfect recall but makes compute grow quadratically with sequence length, and the KV cache grows linearly, eating memory at long contexts. State space models take the opposite approach. They compress all past information into a fixed-size hidden state that gets updated at each timestep. This gives them linear scaling with sequence length which is great for efficiency but that fixed state has to do all the work of “remembering,” which is inherently lossy.

The SSM recurrence looks something like:

where h is the hidden state, A is a transition matrix controlling how the state evolves, B controls how new input enters the state, and a separate matrix C reads out the output. The design of A, B, and the discretization scheme (how you go from continuous to discrete time) determines the model's expressiveness.

Since the state can’t grow, the question becomes:

how do you make that fixed box more powerful?

This is exactly where the Mamba lineage diverges.

• Mamba-1 (late 2023) made these matrices input-dependent (”selective”), meaning A, B, C change based on the current token which is a huge leap that let SSMs actually compete with Transformers on language modeling.

• Mamba-2 (mid-2024) simplified the SSM mechanism to leverage GPU tensor cores better, achieving 2–8x faster training than Mamba-1. The tradeoff was that it reduced the transition matrix A from a diagonal matrix to a scalar-times-identity, making the recurrence simpler but less expressive. This was a deliberate bet that training speed was the primary bottleneck.

• Mamba-3 The LLM landscape has shifted since Mamba-2. Post-training methods like reinforcement learning with verifiable rewards require massive amounts of generated rollouts, and agentic workflows have pushed inference demand enormously. Yet Mamba-2’s simplified recurrence left the decode step memory-bound that is the GPU spends most of its time moving data rather than computing.

  Mamba-3 asks:

  what would an SSM designed with inference in mind look like?

The authors identified three “levers” to pull: make the recurrence more expressive, use a richer transition matrix, and add more parallel work per step, all without significantly increasing decode latency.

1. Exponential-Trapezoidal Discretization

Mamba-3 introduces a new discretization scheme that makes the recurrence formula more expressive. In classical numerical methods, “discretization” is how you convert a continuous-time differential equation into discrete update steps. Mamba-2 used a simpler scheme; the new exponential-trapezoidal approach implicitly applies a convolution-like operation on the input to the hidden state. This is significant because it actually allowed the team to remove the short causal convolution that had been a staple of Mamba-1 and Mamba-2 (and most linear models). That short conv was originally needed for induction-style retrieval capabilities, but the new recurrence, combined with simple biases on the B and C matrices, provides equivalent functionality intrinsically.

2. Complex-Valued State Tracking

Mamba-3 expands state-tracking capabilities by modeling a complex-valued SSM system. Instead of real-valued hidden states, the transition matrix operates in the complex plane. The clever implementation trick is that they express this via RoPE (Rotary Position Embeddings) interpreting complex transitions as rotations. This avoids having to rewrite kernels from scratch for complex arithmetic while giving the state richer dynamics to represent oscillatory or periodic patterns.

3. Multi-Input, Multi-Output (MIMO) SSMs

Instead of the standard single-input, single-output (SISO) SSM, Mamba-3 offers a MIMO variant that runs multiple SSMs in parallel. The key insight here is about the compute-vs-memory tradeoff during inference: current linear models use lots of GPU tensor cores for fast training, but during decoding, each timestep requires so little compute that the hardware sits idle most of the time. MIMO adds more FLOPs per timestep — using those idle cores — so decode latency stays roughly constant even though the model is doing more useful work. The cost shows up in training (which becomes slower) but not in inference. This is an elegant exploitation of the asymmetry between training and inference compute profiles.

Beyond the SSM core, Mamba-3 also modernizes its surrounding architecture. They added QK-normalization (called “BCNorm” in SSM terms) which stabilizes training, bringing Mamba-3 in line with contemporary transformer models. They also switched to interleaved MLP layers following standard transformer conventions. And as mentioned, the short causal convolution is gone.

Let me visualize the practical outcome: how Mamba-3 fits into the bigger picture of hybrid architectures, which the authors predict will be the dominant paradigm going forward.

This hybrid layout captures the Mamba-3 team’s prediction: linear layers will predominantly be used alongside global self-attention layers, the SSM handles the bulk of sequence processing at O(1) memory cost, while a few strategically placed attention layers handle exact retrieval tasks that need the full KV cache.

Mamba-3 is a bet that the age of inference has arrived. Rather than simplifying the SSM to train faster (as Mamba-2 did), it enriches the recurrence via trapezoidal discretization, complex-valued states, and MIMO to make each decode step do more useful work using GPU cycles that were previously idle. The result is a model that is both more accurate and faster at generation than its predecessors, with open-sourced kernels built across three levels of GPU abstraction (Triton, TileLang, CuTe DSL) for maximum hardware performance. You can find the paper at arxiv.org/abs/2603.15569 and the code at the mamba-ssm GitHub repo.