PaperlayerView exact revision on arXiv

Source-linked AI summary

On the Recall Scaling Laws in Mamba: A Theoretical and Mechanistic Study via Hashing

Yuval Koren, Assaf Ben-Kish, Raja Giryes, Lior Wolf, Itamar Zimerman

arXiv:2609.07681v1cs.LGcs.CL

TL;DR

Fixed-size recurrent states constrain how context is retained, making Mamba’s associative-recall mechanism and capacity important to characterize. The paper identifies a hash-like recall circuit and develops scaling laws whose predictions closely align with empirical results.

  • Problem

    The paper asks how Mamba performs associative recall and how its fixed-size recurrent state constrains recall across vocabulary and context sizes.

    Source: §1 · p. 1

  • Method

    The paper uses mechanistic interpretability to identify Mamba’s recall circuit, interprets it as hash-like, and derives dimension-based scaling laws.

    Source: §1 · pp. 1–2

  • Results

    Theoretical predictions closely align with empirical results, and high-probability mean-case recall has a dimension requirement ND = Ω(Nf log V) matching the stated upper scaling.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    For MQAR, the framework predicts the maximum vocabulary and number of facts over which accurate recall is achievable from model dimensions.

    Source: §6 · p. 10

    The analysis uses simplified models and does not fully characterize each full-Mamba component; the gating branch’s theoretical role in recall remains unclear.

    Source: §6 · p. 11

Abstract

from arXiv · show

Associative Recall (AR) is the cognitive ability to learn and retrieve links between items in memory. In NLP, AR is used as a benchmark for evaluating the in-context memory capacity of architectures such as Mamba, and has been found to strongly correlate with language modeling performance. This paper explores AR from the perspective of mechanistic interpretability, aiming to reverse-engineer the exact internal algorithm used by Mamba to perform recall. Our key insight is that Mamba performs recall by implicitly learning linear hash functions, and we identify the low-level circuit that enables this behavior. Building on these findings and inspired by theoretical tools in similarity-preserving hashing, such as the Johnson-Lindenstrauss lemma, we develop a theoretical framework for analyzing AR, which we term Recall Scaling Laws. Given the vocabulary size and the number of facts in context, this framework allows us to (1) predict the embedding and state dimensions required for Mamba to achieve perfect recall, (2) predict recall success probability given the model dimensions, and (3) analyze multi-layer models and multi-head SSM patterns. Empirical results show that our theoretical findings are accurate and predictive, offering insights into how AR capacity scales with vocabulary, state, embedding size, and architecture.

1 Introduction

The paper studies Mamba’s associative-recall capacity, motivated by the memory limits of fixed-size recurrent states, and identifies its recall circuit as hash-like. It develops dimension-based recall scaling guarantees and reports that theoretical predictions align closely with empirical results.

1 Introduction

  • The paper identifies a circuit-level mechanism for Mamba recall and interprets it as implicitly computing similarity-preserving linear hash functions.The analysis draws on hashing theory, including the Johnson–Lindenstrauss lemma.Source: §1 · pp. 1–2
  • Associative-recall benchmarks assess in-context memory in sequence models and are reported to correlate strongly with language-modeling performance.The introduction motivates studying Mamba’s recall against the fixed-size memory of recurrent architectures.Source: §1 · p. 1
  • The framework predicts embedding and state dimensions for perfect recall, recall probability from model dimensions, and behavior in multi-layer and multi-head models.Its scaling analysis relates dimensions to vocabulary size and the number of facts in context.Source: §1 · p. 2
  • N=D=V permits exact recall on every context in the non-compressive worst-case construction.The paper then gives smaller hash-based dimensions for exact worst-case recall and high-probability mean-case recall.Source: §1 · p. 2
  • ND = O(Nf log V) suffices for exact worst-case recall in a compressive hash-based model, while the mean-case guarantee is essentially optimal.The corresponding mean-case lower bound is ND = Ω(Nf log V); required state memory is linear in facts and logarithmic in vocabulary size.Source: §1 · p. 2

2 Background & Related Work

The background defines the multi-query recall task and outlines Mamba’s input-dependent state-space architecture. It situates the paper’s circuit analysis among mechanistic-interpretability and recall-theory work.

2 Background & Related Work

  • The recall-circuit description embeds tokens, uses duplication and copy/shift operations to form pairs, and uses compressed key–query matching and value reconstruction for retrieval.The circuit passage describes projection and compression for matching, followed by output selection and reconstruction.Source: Figure 2 · p. 3, §2 · p. 3
  • Prior work studies gather-and-aggregate retrieval [13], copying theory [26], or JL-based recall bounds [27], while this paper combines mechanistic analysis with tighter scaling laws.The paper’s related-work discussion distinguishes those studies by task focus or lack of analysis of learned mechanisms.Source: §2 · p. 3
  • MQAR presents key–value facts followed by queries that repeat context keys, requiring retrieval of each corresponding value; AR denotes the single-query case.The vocabulary is split into key and value subsets, and remaining query positions may contain padding.Source: §2 · p. 3
  • Mamba combines a selective state-space model with convolution and elementwise operations, using input-dependent recurrent rules at each timestep.Its block maps embedded inputs through projections and a convolution before the state-space operation and output projection.Source: §2 · pp. 3–4
  • Mechanistic interpretability links observable behavior to neural components and compact circuits, which this work applies to Mamba recall.The approach aims to uncover mechanisms responsible for recall in Mamba and linear RNNs.Source: §2 · p. 4

3 Reverse Engineering Mamba on MQAR

The paper reverse-engineers Mamba’s MQAR circuit and tests whether trained models use the proposed mechanism. It interprets recall as compressed hash-table storage and retrieval.

3.1 A Mamba Circuit for MQAR

  • A simplified single-layer circuit achieves perfect MQAR recall with D=N=V, but its dimensions make it impractical.The proposed alternative uses compression and achieves high-probability recall with ND=O(Nf log V).Source: §3.1 · pp. 4–6
  • In the compressive circuit, approximately orthogonal projections implement hashed key matching and value retrieval, with non-matching terms contributing residual error.The required dimension product is O(Nf log V) in the worst case and suffices for high-probability recall.Source: §3.1 · pp. 5–6

3.2 Validating the Circuit: Mechanistic Interpretability Experiments

  • Trained linear Mamba’s effective operators match the proposed circuit’s pattern: Gkq links queries to stored keys, while Gvv links stored values to outputs.The learned nonzero key-query blocks can be more complex than identity blocks and support key-value selectivity.Source: §3.2 · p. 6
  • The trained model’s hidden state behaves like compressed, random-looking storage whose projection back to vocabulary space preserves context information.The paper relates this to a hash-table interface: hash, store, match, and retrieve.Source: §3.2 · p. 7
  • A simplified Mamba block with a length-2 Conv1D achieves nearly perfect recall, but removing the Conv1D causes recall to fail completely.The proposed circuit assigns the Conv1D a copy-and-shift role.Source: §3.2 · pp. 7–8

4 Mamba Recall Scaling Laws

The paper derives dimension and probability scaling laws for associative recall, then extends the analysis to deeper and multi-head Mamba architectures.

4.1 JL-Based Recall Scaling Laws

  • JL-based recall laws give dimension conditions for perfect recall in associative recall and MQAR tasks.The analysis models successive embeddings and state-space compression as approximately orthogonal hashed representations.Source: §4.1 · p. 8

4.2 Probabilistic Recall Scaling Laws

  • Probabilistic laws account for cancellation among compression errors, while a lower bound gives ND=Ω(Nf log V) for high-probability Mamba recall.The lower bound is a special case of an information-theoretic constraint for finite-state recurrent models.Source: §4.2 · pp. 8–9

4.3 Multi-Layer and Multi-Head Scaling Laws

  • For simplified multi-layer Mamba, Theorem 4.6 provides a scaling law for recall probability.The theorem covers associative recall tasks and MQAR.Source: §4.3 · p. 9
  • With fixed model dimensions, MVA matches single-head recall probability, MQA and MKA are no better, and MHA performs at least as well as single-head.The ordering is pMQA=pMKA≤psingle=pMVA≤pMHA.Source: §4.3 · p. 9

5 Scaling Laws Experiments

Experiments show that the proposed scaling laws closely track recall across trained linear and full single-layer models, and extend to multi-layer and multi-head settings. The results also connect recall trade-offs to effective state size and head design.

5.1 Single-Layer Model: Accuracy Grids

  • Theory closely predicts trained linear-model accuracy grids, and the trained full model closely matches the linear model, supporting the simplification’s preservation of core recall behavior.Across the grids, accuracy also exhibits the predicted inverse trade-off between embedding dimension D and state size N.Source: §5.1 · p. 10

5.2 Single-Layer Model: Scaling Curves

  • Both AR and MQAR follow precall(x) ≈ Φ(x − √2 log V) after rescaling; MQAR’s larger variance makes high recall harder at the same D and N.The full nonlinear model follows a similar law, with afull ≈ 1/2 alinear estimated heuristically; results also align at vocabularies up to V = 50K and contexts up to L = 4K.Source: §5.2 · p. 10

5.3 Multi-Layer Model Experiments

  • In multi-layer models, recall accuracy depends on effective state size Neff = ΛN rather than on layer count Λ and state size N separately.Across depths, accuracy collapses onto a single curve closely following precall(x) ≈ Φ(x − √2 log V).Source: §5.3 · p. 10

5.4 Multi-Head Model Experiments

  • Among tested multi-head patterns, only MHA improves recall by enlarging both key and query dimensions, at the cost of more trainable parameters; MQA and MKA perform worse under shared-value compression.This pattern is consistent with Mamba-2 perplexity ablations: MQA and MKA perform worse than MVA, while MVA is stronger when parameter count is controlled.Source: §5.4 · p. 10

6 Conclusions

The paper identifies Mamba’s associative-recall algorithm and develops a framework that predicts recall capacity from model dimensions. Its predictions closely align with empirical results, while the analysis leaves some full-architecture mechanisms unresolved.

6 Conclusions

  • The paper identifies Mamba’s associative-recall algorithm and develops a theory predicting the maximum vocabulary size and number of facts for accurate MQAR recall from model dimensions.The theoretical predictions closely align with empirical results.Source: §6 · p. 10
  • The analysis uses simplified models and does not fully characterize each full-Mamba component’s contribution; in particular, the gating branch’s theoretical role remains unclear.The paper proposes extending the analysis to architectures including xLSTM, RWKV, and DeltaNet.Source: §6 · p. 11

A.1 Scaling Laws Experiments

The appendix presents scaling-law experiments, larger-vocabulary results, and mechanistic analyses of trained Mamba on MQAR. It also supplies notation and theoretical background for recall circuits and dimension-reduction arguments.

A.1 Scaling Laws Experiments

  • Across AR and MQAR, the scaling curves follow precall(x) ≈ Φ(x − b), with b ≈ √2 log V; MQAR has larger variance than AR, while the full model follows a similar law.The figure compares fixed-weight, theoretical, trained linear, and trained full-model accuracy grids across MQAR regimes.Source: Figure 8 · p. 16

A.2 Long Context and Large Vocabulary

  • For vocabularies from V = 1000 to V = 50000, fitted slopes of 1.18–1.56 differ from the theoretical slope of 1, consistent with the theory being an approximation.The experiments use context lengths L = 1000 to L = 4000.Source: §A.2 · p. 18, Figure 11 · p. 18

A.3 Reverse Engineering Mamba on MQAR

  • The mechanistic figures characterize hidden-state compression, approximate attention maps, Conv1D value-copy and key-shift patterns, and query-key matching followed by value retrieval.The invariant-operator analysis states that the core operation compares query qt with key kτ−1 and outputs value vτ.Source: Figure 13 · p. 19, Figure 14 · p. 20, Figure 15 · p. 21

B Notation and Dimensions

  • The appendix provides notation and tensor dimensions for MQAR, Mamba, and the SSM state-update rule, with further hidden-state and key-value-table dimensions referenced in Appendix D.2.It also points to Appendix H for Mamba-2 dimensions.Source: B Notation and Dimensions · p. 23, §B.1 · p. 23, §B.2 · p. 23

C.1 Johnson-Lindenstrauss (JL) Lemma

  • The JL argument uses random projections and concentration to preserve inner products; for one-hot vectors, direct Gram-entry bounds sharpen the constant to c = 4.The lemma’s stated dimension requirement is logarithmic in the number of vectors.Source: §C.1 · pp. 24–25

C.2 Non-Compressive Recall Circuit

  • A simplified single-layer Mamba with D = N = V, expand = 2, and Dconv = 2 can achieve perfect MQAR recall using a non-compressive circuit.The construction uses previous-token and current-token projections as key and query extractors, then retrieves the corresponding value.Source: §C.2 · pp. 25–26

C.3 Compressive Recall Circuit

The compressive circuit uses successive embeddings to hash and recover keys and values, yielding high-probability recall with dimensions that scale with facts and vocabulary. The appendices interpret this computation through operators, hidden-state tables, and approximate attention.

C.3 Compressive Recall Circuit

  • A single-layer simplified Mamba can solve MQAR with high probability when ND = O(Nf log V), using a compressive circuit based on the same key-value recall mechanism.The stated setting has context length L = 4Nf; exact quantitative laws are deferred to Lemma 4.1 and Theorem 4.2.Source: §C.3 · pp. 26–27

D.1 Invariant Operators

  • Invariant operators isolate query-key matching at (τ−1, t), then value retrieval at (t, τ), ignoring unrelated token positions.The matching operation compares query qt with key kτ−1; the retrieval operation selects the associated value.Source: §D.1 · p. 29

D.2 SSM Hidden State and Key-Value Tables

  • The ideal hidden state stores context facts as a key-value array, while the compressive circuit approximates that representation through its hidden state.The ideal array marks the value associated with each key, if that fact exists.Source: §D.3 · p. 30

D.3 Hidden State as a Hash Table

  • The hidden state is interpreted as a hash table: linear projections hash keys and queries, the state stores hashed key-value pairs, and inner products match queries to keys.The paper characterizes the internal representation as not directly interpretable and the final output as fairly interpretable.Source: §D.3 · pp. 30–31

D.4 Inspecting the Hash Table

  • In the general rotated circuit, key and query compression is distributed across SB, SC, Pin, and Wconv, so the standalone operator F cannot be read from one model weight.The analysis instead uses net operations Πk,in and Πq,in to combine these transformations.Source: §D.4 · p. 31

D.5 Noisy Attention Maps

  • Trained compressive models’ attention maps approximate the ideal circuit’s maps, which attend to prior tokens whose keys match the current query.This provides a circuit-level resemblance between ideal and trained compressive regimes.Source: §D.5 · p. 31

E Designed Model Baseline

  • The designed model is a sub-optimal baseline for associative recall, chosen because its weight construction is simpler to analyze mathematically.Its analysis is used alongside trained-model comparisons in Section 5 and Figure 1.Source: E Designed Model Baseline · p. 31

E.1 Designed Model Weights

  • The baseline constructs E from normalized Gaussian columns and F from a scaled orthonormal basis, then defines the compressed embedding as Ẽ = FE.For large V, one realization’s empirical Gram-matrix statistics closely match ensemble statistics.Source: §E.1 · pp. 31–32

F.1 JL-Based Scaling Laws

This section develops and refines JL-based recall scaling laws for simplified Mamba models, deriving conditions for perfect associative recall and probabilistic laws for designed and trained models. It also presents a unified approximation, an information-theoretic lower bound on recurrent state size, and a simplified multi-layer model with residual connections.

F.1 JL-Based Scaling Laws

  • For perfect recall, AR requires εv + εk + 2Nf εvεk < 1/2, while MQAR requires εv + εk + Lεvεk < 1/2, with εv, εk < 1.The JL framework also requires D, N > c log V.Source: §F.1 · p. 34

F.2 Probabilistic Scaling Laws: Designed Model

  • For a designed single-layer simplified Mamba, the analysis approximates AR and MQAR success probabilities in the limits Nf ≫ N, D and V ≫ 1.For MQAR, overall success is averaged across queries after accounting for query position.Source: §F.2 · pp. 34, 37

F.3 Probabilistic Scaling Laws: Trained Model

  • Orthogonal key and value embedding subspaces let a trained model ignore non-key-value pairs, so only key-value pairs contribute to recall output.The trained-model probabilistic scaling law is derived for Nf ≫ N, D and V ≫ 1.Source: §F.3 · p. 37

F.4 Unified Recall Scaling Laws

  • For a single-layer simplified Mamba solving associative recall tasks, relaxing the Nf ≫ N, D approximation and linearly approximating Φ for moderate ND/Nf gives 2.91 < a < 3, so a ≈ 3 is accurate.Source: §F.4 · pp. 37, 39
  • The refined denominator is close to the coarse approximation 5/4Nf + N, up to an additive term that yields higher-order corrections.The model-dimension remark states conditions for achieving recall success probability greater than 1 − δ, but the supplied passage does not include the conditions themselves.Source: §F.4 · p. 39

F.5 Information-Theoretic Lower Bound for Model Dimensions

  • A recurrent model solving AR with success probability greater than 1 − δ requires bS = Ω(Nf log V), a bound that also applies to full nonlinear Mamba.The bound assumes a fixed-size state with S entries at b bits of precision and also holds for MQAR when L = 4Nf.Source: §F.5 · p. 40

G.1 Multi-Layer Mamba

  • The simplified multi-layer Mamba model removes gating, discretization, nonlinearities, biases, and normalization, while retaining residual connections because they matter for multi-layer recall.The model is detailed in Algorithm 3 and is adapted from the single-layer model in Algorithm 1.Source: §G.1 · p. 40

G.2 Multi-Layer Recall Circuit and Scaling Laws

A simplified Mamba with Λ layers can match a single-layer model whose effective state is ΛN, preserving its recall probability. The section also describes how Mamba-2 shares components across SSM heads in several patterns.

G.2 Multi-Layer Recall Circuit and Scaling Laws

  • Each layer can add an independent associative-recall result, with residual connections passing the accumulated output to the next layer.The construction extends this process by induction across layers.Source: §G.2 · pp. 40, 43
  • A simplified Λ-layer Mamba solves recall with the same probability as a single-layer model with effective state size N_eff = ΛN.The multi-layer construction is equivalent to a single layer with the larger state size.Source: §G.2 · pp. 41, 43
  • In the two-layer construction, orthogonal key and value subspaces let the second layer act on embedded keys while zeroing embedded values.The layers can use different compressive-recall attention matrices because their projection matrices may differ.Source: §G.2 · p. 42

H.1 Multi-Head Patterns

  • Mamba-2 head patterns differ in which projected components they share: MVA shares B and C, MKA shares C, MQA shares B, and MHA shares none.The paper relates these choices to sharing V, K, and Q across Transformer attention heads.Source: §H.1 · p. 44, Table 5 · p. 44

H.2 Multi-Head Recall Circuits

The multi-head analysis compares recall capacity across Mamba-2 head patterns: MHA can benefit from splitting, while MQA, MKA, and MVA do not gain the same way. The notation section defines a simplified Mamba-2 setup with separate projections and convolutions for its SSM sequences.

H.2 Multi-Head Recall Circuits

  • MVA is equivalent for associative recall to a single-head model with the same dimensions D and N.The stated equivalence applies to AR or MQAR tasks.Source: §H.2 · p. 45
  • MQA and MKA are equivalent to single-head models with D′ = D/M and N′ = N, so splitting heads might hurt recall despite a larger query or key dimension MN.For MKA, the analysis requires shared queries; the quantitative analysis is otherwise the same as for MQA.Source: §H.2 · pp. 45–46
  • MHA can improve recall by effectively increasing state size to MN, although this success probability is an upper bound because optimization is not guaranteed to find those weights.The analysis attributes the effective increase to reduced variance in independent SSM compression noise.Source: §H.2 · pp. 46–47

H.3 Notation and Dimensions

  • The simplified Mamba-2 notation omits the state update matrix A, assumes it is identity, and separates the projections and convolution kernels that produce x̂, B, and C.The model uses no expansion for input projection; implementation may fuse the separately represented projections and kernels.Source: §H.3 · p. 47

I.1 Recall Circuit Minimality Ablation

Ablations identify convolution as the most critical component of the tested single-layer Mamba recall circuit, while repeated-key experiments show that depth is needed to retrieve the last value for a key.

I.1 Recall Circuit Minimality Ablation

  • Without convolution, MQAR accuracy falls to approximately 1/Nf because the model cannot bind keys to values and can only guess among the stored values.Here Nf = 16, so guessing among the stored values scores approximately 0.06; the ablation results are summarized in Table 8 and Figure 17.Source: §I.1 · p. 48, Table 8 · p. 48, Figure 17 · p. 49

I.2 MQAR with Repeated Keys

  • With repeated keys, a single layer performs at approximately 1/Nk accuracy, while two layers can overwrite stale values and retrieve the last occurrence.The proposed mechanism is that the first layer reads the stored value and the second writes the difference between the new and old values; Figure 18 summarizes the experiment.Source: §I.2 · p. 49, Figure 18 · p. 49

J Similarities and Differences with Understanding Input Selectivity in Mamba

The comparison argues that this work derives tighter, empirically supported recall bounds than Huang et al. [27], using a compressive circuit that compresses both keys and values. Mechanistic experiments support distributed storage of values across hidden-state coordinates rather than one key coordinate per value.

J Similarities and Differences with Understanding Input Selectivity in Mamba

  • The paper uses “Mamba” to mean Mamba-1, distinguishing it from the Mamba-2 architecture discussed separately.The comparison also notes that this work extends to multi-head recall circuits for Mamba-2.Source: J Similarities and Differences with… · p. 50
  • For a fixed number of facts, the proposed circuit achieves MQAR with ND = O(log V), improving on Huang et al.’s [27] dimensions, which leave keys uncompressed.The construction applies the Johnson–Lindenstrauss lemma to both values and keys; the compared bounds are D = O(Vk + log Vv) and N = Vk.Source: J Similarities and Differences with… · p. 50
  • Mechanistic experiments support distributed storage: writing one key-value pair affects all key coordinates, unlike the single-coordinate-per-key storage implied by Huang et al. [27].The tested model uses relatively small dimensions N = 32 and D = 64 for Vk = 512.Source: J Similarities and Differences with… · p. 50
  • The authors attribute the scaling dependence on both N and D to noisy retrieval, with value reconstruction noise and noisy query-key matching contributing distinct factors.The comparison describes the resulting dependence as performance improving with ND/Nf.Source: J Similarities and Differences with… · p. 51
  • Ablations and comparisons with the full gated model indicate that gating is not a bottleneck for MQAR recall; model dimensions D and N are identified as the main bottleneck.The gate-free simplified model and full gated model perform quite similarly on MQAR.Source: J Similarities and Differences with… · pp. 50–51

K.4 Compute and Environment

The paper reports its compute setup and explains why it uses best-of-seeds accuracy to assess whether recall solutions exist. Seed-level analyses show a consistent recall boundary with some isolated optimization failures.

K.4 Compute and Environment

  • Across the experiments, total compute was roughly 10 GPU-days; full Mamba grids took 2.3 ± 1.0 days per grid, compared with 8.0 ± 5.5 hours for simplified linear-model grids.The experiments used Nvidia B200 GPUs.Source: §K.4 · p. 51

K.5 Seeds Aggregation

  • Best-of-seeds accuracy is used because the analysis targets the existence of a recall solution rather than the probability that optimization finds it; seed-level grids show a consistent boundary with isolated failures.A seed counts as converged at accuracy ≥0.1; the figures compare per-seed accuracy grids, convergence fractions, and convergence rates.Source: §K.5 · p. 52, Figure 19 · p. 52, Figure 20 · p. 52

L.1 Impact Statement

  • The authors anticipate no direct negative societal impacts beyond those generally associated with LLMs.Source: §L.1 · p. 53

L.2 Reproducibility Statement

The statement describes the paper’s reproducibility materials, identifies its open-source-asset license table, and discloses how large language models were used.

L.2 Reproducibility Statement

  • The authors provide source code for the main experiments and report that datasets, models, and metrics are explicitly stated to support transparency and reproducibility.Configuration descriptions are available in the public repository; implementation details are in Appendix K.Source: §L.2 · p. 53

L.3 Licenses

  • Table 10 lists the open-source assets used in the work and their licenses.Source: Table 10 · p. 53

L.4 The Use of Large Language Models (LLMs)

  • ChatGPT and Claude assisted with writing, editing, and coding, but not with research ideation, design, implementation, analysis, or interpretation.Their writing role was limited to polishing presentation, including rephrasing and improving clarity and flow.Source: §L.4 · p. 53

Related papers

Loading 2609.07681v1…