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

> Mamba models learn linear hash functions for associative recall, requiring state memory scaling as ND = Θ(Nf log V) — linear in facts, logarithmic in vocabulary.

- **Source:** [arXiv](https://arxiv.org/abs/2609.07681)
- **Published:** 2026-10-10
- **Permalink:** https://picx.dev/p/hCyKV6
- **Whiteboard:** https://picx.dev/p/hCyKV6/image

## Summary

## Summary (Overview)

- This paper reverse-engineers the internal algorithm used by Mamba models to perform Associative Recall (AR) and Multi-Query Associative Recall (MQAR), identifying that Mamba implicitly learns **linear hash functions** via a similarity-preserving mechanism.
- The authors develop a theoretical framework called **"Recall Scaling Laws"** that predicts the embedding dimension $D$ and state dimension $N$ required for perfect recall given vocabulary size $V$ and number of facts $N_f$.
- The key theoretical result shows that high-probability recall requires model dimensions satisfying $ND = \Omega(N_f \log V)$ — the state memory scales **linearly in the number of facts** and **logarithmically in vocabulary size**.
- The framework extends to **multi-layer** models (where depth $\Lambda$ multiplies effective state size) and **multi-head** SSM patterns (where MHA improves recall but MQA/MKA degrade it).
- Extensive empirical validation confirms the theoretical predictions across linear and full nonlinear Mamba models, with accuracy collapsing onto predicted one-dimensional scaling curves.

---

## Introduction and Theoretical Foundation

### Background

Transformers suffer from linear memory growth during decoding due to key-value caching. Recent architectures like Mamba [1], RWKV [2], and other linear RNNs use a **fixed-size recurrent state**, enabling constant memory complexity. However, this fixed-size state imposes inherent information compression limitations.

The paper addresses a core question: **How does the fixed-size state of Mamba limit its recall capabilities, and what is the exact scaling relationship?**

### Theoretical Foundations

The analysis leverages two key theoretical tools:

1. **Johnson–Lindenstrauss (JL) Lemma**: States that points in a high-dimensional space can be embedded into a lower-dimensional space while approximately preserving pairwise distances. Formally, for a set of $V$ points, an embedding into $\mathbb{R}^k$ with $k = O(\log V / \varepsilon^2)$ preserves distances up to factor $(1 \pm \varepsilon)$ with high probability.

2. **Mechanistic Interpretability**: The approach of reverse-engineering neural networks by identifying specific "circuits" — compact computational units within the network that implement identifiable algorithms.

### The MQAR Task

The Multi-Query Associative Recall task partitions vocabulary $\mathcal{V}$ into key vocabulary $\mathcal{V}_k$ and value vocabulary $\mathcal{V}_v$, each of size $V/2$. A prompt contains:
- **Context**: $N_f$ non-repeating key-value pairs $(k_i, v_i)$
- **Query section**: $N_f$ query tokens (duplicates of context keys) plus padding

The model must retrieve the value $v^*$ corresponding to each query key $k^*$.

### Mamba Architecture

The Mamba block operates as (simplified):

$$
\hat{x} = \text{SiLU}(\text{Conv1D}(\text{Linear}(x^e))), \quad \hat{z} = \text{SiLU}(\text{Linear}(x^e)), \quad \hat{y} = \text{SSM}(\hat{x}) \odot \hat{z}
$$

with the SSM update:

$$
h_t^d = \bar{A}_t^d \odot h_{t-1}^d + \hat{x}_t^d \bar{B}_t^d, \quad \hat{y}_t^d = h_t^d \cdot C_t
$$

---

## Methodology

### Minimal Model and Simplified SSM

The authors construct a simplified single-layer Mamba with gating, discretization, nonlinearities, biases, and normalization removed. Without discretization and with $A = I$, the SSM becomes:

$$
h_t = h_{t-1} + \hat{x}_t B_t^\top, \quad \hat{y}_t = h_t C_t
$$

### The Ladder of Theoretical Guarantees

The paper establishes four levels of recall guarantees:

| Level | Guarantee | Dimension Requirement | Type |
|-------|-----------|----------------------|------|
| 1 | Exact recall (worst-case) | $N = D = V$ | Non-compressive (Thm. 3.1) |
| 2 | Exact recall (worst-case) | $\sqrt{ND} = O(N_f \log V)$ | Hash-based (Lem. 4.1) |
| 3 | High-probability recall | $ND = O(N_f \log V)$ | Mean-case (Thm. 4.2) |
| 4 | **Lower bound** (necessity) | $ND = \Omega(N_f \log V)$ | Information-theoretic (Lem. 4.5) |

### Mechanistic Interpretability Validation

The authors define **invariant operators** that are robust to orthogonal transformations of weights:

$$
\hat{E}_{\text{in}} = \left(\text{diag}(W_{\text{conv}}^0) P_{\text{in}} E \mid \text{diag}(W_{\text{conv}}^1) P_{\text{in}} E\right), \quad \Pi_{v,\text{out}} = E^\top P_{\text{out}}
$$

$$
\Pi_{k,\text{in}} = S_B \hat{E}_{\text{in}}, \quad \Pi_{q,\text{in}} = S_C \hat{E}_{\text{in}}, \quad G_{vv} = \Pi_{v,\text{out}} \Pi_{v,\text{in}}, \quad G_{kq} = \Pi_{k,\text{in}}^\top \Pi_{q,\text{in}}
$$

The model output becomes:

$$
y_t = \sum_{\tau=0}^{t} G_{vv} \xi_\tau \xi_\tau^\top G_{kq} \xi_t
$$

where $\xi_t = \binom{x_{t-1}}{x_t}$ are input token pairs.

### Hidden State as Hash Table

The ideal non-compressive circuit stores facts as an outer product:

$$
H_t \equiv P_{\text{out}} h_t = \sum_{n=1}^{N_f} v_n k_n^\top, \quad y_t = H_t q_t
$$

This is a $V \times V$ table where entries are 1 where facts exist. The compressive version uses:

$$
H_t' = \sum_{n=1}^{N_f} v_n' k_n'^\top, \quad y_t = E^\top \sum_{n=1}^{N_f} v_n' k_n'^\top q_t'
$$

with compressed tokens $v_n' = Ev_n$, $k_n' = \tilde{E}k_n$, $q_t' = \tilde{E}q_t$.

---

## Empirical Validation / Results

### Key Theoretical Results

**Theorem 3.1** (Perfect non-compressive recall): A single-layer simplified Mamba with $D = N = V$, $\text{expand} = 2$, $D_{\text{conv}} = 2$ perfectly solves MQAR (recall probability = 1).

**Theorem 3.2** (Efficient compressive recall): With $ND = O(N_f \log V)$, $\text{expand} = 2$, $D_{\text{conv}} = 2$, a single-layer Mamba solves the task with high probability.

**Theorem 4.2** (Trained model recall scaling laws):

$$
p_{\text{AR}} \approx \Phi\left(\sqrt{\frac{ND}{N_f}} - \sqrt{2\log V}\right)
$$

$$
p_{\text{MQAR}} \approx \frac{1}{L - 2N_f} \sum_{t=2N_f}^{L} \Phi\left(\sqrt{\frac{ND}{\frac{1}{2}N_f + \frac{1}{4}t}} - \sqrt{2\log V}\right)
$$

**Remark 4.3** (Unified form):

$$
p_{\text{recall}} \approx \Phi\left(\sqrt{\frac{ND}{aN_f + N}} - \sqrt{2\log V}\right)
$$

where $a = 1$ for AR and $a \approx 5/4$ for MQAR.

**Theorem 4.6** (Multi-layer):

$$
p_{\text{recall}} \approx \Phi\left(\sqrt{\frac{\Lambda ND}{aN_f + \Lambda N}} - \sqrt{2\log V}\right)
$$

**Theorem 4.7** (Multi-head): Given fixed $D, N$:

$$
p_{\text{MQA}} = p_{\text{MKA}} \leq p_{\text{single}} = p_{\text{MVA}} \leq p_{\text{MHA}}
$$

### Empirical Findings

**Figure 1 results** — Accuracy grids show:
- Theoretical predictions (column b) closely align with trained linear models (column c)
- Full nonlinear models (column d) match linear models, confirming simplification preserves core recall behavior
- All columns exhibit the predicted inverse $D$–$N$ tradeoff

**Figure 4 results** — Scaling curves:
- Accuracy collapses onto one-dimensional curves of the form $p_{\text{recall}}(x) \approx \Phi(x - b)$
- The scaling variable is $x = 1/\sigma$ where $\sigma^2 = \frac{aN_f}{ND} + \frac{1}{D}$
- $b \approx \sqrt{2\log V}$
- Full nonlinear Mamba fits with $a_{\text{full}} \approx \frac{1}{2}a_{\text{linear}}$, indicating improved performance under the same scaling law

**Multi-layer results** (Figure 5):
- Recall accuracy depends only on **effective state size** $N_{\text{eff}} = \Lambda N$
- Accuracy collapses onto a single curve independent of depth $\Lambda$

**Multi-head results** (Figure 6):
- MHA improves recall by increasing effective state size
- MQA and MKA are weaker due to shared-value compression across heads
- MVA matches the single-head baseline

### Mechanistic Validation

- **Hidden state inversion** (Figure 3): Projecting the hidden state back to vocabulary space reveals $H_t'' \approx H_t$, confirming a compression-decompression scheme
- **Conv1D as copy-shift**: $G_{vv}\xi_\tau \approx x_\tau$ and $G_{kq}^\top\xi_\tau \approx \xi_{\tau-1}$, verified empirically
- **Ablation**: Removing Conv1D causes complete failure of recall

---

## Theoretical and Practical Implications

### Theoretical Implications

1. **Optimality of scaling**: The upper bound $ND = O(N_f \log V)$ is matched by an information-theoretic lower bound $ND = \Omega(N_f \log V)$, establishing that Mamba's recall capacity is essentially optimal for fixed-state recurrent models.

2. **Hash-table interpretation**: The paper provides strong evidence that Mamba learns similarity-preserving linear hash functions, connecting neural network behavior to classical data structures.

3. **Multi-head design guidance**: The result that MQA/MKA degrade recall while MHA improves it provides theoretical justification for architectural choices in Mamba-2.

### Practical Implications

- **Model sizing**: Given a target vocabulary size and number of facts, practitioners can predict required dimensions for reliable recall
- **Architecture selection**: MVA is validated as the stronger multi-head design when controlling for parameter count, consistent with real-world perplexity results
- **Diagnostic tool**: The scaling laws provide a principled way to evaluate whether a model is memory-limited or has other bottlenecks

### Connection to Language Modeling

The trends in Thm. 4.7 are consistent with perplexity ablations for full Mamba-2 (Dao & Gu, 2024): MQA and MKA both yield worse perplexity than MVA, matching the finding that shared-value compression degrades recall. This suggests the analysis captures tradeoffs observed in real-world NLP tasks.

---

## Conclusion

### Main Takeaways

1. Mamba performs associative recall by **implicitly learning linear hash functions**, verified through mechanistic interpretability
2. The required state memory scales as $ND = \Theta(N_f \log V)$ — **linear in facts, logarithmic in vocabulary**
3. The theoretical framework accurately predicts recall probability across model configurations
4. Multi-layer models benefit through effective state size $\Lambda N$; multi-head patterns have distinct recall characteristics

### Future Directions

- Extend analysis to other architectures: xLSTM, RWKV, DeltaNet
- Investigate how architectural modifications impact recall capabilities
- Characterize the theoretical role of the gating branch in recall (currently unclear)
- Detailed investigation of why full models achieve $a_{\text{full}} \approx \frac{1}{2}a_{\text{linear}}$ (improved performance under same scaling law)

### Limitations

- Analysis relies on simplified models; the full contribution of each Mamba component (e.g., gating) is not fully characterized
- The heuristic determination of $a_{\text{full}}$ for nonlinear models requires further theoretical justification

---

_Markdown view of https://picx.dev/p/hCyKV6, served by PicX — AI-generated visual whiteboard summaries of research papers._