Mamba models underperform Transformers on in-context retrieval. This gap is usually attributed to Mamba's fixed-size recurrent state, which loses information as the context grows. We challenge this assumption by focusing on query-first retrieval, a setup in which the query precedes the context and state capacity should not be the binding constraint. We show that even in this setting Mamba lags substantially behind Transformers, despite its selection mechanism, which was designed to distinguish relevant content from distractors. We then demonstrate that the Mamba selection gates have sufficient control over the state updates to influence the model output: intervening on specific gates alone can swing accuracy between extremes. However, as gates are input-dependent, they require the incoming token's representation to encode its relevance to the query. We therefore probe token representations across layers, and find that the relevance signal emerges only gradually with model depth, and that the gates begin to reflect it only several layers deeper. Crucially, we find that gates in retrieval-critical layers are not yet aligned with relevance. The selection mechanism is thus expressive enough to solve the task, but it commonly operates on a relevance signal that has not yet arrived. Mamba's retrieval bottleneck therefore extends beyond state capacity: even when capacity is sufficient, the delayed emergence of relevance information across depth often prevents successful retrieval.
Figures & tables
Figure 1: Mechanistic overview of query-first retrieval in Mamba Guided by the design principles behind Mamba’s architecture, we trace the flow of information across the layer and token-sequence axes, examining under a magnifying glass the mechanisms involved at query-first retrieval in Mamba, in search of the source of its degraded performance. Specifically, we zoom in on the following processes: ( a ) As token representations are mixed across layers, a relevance signal is encoded. ( b ) The selection gate responds to this signal by ‘opening and closing’, selectively writing relevant content to state while suppressing distractors. ( c ) The recurrent state stores the target information while being repeatedly updated, which requires protecting it from decay. ( d ) Finally, the target information is extracted from the state to generate the output token.
Figure 2: A local relevance flag in the target entry enables Mamba to identify and encode relevant information earlier, substantially improving long-context retrieval. ( a ) Evaluating the model on the query-first phonebook dataset, with and without a local relevance flag. The flag drastically improves Mamba’s accuracy on longer phonebooks. ( b ) Probing recurrent states for the target information reveals that the local relevance flag enables target decodability to emerge in earlier layers and with greater strength. ( c ) Calculating the difference between Mamba’s selection gates for target and distractor tokens demonstrates that the selection gate distinguishes token relevance earlier and more effectively when the local relevance flag is included. Showing results for Nvidia-Mamba; See fig. 6 for Codestral-Mamba results.
Figure 3: Manipulating the gates corresponding to the target information tokens governs performance. ( a ) Patching the unflagged prompt with flagged gating values recovers full flagged accuracy, whereas the reverse substitution reduces flagged accuracy to unflagged levels. ( b ) Oracle-scaling the target information tokens gates of the unflagged prompt can bring Mamba to near-perfect or near-zero accuracy.
Figure 4: Relevance information emerges gradually through the network, with delays both in contextualizing token representations and in aligning the gates with the resulting relevance signal. ( a ) Layer-wise linear probing shows that relevance becomes readily decodable from token representations much later when a local relevance flag is not present. ( b ) Layer-wise probes reveal a further delay: relevance becomes decodable from the gates only several layers after it emerges in the corresponding representations. ( c ) Cumulative layer-wise patching of either target information gates or representations shows that accuracy gains are concentrated in the early-to-middle layers, including layers where relevance is already decodable from representations but not yet adequately reflected in the gates.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: A sample of unflagged vs. flagged datasets. A single word addition, encoded to merely zero to two tokens, improves Mamba’s accuracy substantially.
Figure 6: Codestral-Mamba exhibits trends similar to Nvidia-Mamba when a local relevance flag is added to the target entry: ( a ) overall model accuracy improves, ( b ) the decodability of the target information in the recurrent state becomes stronger, and ( c ) selection gates better distinguish between target and distractor entries.
Recurrent linear attention models (RLAs) such as Mamba offer efficient linear-time sequence modeling as an alternative to Transformers, yet their fixed-capacity recurrent states limit long-sequence modeling. Drawing inspiration from hierarchical human memory, we propose Hierarchical Memory Mamba (HMM) to address this limitation. Building upon a pre-trained Mamba backbone, HMM integrates a lightweight working memory that extracts slow paragraph-level semantics (PLS) from the fast sensory memory embedded in the backbone's hidden states. The PLS is subsequently compressed into persistent long-term memory for task-relevant retrieval. The hierarchical processing of semantic information overcomes the representation bottleneck of RLAs and endows HMM cross-task generalization through parametric learning, which is not observed in other long-context enhanced Mamba variants. Evaluations on Passkey Retrieval and LongBench-E tasks demonstrate that HMM improves retrieval success by 34.3--37.1% and reasoning accuracy by 1.6--14.2% over strong Mamba-based models, while adding only 2% extra parameters and with minimal training overhead.
Qinwen Wang, Jieping Luo, Aoxiang Qin +5
The Chinese University of Hong Kong, Shenzhen · Hainan Bielefeld University of Applied Sciences · BrainGalaxy
Mamba's recurrent state h_t is, by construction, a compressed summary of every token seen so far. This raises a tempting hypothesis: if we extract token-level outputs y_t at fixed patch boundaries, we obtain semantic sentence summaries for free, with no pooling head, no fine-tuning, and no [CLS] token. We test this hypothesis carefully. Across five benchmarks (SST-2, CoLA, MRPC, STS-B, IMDb), we compare four strategies for extracting frozen sentence representations from a pretrained Mamba-130M backbone under a strict frozen-feature probing protocol, using three random seeds where computationally feasible. The results do not support the hypothesis: patch boundary readouts do not consistently outperform simple mean pooling. We identify and quantify two structural pathologies: severe anisotropy (mean pairwise cosine similarity 0.9999, std 0.000044) and representational collapse in the raw final SSM state (MCC = 0.000 on CoLA across all three seeds, confirmed via confusion matrix). We further propose orthogonal injection, a modified recurrence that constrains new information per
Can the global-workspace account of transformer representations extend to state-space language models? We fit Jacobian lenses to the residual streams and recurrent states of Mamba-1, Mamba-2 and Mamba-3, using the original 1000-prompt recipe. Joint residual--state readouts improve recovery of known intermediate concepts over the residual lens on at least five of six task families in every tested Mamba checkpoint. On Mamba-2, state alone exceeds residual and logit lenses on all six families; a normalised joint readout improves on both components on five. Temporal maps and word-list experiments show earlier content remaining state-readable as residual visibility changes. We also propose sign-guarded steering, which improves target top-five success over coordinate exchange on matched verbal-report trials in five models. Recurrent state alone supports this verbal access. These gains do not extend consistently to relational answers: guarded edits often output the edited concept itself, and Mamba-3's joint edits can disrupt successful state-only redirection. Recurrent state thus provides a complementary carrier of workspace content, whose recovery, persistence and causal uses require separate measurements.