Mamba and State Space Models: Interview Questions and Answers

SSM & Mamba Fundamentals

Q1 [Basic] What is a State Space Model?

Q: What is a State Space Model (SSM) and how does it apply to sequence modeling?

A: A State Space Model is a mathematical framework that maps an input sequence to an output sequence through a latent hidden state. In its continuous form, it is defined by two equations: a state equation $\frac{dx}{dt} = Ax + Bu$ (describing how the hidden state evolves over time) and an output equation $y = Cx + Du$ (describing how the output is generated from the state). Here, $A$ is the state transition matrix, $B$ maps inputs to state updates, $C$ maps the state to outputs, and $D$ is an optional skip connection.

For sequence modeling on discrete data (text, audio samples), the continuous system is discretized into a recurrence: $x_t = \bar{A}x_{t-1} + \bar{B}u_t$ and $y_t = Cx_t$. This recurrence makes SSMs computationally similar to RNNs: at inference, each step requires only a fixed-size state update, giving $O(1)$ memory per step regardless of sequence length.

A key mathematical property is that the recurrence can be equivalently computed as a convolution: $y = K * u$, where $K$ is a convolutional kernel derived from the SSM parameters. This duality allows SSMs to be trained efficiently via FFT-based convolution (parallelizable across the sequence) while being deployed efficiently as a recurrence (constant memory at inference). This training-inference duality is one of SSMs’ core advantages over both Transformers (which are slow at inference for long sequences) and RNNs (which are slow to train).

Q2 [Basic] What is the HiPPO matrix and why does it matter?

Q: What is the HiPPO matrix, and what role does it play in modern SSMs?

A: HiPPO (High-order Polynomial Projection Operators, Gu et al., 2020) is a principled method for initializing the state transition matrix $A$ in an SSM so that the hidden state optimally compresses the history of the input signal. Specifically, the HiPPO framework constructs an $A$ matrix such that the hidden state at each time step contains the coefficients of the best polynomial approximation (in the Legendre basis) of the input seen so far, up to order $N$.

The practical consequence is that the HiPPO-initialized $A$ matrix gives the SSM a structured way to represent both recent and distant history, with each polynomial basis function capturing information at a different timescale. Without HiPPO initialization, SSMs initialized with random $A$ matrices suffer from the same vanishing or exploding gradient problems as vanilla RNNs — they quickly lose information from the distant past.

HiPPO is the foundational memory mechanism that enabled all subsequent SSM work including S4 and Mamba. While Mamba does not use HiPPO’s full continuous-time construction, the structured initialization philosophy — ensuring $A$ encodes a useful inductive bias for long-range memory — carries through to Mamba’s design.

Q3 [Basic] What are the core contributions of S4?

Q: What problem did S4 solve, and what were its key innovations?

A: S4 (Structured State Space Sequence Model, Gu et al., 2021) addressed the computational bottleneck of applying SSMs to long sequences. Prior SSMs had $O(n \cdot N^2)$ computation per sequence (where N is the state dimension) because computing the SSM convolution kernel required dense matrix operations involving $A$. S4’s key insight was to parameterize $A$ as a Diagonal Plus Low-Rank (DPLR) matrix. Under this structure, the SSM convolution kernel can be computed efficiently using the Cauchy kernel — reducing the complexity to $O(n \log n)$ via FFT.

S4 achieved state-of-the-art results on the Long Range Arena benchmark, a suite of tasks requiring modeling dependencies across thousands of time steps, far outperforming Transformers and previous SSMs. It demonstrated that SSMs with principled structure could compete with attention-based models on long-range tasks.

The limitation of S4 is that its parameters ($A$, $B$, $C$) are fixed across all input positions — it is a linear, time-invariant system. This means S4 processes every input the same way regardless of content. It cannot selectively attend to relevant tokens or ignore irrelevant ones, which is a fundamental capability of Transformer attention. This limitation motivated Mamba’s selective mechanism.

Q4 [Basic] What is Mamba’s selective mechanism?

Q: What does it mean that Mamba uses a “selective” state space model, and why is this important?

A: In S4 and earlier SSMs, the matrices $A$, $B$, and $C$ are time-invariant: they are fixed parameters that do not depend on the input. This means the model applies the same transformation to every input token regardless of its content — it cannot decide to “pay attention” to some tokens more than others.

Mamba (Gu & Dao, 2023) introduces selectivity by making the parameters $B$, $C$, and $\Delta$ (the discretization step size) functions of the input. Concretely, for each token $x_t$, the model computes $B_t = \text{Linear}(x_t)$, $C_t = \text{Linear}(x_t)$, and $\Delta_t = \text{softplus}(\text{Linear}(x_t))$. The $A$ matrix remains fixed (initialized with a structured diagonal form), but all other parameters now vary per token.

The intuition is similar to attention: the model can learn to “open the gate” (large $\Delta$, large state update) when it encounters a token it wants to remember, and to “close the gate” (small $\Delta$, slow state change) when it wants to preserve existing memory. This content-based selectivity is what allows Mamba to selectively filter information and reason about which parts of the context are relevant for the current prediction — a capability that S4 completely lacked.

Q5 [Basic] What are B, C, and Δ in Mamba and what does each control?

Q: What roles do the parameters B, C, and Δ play in Mamba’s selective state space model?

A: In Mamba’s SSM, $B$, $C$, and $\Delta$ are the three input-dependent parameters that implement selectivity. Each controls a different aspect of information flow through the hidden state.

$B$ is the input projection matrix that determines how much the current input $x_t$ influences the state update. A large $B_t$ entry for a particular dimension means the input strongly writes into that dimension of the state. By making $B$ input-dependent, Mamba can learn to selectively encode certain types of information into the state based on content.

$C$ is the output projection matrix that determines how the current state is read out to produce the output $y_t$. By making $C$ input-dependent, the model can selectively retrieve different aspects of the accumulated state depending on what the current token needs.

$\Delta$ (delta) is the discretization step size, which controls the timescale of the state update. Large $\Delta$ collapses the continuous SSM to a large discrete step, making the new state heavily influenced by the current input (focus on present). Small $\Delta$ corresponds to a small discrete step, meaning the state changes slowly and retains more of its previous value (focus on history). $\Delta$ is the most critical selectivity parameter: it acts as a learned gate that modulates the tradeoff between remembering the past and responding to the present input, playing a role analogous to the forget gate in an LSTM.

Q6 [Advanced] How does parallel scan enable efficient training in Mamba?

Q: Mamba uses a recurrence at its core. How does it avoid the sequential bottleneck of RNNs during training?

A: The fundamental challenge of training recurrent models is that naive recurrence is inherently sequential: $x_t$ depends on $x_{t-1}$, which depends on $x_{t-2}$, creating a chain of $n$ dependencies that cannot be parallelized across the sequence dimension. This is why RNNs are slow to train compared to Transformers, which compute all positions simultaneously.

Mamba resolves this using the parallel scan algorithm (also called parallel prefix sum). The key insight is that Mamba’s state recurrence — $x_t = \bar{A}_t x_{t-1} + \bar{B}_t u_t$ — is a linear recurrence, and linear recurrences are associative operations. Any associative operation can be computed over all prefix subsequences simultaneously using a divide-and-conquer approach with $O(\log n)$ depth and $O(n)$ total work, rather than $O(n)$ depth for the sequential approach.

Concretely, parallel scan proceeds by first combining pairs of adjacent elements, then pairs of pairs, and so on, building up all prefix reductions in $O(\log n)$ rounds of parallel computation. For a sequence of length $n$, this requires $O(n)$ total operations but only $O(\log n)$ sequential steps — enabling full GPU parallelism. Mamba implements this as an optimized CUDA kernel that runs the parallel scan entirely in fast on-chip SRAM, avoiding expensive reads from slow HBM. The result is training throughput competitive with optimized Transformer implementations despite the underlying recurrent structure.

Q7 [Advanced] How is Mamba’s hardware-aware algorithm similar to FlashAttention?

Q: What is Mamba’s hardware-aware algorithm, and how does it compare to FlashAttention in its approach?

A: Both Mamba and FlashAttention solve the same class of problem: a mathematically straightforward computation whose naive implementation is bottlenecked by memory bandwidth rather than arithmetic throughput on modern GPUs. The solution in both cases is to restructure the computation to minimize reads from slow HBM (GPU global memory) by keeping intermediate results in fast on-chip SRAM.

In standard SSM computation, the intermediate states $x_t$ for all $t$ would be materialized in HBM — requiring $O(n \cdot d_\text{state})$ memory and $O(n)$ round-trips between SRAM and HBM. Mamba’s hardware-aware algorithm instead fuses the entire forward pass into a single CUDA kernel: input $u_t$, the discretized parameters $\bar{A}_t$ and $\bar{B}_t$, and the output $y_t$ are all handled within SRAM tiles, and states are never written back to HBM. During the backward pass, states are recomputed from the stored inputs rather than loaded from HBM — trading arithmetic (cheap) for memory bandwidth (expensive), exactly as FlashAttention does.

The practical outcome is the same in both cases: the actual memory usage is much lower than the naive implementation, and the wall-clock speed is significantly faster due to reduced I/O. FlashAttention reduces Transformer attention’s $O(n^2)$ memory to $O(n)$; Mamba’s hardware-aware kernel reduces the SSM’s memory overhead similarly. Both have become the de facto standard implementations in their respective model families.

Q8 [Advanced] How does Mamba achieve linear time complexity, and what are the caveats?

Q: Mamba is described as having linear-time complexity. What does this mean precisely, and what are the trade-offs?

A: Mamba’s linear complexity claim applies primarily to inference (autoregressive generation). At each inference step, Mamba updates a fixed-size hidden state of dimension $d_\text{model} \times d_\text{state}$ using the current input — an $O(d)$ operation independent of the sequence length $n$. Generating $n$ tokens therefore requires $O(n \cdot d)$ total work and $O(d)$ memory (the state is overwritten at each step, not accumulated). By contrast, a Transformer decoder requires $O(n)$ work per step (attending to the KV cache) and $O(n \cdot d)$ memory for the KV cache — both growing with sequence length.

For training, Mamba uses the parallel scan approach, which has $O(n \log n)$ work — not strictly linear, but substantially better than $O(n^2)$ for Transformers. The memory during training is $O(n \cdot d_\text{state})$ for storing intermediate states needed for the backward pass, though the hardware-aware kernel reduces this in practice by recomputing states rather than storing them.

An important caveat is that Mamba’s fixed-size state is a double-edged sword. The $O(1)$ inference memory is only possible because the model compresses the entire history into a state of fixed dimension $d_\text{state}$. This compression is lossy: unlike a Transformer’s KV cache, which stores all past tokens exactly and can retrieve any of them with full precision, Mamba’s state is a summary that may not faithfully represent all past information. This is the fundamental trade-off — efficiency at the cost of perfect recall — and it is why Mamba tends to underperform Transformers on tasks requiring precise retrieval of specific past tokens.

Q9 [Basic] What are the core improvements in Mamba2?

Q: How does Mamba2 differ from the original Mamba, and what problem does it solve?

A: Mamba2 (Dao & Gu, 2024, “Transformers are SSMs”) introduces a key architectural simplification: it restricts the $A$ matrix to be a scalar multiple of the identity ($A = -\exp(a) \cdot I$, where $a$ is a scalar learned per channel). In the original Mamba, $A$ is a full diagonal matrix with one learned value per state dimension. This restriction in Mamba2 may seem like a loss of expressiveness, but it enables a critical theoretical connection: it allows the SSM computation to be rewritten as a structured matrix multiplication that is provably equivalent to a form of linear attention.

This equivalence — the State Space Duality (SSD) — has two practical consequences. First, Mamba2’s SSD layer is more amenable to tensor parallelism across multiple GPUs, since it can be expressed as block-diagonal matrix operations rather than custom scan kernels. This improves training throughput at large scale. Second, Mamba2 can be implemented either as an SSM (efficient for long sequences) or as a form of attention (efficient for short sequences using FlashAttention), choosing whichever is faster for the given context length.

Empirically, Mamba2 achieves comparable or better language modeling quality than Mamba at equivalent parameter counts, with higher training throughput on multi-GPU setups. It represents a convergence of the SSM and attention literature, providing a unifying framework rather than treating them as entirely separate paradigms.

Q10 [Advanced] What is State Space Duality and how does it unify SSMs with attention?

Q: What is the State Space Duality (SSD) framework introduced in Mamba2, and what does it reveal about the relationship between SSMs and attention?

A: State Space Duality (SSD) is the theoretical framework introduced in the Mamba2 paper showing that, under certain structural constraints, SSMs and a form of linear attention are two different ways to compute the same function. Specifically, when $A$ is restricted to a scalar times identity (as in Mamba2), the SSM computation $x_t = A_t x_{t-1} + B_t u_t$, $y_t = C_t^\top x_t$ can be reorganized into a matrix form $Y = M \cdot (B^\top U)$, where $M$ is a structured lower-triangular matrix (causal mask) determined by the $A$ scalars, and $B$, $C$, $U$ are the input projections. This is precisely the form of linear attention with a specific structured mask.

This duality has deep implications. Standard softmax attention uses an $n \times n$ attention matrix $A_{ij} = \exp(q_i^\top k_j) / Z$; linear attention approximates this by removing the softmax and computing attention as a matrix product, which can be reorganized into a recurrence. The SSD framework shows that SSMs like Mamba2 are a specific instantiation of this linear attention family with a particular structured (causal, decaying) attention pattern.

The practical consequence is architectural flexibility: the SSD layer can be implemented as a recurrence (for autoregressive inference, $O(1)$ memory), as a chunked algorithm (for training, processing blocks of tokens in SRAM), or in principle as attention (for very short sequences). This makes Mamba2 more hardware-friendly than Mamba’s custom selective scan and opens the door to hybrid implementations that use the most efficient algorithm for each context length.

Mamba vs Transformer

Q11 [Basic] How do Mamba and Transformer compare in efficiency?

Q: What are the key efficiency differences between Mamba and Transformer architectures?

A: The efficiency comparison between Mamba and Transformer depends on whether you are comparing training or inference, and at what sequence length.

During training, Transformer attention has $O(n^2)$ time and $O(n^2)$ memory complexity with respect to sequence length $n$ — the attention matrix is $n \times n$. Mamba’s parallel scan has $O(n \log n)$ time and $O(n \cdot d_\text{state})$ memory, which is substantially better for long sequences. In practice, for sequences up to around 2,000 tokens, highly optimized Transformer kernels (FlashAttention) can match or outpace Mamba; beyond 8,000 tokens, Mamba’s advantage becomes pronounced.

During inference (autoregressive generation), the gap is even more significant. A Transformer decoder must maintain a KV cache that grows as $O(n \cdot d_\text{model})$ — storing all past key and value vectors. Attending to this cache at each step costs $O(n)$ per token, and memory grows unboundedly with context length. Mamba at inference maintains a fixed-size hidden state of $O(d_\text{state})$ regardless of how many tokens have been processed, and each generation step costs $O(d)$ — truly constant in sequence length. This makes Mamba particularly attractive for long-context inference and streaming applications where Transformer’s KV cache memory becomes prohibitive.

Q12 [Basic] How do Mamba and Transformer differ in their inductive biases?

Q: What are the fundamental inductive bias differences between Mamba and Transformer?

A: Transformer attention has minimal sequential inductive bias: it treats all positions symmetrically and uses explicit positional encodings to inject order information. Attention can, in principle, assign equal weight to any two positions regardless of their distance. This makes Transformers flexible but requires learning positional relationships from data.

Mamba has a strong sequential inductive bias through its recurrent structure: information flows from left to right (in standard causal models), and the influence of a past token on the current output decays as a function of the learned $\Delta$ parameters and the intervening tokens. Recent tokens naturally have stronger influence on the hidden state than distant tokens, unless the model explicitly learns to preserve them. This recency bias is implicit in the architecture, unlike Transformers where recency must be learned through positional encodings and attention patterns.

Additionally, Mamba processes information through a fixed-size bottleneck (the hidden state), which forces it to compress context — a useful inductive bias for tasks requiring summarization, but a limitation for tasks requiring exact recall of specific past tokens. Transformers, by caching all past key-value pairs, have no such compression bottleneck and can retrieve any past token with full precision. This fundamental difference in memory structure is the main driver of the performance trade-offs observed between the two architectures.

Q13 [Advanced] What are Mamba’s advantages and limitations on long-sequence tasks?

Q: Where does Mamba excel and where does it struggle compared to Transformers on long-sequence tasks?

A: Mamba’s primary advantage on long sequences is efficiency: both training memory (linear vs. quadratic) and inference memory (constant vs. linear) scale dramatically better than Transformers. This makes Mamba feasible for sequence lengths of 100K tokens or more that would be prohibitively expensive for standard Transformers. Domains where this matters include genomics (very long DNA sequences), audio processing (long waveforms), and long-document language modeling.

For tasks that benefit from compression of long contexts — such as summarization, extraction of global patterns, or time-series forecasting — Mamba’s recurrent structure is a good fit. The hidden state naturally accumulates a compressed representation of the entire history, and the selective mechanism allows relevant information to be preserved while irrelevant signals are filtered out.

However, Mamba has a well-documented weakness on tasks requiring precise retrieval of specific tokens from the distant past. Research on “needle in a haystack” benchmarks (finding a specific fact buried in a long document) and associative recall tasks (given a key, retrieve the associated value seen earlier) shows that Mamba’s fixed-size state struggles to retain all past information faithfully. Transformer attention, which stores all past tokens in the KV cache, can retrieve any past token exactly. Mamba must hope that the relevant information was encoded into its finite-dimensional state and not overwritten by subsequent inputs. This is the core limitation that makes Mamba less suitable for retrieval-augmented generation and multi-hop reasoning over long contexts.

Q14 [Advanced] When should you choose Mamba over Transformer, and vice versa?

Q: In practice, how do you decide between a Mamba-based and a Transformer-based architecture?

A: The decision hinges on three factors: sequence length, the nature of the task, and deployment constraints.

For very long sequences (beyond 8K–16K tokens), Mamba’s linear inference memory and sub-quadratic training complexity make it practically feasible in contexts where Transformers become prohibitively expensive. Applications in genomics, audio, long-context document understanding, and streaming inference are natural fits. If your use case requires processing 100K+ tokens and you do not have access to very large GPU clusters, Mamba (or a Mamba-based hybrid) is likely the only viable option.

For tasks requiring precise token-level retrieval — such as few-shot in-context learning, multi-hop question answering over long contexts, or retrieval-augmented generation — Transformer attention’s exact KV cache has a structural advantage. Mamba’s compressed state may lose the specific information needed for retrieval, leading to lower accuracy. Similarly, if you need to leverage large pretrained Transformer checkpoints (which vastly outnumber available Mamba checkpoints), fine-tuning a Transformer is the pragmatic choice.

For general-purpose language modeling at moderate context lengths (2K–32K), the two architectures are more competitive, and hybrid designs (interleaving Mamba and attention layers) often provide the best trade-off: Mamba layers handle efficient compression across most of the sequence, while attention layers provide precise retrieval capability at key points. When in doubt, a hybrid with a small fraction of attention layers (e.g., 1 in 8) captures most of Mamba’s efficiency gains while recovering most of Transformer’s retrieval ability.

Q15 [Advanced] Can Mamba replace Transformer?

Q: Is Mamba a replacement for Transformer, or are they complementary?

A: As of 2024, Mamba is best understood as complementary to Transformer rather than a replacement. At equivalent parameter counts and training compute, Mamba achieves competitive language modeling perplexity with Transformers and outperforms them on throughput-sensitive long-context benchmarks. However, Transformers retain advantages in instruction following, few-shot in-context learning, and precise retrieval tasks — capabilities that are central to the most practically valuable applications of large language models today.

The ecosystem gap is also significant: Transformers have years of pretraining scale, infrastructure tooling, and fine-tuning research behind them. Mamba’s largest pretrained models (as of early 2024) are at the 3B parameter scale, while the leading Transformer models are at hundreds of billions of parameters. Closing this gap requires substantial investment that the community is only beginning to make.

The most promising near-term direction is hybrid architectures. Models like Jamba (AI21 Labs) and Zamba interleave Mamba layers with a small number of Transformer attention layers, achieving most of Mamba’s efficiency while recovering Transformer’s retrieval and in-context learning capabilities. These hybrids demonstrate that the two architectures are not competing paradigms but complementary tools that can be combined. The longer-term question — whether pure SSM architectures can match Transformers at scale with appropriate training — remains open, but the hybrid approach already delivers practical benefits today.

Applications and Extensions

Q16 [Basic] How does Mamba perform in language modeling?

Q: How does Mamba compare to Transformer-based language models in practice?

A: Mamba shows strong performance in language modeling, particularly as sequence length increases. At the 1B–3B parameter scale with equivalent training tokens, Mamba achieves perplexity comparable to Transformer baselines like GPT-NeoX and achieves better throughput (higher tokens per second) at longer context lengths due to its linear inference complexity. The original Mamba paper demonstrated that Mamba-3B matches or exceeds GPT-NeoX-3B on standard language modeling benchmarks.

Mamba’s strengths in language modeling are most apparent on tasks that benefit from long context: long document summarization, code completion with large codebases as context, and structured prediction over long inputs. Its weaknesses appear on few-shot prompting tasks (in-context learning), where Transformer models consistently outperform Mamba at equivalent scale. This gap is attributed to Mamba’s fixed-size state being unable to maintain all few-shot examples with equal fidelity.

In terms of practical deployment, Mamba offers a meaningful advantage for applications that require serving long contexts with limited GPU memory — the constant KV-cache equivalent (fixed-size state) means that context length does not increase memory requirements at inference, enabling larger batch sizes and lower latency for long-context requests.

Q17 [Basic] How is Mamba applied to vision tasks (Vision Mamba)?

Q: How does Mamba handle 2D image data, and what are the approaches for applying it to vision?

A: Applying Mamba to images requires addressing a fundamental mismatch: Mamba’s SSM is designed for 1D sequences, but images are 2D spatial structures where locality and 2D positional relationships matter. The most direct approach, followed by Vision Mamba (Zhu et al., 2024), is to tokenize images into patches (as in ViT) and flatten them into a 1D sequence, then apply a bidirectional Mamba model.

Bidirectionality is important for vision because image patches have spatial neighbors in all directions, not just the causal (left-to-right) direction that standard Mamba uses. Vision Mamba introduces a bidirectional SSM by running two Mamba models in opposite scanning directions and combining their outputs, allowing each patch to aggregate information from both preceding and following patches in the scan order. Alternative scanning strategies have also been proposed: VMamba scans in four directions (horizontal, vertical, and two diagonals), and LocalMamba uses window-based local scanning to reduce complexity for high-resolution inputs.

In terms of performance, Vision Mamba achieves competitive accuracy with ViT models of comparable parameter count on ImageNet classification, with lower memory usage for high-resolution images where ViT’s quadratic attention becomes expensive. For dense prediction tasks (detection, segmentation), hierarchical variants that produce multi-scale feature maps (analogous to Swin Transformer) show the most practical promise.

Q18 [Advanced] What are hybrid Mamba + Transformer architectures and why are they promising?

Q: What is the motivation for hybrid architectures like Jamba, and how do they work?

A: Hybrid architectures interleave Mamba (SSM) layers with Transformer attention layers in the same model, aiming to capture the complementary strengths of both: Mamba’s efficiency and compression for the majority of layers, and Transformer attention’s precise retrieval capability for a small fraction of layers.

Jamba (AI21 Labs, 2024) is the most prominent example. It interleaves one Transformer attention layer for every seven Mamba layers, along with Mixture-of-Experts (MoE) feed-forward layers. The result is a model that supports a 256K-token context window with significantly better throughput than a pure Transformer at equivalent quality — the Mamba layers handle the bulk of sequence processing cheaply, while the attention layers provide the precise global retrieval that Mamba’s state cannot guarantee. Jamba achieves roughly 3x the throughput of a comparable Transformer at long context lengths.

The key design insight is that full global attention is not needed at every layer — most of the benefit of attention (exact retrieval, in-context learning) can be captured with a small fraction of attention layers, while Mamba layers handle the rest more efficiently. This is analogous to how Swin Transformer uses local attention in most layers and only occasionally computes cross-window interactions: not every layer needs the most expensive mechanism. The hybrid approach also makes Mamba-based models more competitive on in-context learning benchmarks, where pure Mamba models currently lag behind pure Transformers.

Q19 [Advanced] What is Mamba’s in-context learning capability and how does it compare to Transformer?

Q: Can Mamba perform in-context learning (ICL), and how does its ICL ability compare to Transformer?

A: In-context learning refers to a model’s ability to perform a new task given only a few input-output examples in the prompt, without any weight updates. This capability is widely regarded as a defining feature of large Transformer language models and is critical for practical deployment.

Mamba does exhibit in-context learning ability — it can use examples in the prompt to improve its outputs on novel tasks. However, research consistently shows that Mamba’s ICL ability degrades more rapidly than Transformer’s as the number of in-context examples increases. The proposed explanation is mechanistic: Transformer attention can retrieve any past example from the KV cache with full precision, making it straightforward to apply patterns from earlier examples to the current query. Mamba must compress all examples into its fixed-size state, and with many examples, the state becomes a lossy summary — earlier examples are more likely to be overwritten or diluted by later ones.

This limitation has practical implications: Mamba-based models may need larger parameter counts or specialized training to match Transformer ICL performance. Hybrid architectures with a small fraction of attention layers recover most of the ICL gap, as the attention layers maintain an exact record of the in-context examples while Mamba layers handle the efficient compression of the surrounding context. As a result, hybrid models like Jamba show much stronger ICL performance than pure Mamba models at equivalent parameter counts.

Q20 [Advanced] What are the current limitations of Mamba and future research directions?

Q: What are the main open problems and limitations of Mamba, and where is the field headed?

A: Mamba’s most fundamental limitation is its fixed-size state as a lossy compressor of history. This manifests as weaker performance on associative recall, needle-in-a-haystack retrieval, and multi-hop reasoning over long contexts — tasks where Transformer’s exact KV cache has a structural advantage. This is not a training or scale issue that will automatically resolve with more data; it is an architectural property of finite-state recurrence.

A second limitation is the relative immaturity of the ecosystem. As of early 2024, the largest available Mamba pretrained models are at the 3B parameter scale, while state-of-the-art Transformer models are at 70B–700B parameters. The scaling behavior of Mamba at very large model sizes and training budgets is not yet well-characterized. It is unclear whether the perplexity gap between Mamba and Transformer observed at 3B parameters will widen or close at 70B+ parameters.

On the research frontier, several directions are active. Hybrid architectures (combining Mamba and attention) are showing the most immediate practical results. Extending Mamba to 2D and higher-dimensional data (vision, video, graphs) requires non-trivial adaptations to the 1D scanning structure. The SSD framework introduced in Mamba2 opens theoretical connections to linear attention that may yield further architectural improvements. Hardware-efficient kernels for Mamba on non-NVIDIA hardware (TPUs, AMD GPUs) remain underdeveloped. Finally, the theoretical expressivity of SSMs — what functions they can and cannot represent as a function of state size and depth — is an active area that will inform better architecture design.

Quick Reference

Resources

• Gu & Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces (2023)
• Dao & Gu, Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (Mamba2, 2024)
• Gu et al., Efficiently Modeling Long Sequences with Structured State Spaces (S4, 2021)
• Gu et al., HiPPO: Recurrent Memory with Optimal Polynomial Projections (2020)
• Zhu et al., Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model (2024)
• Lieber et al., Jamba: A Hybrid Transformer-Mamba Language Model (2024)