cs.CLOct 14, 2026

Beyond State Capacity: Why Retrieval Breaks in Mamba

Authors: Nir Yarden, Amit Ben-Artzy, Roy Schwartz

Organizations: The Hebrew University of Jerusalem

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

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.
More in this topic

In-Context Retrieval

In-context retrieval studies how language models find and use relevant facts, examples, or documents supplied within their active context.

See all 29 papers

Explore similar work

CardsList

  1. Mamba with Hierarchical Memory: Solving Representation Bottleneck in Long Sequence Modeling

    Aug 3, 2026Qinwen Wang, Jieping Luo, Aoxiang Qin +5Long-Context RetrievalMemory-Augmented Neural Networks

    The Chinese University of Hong Kong, Shenzhen · Hainan Bielefeld University of Applied Sciences · BrainGalaxy

  2. Lost in State Space: Probing Frozen Mamba Representations

    Apr 30, 2026Bhagyashree Wagh, Akash SinghRepresentation GeometryRepresentation Collapse

    University of Washington · University of Florida

  3. Residual Streams Read, Recurrent States Remember: The Global Workspace in Mamba Models

    Sep 26, 2026Wenlong Wang, Fergal ReidState Space Language ModelsLLM Interpretability

    Fin AI Research