Paper 2510.03279
MemMamba: Rethinking Memory Patterns in State Space Model
- Published
- Sep 2025
- Research lab
- Independent
- Citations
- 2
- GitHub
- Not linked
01 In brief
Summary
MemMamba is a novel architecture for long-sequence modeling that addresses the memory decay problem in state space models (SSMs) like Mamba.
The paper first analyzes Mamba's memory mechanism, showing that its state update leads to exponential decay of early information both within layers (horizontal) and across layers (vertical), quantified by new metrics: Expected Token Memory Fidelity (ETMF) and Expected Cross-Layer Memory Fidelity (ECLMF).
To mitigate this, MemMamba introduces a Note Block that compresses and stores important tokens into a state pool, and uses cross-token attention (every layer) and cross-layer attention (every p layers) to restore forgotten information.
This design maintains linear complexity O(n).
Experiments on PG19 language modeling show MemMamba achieves stable perplexity (17.35) at 60k tokens, where Mamba and DeciMamba fail.
On Passkey Retrieval, it maintains 90% accuracy at 400k tokens, outperforming baselines.
It also excels in cross-document retrieval under noise.
MemMamba delivers a 48% inference speedup over Transformers.
Theoretical analysis proves linear complexity, BIBO stability, and non-vanishing gradients.
The work provides a new paradigm for ultra-long sequence modeling by breaking the complexity-memory trade-off.
02 From the paper
Abstract
With the explosive growth of data, long-sequence modeling has become increasingly important in tasks such as natural language processing and bioinformatics. However, existing methods face inherent trade-offs between efficiency and memory. Recurrent neural networks suffer from gradient vanishing and explosion, making them hard to scale. Transformers can model global dependencies but are constrained by quadratic complexity. Recently, selective state-space models such as Mamba have demonstrated high efficiency with O(n) time and O(1) recurrent inference, yet their long-range memory decays exponentially. In this work, we conduct mathematical derivations and information-theoretic analysis to systematically uncover the memory decay mechanism of Mamba, answering a fundamental question: what is the nature of Mamba's long-range memory and how does it retain information? To quantify key information loss, we further introduce horizontal-vertical memory fidelity metrics that capture degradation both within and across layers. Inspired by how humans distill and retain salient information when reading long documents, we propose MemMamba, a novel architectural framework that integrates state summarization mechanism together with cross-layer and cross-token attention, which alleviates long-range forgetting while preserving linear complexity. MemMamba achieves significant improvements over existing Mamba variants and Transformers on long-sequence benchmarks such as PG19 and Passkey Retrieval, while delivering a 48% speedup in inference efficiency. Both theoretical analysis and empirical results demonstrate that MemMamba achieves a breakthrough in the complexity-memory trade-off, offering a new paradigm for ultra-long sequence modeling.