Source-linked AI summary
MemMamba: Rethinking Memory Patterns in State Space Model
Youjin Wang, Yangjingyi Chen, Jiahao Yan, Jiaxuan Lu, Xiao Sun
TL;DR
Long-sequence models face a trade-off between memory and efficiency, and Mamba’s long-range memory decays despite its efficient computation. The paper analyzes this decay and introduces memory-fidelity measures, then proposes MemMamba, which improves long-sequence benchmark performance while retaining linear scaling and achieving a 48% speedup.
Problem
The paper asks how Mamba’s memory evolves with distance and depth, because long-range memory decay limits its ability to retain information.
Source: §1 · p. 2, §3.1 · p. 3Method
The paper analyzes Mamba’s memory decay, defines token-level and cross-layer fidelity measures, and proposes MemMamba with state summarization and cross-layer and cross-token attention.
Source: §1 · p. 2, §3.4 · p. 4, §4.2 · p. 6Results
MemMamba achieves state-of-the-art results across long-sequence benchmarks while retaining linear time and space scaling, with a 48% speedup over baseline architectures.
Source: §1 · p. 2, §6 · p. 10Takeaways & Limitations
The results support integrating structured memory into state-space models without sacrificing efficiency on the reported long-sequence benchmarks.
Source: §6 · pp. 9–10The Mamba state-transition analysis assumes |A| < 1 to guarantee BIBO stability.
Source: §3.1 · p. 4
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
The paper examines how long-sequence models balance memory and efficiency, analyzing Mamba’s memory decay and proposing MemMamba to retain salient information. It reports broad benchmark gains alongside faster inference.
1 INTRODUCTION
- Long-sequence modeling must capture ultra-long-range dependencies efficiently, while RNNs face vanishing or exploding gradients and Transformers incur quadratic sequence-length complexity.The paper describes this as an architectural impasse between expressiveness and scalability.Source: §1 · pp. 1–2
- Mamba’s early information decays exponentially during intra-layer recursion and inter-layer propagation, constraining long-range memory despite computational stability.The paper introduces horizontal–vertical memory fidelity to quantify token-level semantic transmission and cross-layer information coupling.Source: §1 · p. 2
- At 60k tokens, PG19 perplexity is 17.35 and Passkey Retrieval accuracy is 90% at 400k tokens.The PG19 result is contrasted with Mamba and DeciMamba of similar parameter scale, which collapse completely; cross-document retrieval also favors MemMamba under noisy conditions.Source: §1 · p. 2
- MemMamba combines lightweight state summarization with cross-layer and cross-token attention to preserve and reuse salient information while maintaining linear complexity.The design is inspired by how people take notes while reading long texts.Source: §1 · p. 2
- MemMamba consistently outperforms Mamba variants and Transformer baselines across language modeling, sparse retrieval, and cross-document reasoning, with 48% inference speedup.The contribution summary reports the cross-task pattern and speedup.Source: §1 · p. 2
2 RELATED WORK
Prior work develops state-space models and other long-sequence architectures for efficiency, memory, and broader applications. The paper identifies memory behavior in Mamba as an underexplored question.
2.1 STATE SPACE MODELS
- Mamba uses selective state-space modeling for efficient sequence modeling, while related efforts extend its architecture, improve computational efficiency, or broaden its applications.Examples include bidirectional state updates in BiMamba and parallelization and caching in FastMamba.Source: §2.1 · pp. 2–3
2.2 LONG-SEQUENCE MODELING
- Long-sequence approaches include external-memory models, compressed-memory Transformers, and sparse-attention variants, with trade-offs in information loss, task scope, or ultra-long-sequence handling.The cited discussion notes that Compressive Transformer risks information loss, while Longformer and BigBird struggle with ultra-long sequences.Source: §2.2 · p. 3
- Mamba extensions target longer contexts, multimodal distillation, reduced memory overhead, or stability, but the passage notes specific risks and remaining copying or in-context-learning gaps.Examples include DeciMamba’s risk of fine-grained information loss and Mamba-2’s lag behind Transformers on copying or in-context learning.Source: §2.2 · p. 3
- Despite progress in structural design and context extension, how models remember and forget critical information remains largely unexplored.This paper focuses on Mamba’s memory patterns and long-range forgetting.Source: §2.2 · p. 3
3 INVESTIGATION OF MEMORY PATTERNS
The investigation traces memory loss in Mamba and Transformers, then frames retention as both a token-level and cross-layer problem. These analyses motivate complementary measures of memory fidelity.
3.1 MEMORY DECAY IN MAMBA
- For an input k steps in the past, Mamba’s contribution decays exponentially as k grows, so early inputs are nearly forgotten.The analysis assumes a stable state transition matrix with |A| < 1; the contribution is expressed through A^k.Source: §3.1 · p. 4
3.3 MEMORY DECAY IN TRANSFORMER
- Transformer self-attention models global dependencies but has quadratic sequence-length complexity; practical windowed or sparse approximations discard information outside the window.For ultra-long sequences, the passage states EML ≤ w ≪ n, limiting effective long-range dependency capture.Source: §3.3 · p. 4
3.4 HORIZONTAL AND VERTICAL MEMORY FIDELITY
- The Horizontal–Vertical Memory Fidelity Framework separates within-layer token loss from information loss across layers.ETMF measures preservation of token semantics across time steps; ECLMF measures information preservation across layers, and together they offer complementary views of degradation.Source: §3.4 · p. 4
- MemMamba’s overall workflow stacks block layers that preserve context with a Note Block and enable long-range interaction through sparse cross-layer attention.Each block combines SSM updates, cross-token attention, and periodically triggered cross-layer attention.Source: Figure 1 · p. 5, Figure 2 · p. 5
4 METHOD
MemMamba combines state summarization with cross-token and periodic cross-layer attention to preserve salient information across long sequences. Its Note Block stores compressed summaries, while the MemMamba Block integrates memory interactions with SSM updates.
4 METHOD
- MemMamba layers combine SSM updates with cross-token attention at every layer and cross-layer attention activated every p layers.This sparse schedule balances memory retention and computational efficiency.Source: §4 · p. 6
4.1 NOTE BLOCK
- The Note Block scores tokens, compresses those above a threshold into a state pool, and retains only high-information summaries using FIFO or priority-based replacement.The pool has limited capacity; dimensionality reduction can use linear projection or pooling.Source: §4.1 · p. 6
4.2 MEMMAMBA BLOCK
- The MemMamba Block combines SSM updates, threshold-triggered cross-token attention, and periodically triggered cross-layer attention.The state pool is checked against the current SSM state; when forgetting is detected, retrieved summaries are fused with the current input.Source: §4.2 · p. 6, §4 · p. 6
5 EXPERIMENTAL ANALYSIS
Across language modeling, sparse retrieval, and cross-document reasoning, MemMamba improves on the reported baselines while preserving linear complexity. The experiments also report faster inference and compare its memory mechanisms and fidelity metrics.
5.1 COMPARISON WITH BASELINES
- MemMamba outperforms all baselines at most context lengths and shows much stronger robustness and stability on ultra-long 30k–60k token sequences, although performance degrades for all models.For passkey retrieval, MemMamba accurately retrieves key information when the target is more than 200k tokens from the prediction point, while Mamba and Pythia fail. It also leads on document retrieval under high-noise conditions.Source: §5.1 · p. 8
- MemMamba consistently improves on state-of-the-art baselines across language modeling, sparse retrieval, and cross-document reasoning while retaining linear complexity.It shows stronger scalability than Transformer-based models on ultra-long sequences and stronger long-term memory retention than Mamba variants.Source: §5.1 · p. 8
5.2 ABLATION STUDIES
- At 1.5k tokens, MemMamba’s PPL is 19.35; at 60k tokens, it fluctuates within 17.33–18.25, while removing both memory mechanisms produces a stark contrast.On the same hardware, its end-to-end latency is 0.52× that of the Transformer, a 48% speedup; complexity is O(n+m).Source: §5.2 · pp. 8–9
- MemMamba significantly outperforms Mamba and DeciMamba on both ETMF and ECLMF, whose higher scores indicate better retention of early token information.DeciMamba has a slight advantage in extremely long-range cross-layer transmission, but is less stable.Source: §5.2 · p. 9
5.3 PROOF OF LINEAR COMPLEXITY
- MemMamba’s computational cost scales as O(n · d), compared with O(n^2d) for Transformers, by keeping the state dimension and attention pool size constant.The proof passage identifies the state dimension ds and pool size k as the constrained quantities.Source: §5.3 · p. 9
6 CONCLUSION
MemMamba augments Mamba with dynamic state summarization and lightweight cross-layer and cross-token attention. The paper reports strong long-sequence benchmark results and linear time and space scaling, with future work aimed at broader settings.
6 CONCLUSION
- MemMamba combines dynamic state summarization with lightweight cross-layer and cross-token attention to address memory decay while preserving efficiency.The conclusion describes future extensions to multimodal settings, retrieval-augmented systems, and scaling as a foundation model.Source: §6 · pp. 9–10
- MemMamba achieves state-of-the-art results on PG19, Passkey Retrieval, and multi-document reasoning while retaining linear time and space scaling.The conclusion reports a 48% speedup over baseline architectures.Source: §6 · p. 10
- Future work will explore multimodal settings, retrieval-augmented systems, and scaling MemMamba as a foundation for efficient, high-fidelity memory.Source: §6 · p. 10
A.1 ON THE FORGETTING OF CRITICAL INFORMATION
The passage characterizes critical-information loss as an irreversible effect of compression or approximation, with accumulated loss leading to key information disappearing.
A.1 ON THE FORGETTING OF CRITICAL INFORMATION
- Compression or approximation introduces irreversible loss of critical information.The passage frames this in terms of an information-theoretic compression limit.Source: §A.1 · p. 12
- Accumulated information loss leads to the disappearance of key information.Source: §A.1 · p. 12
- Transformers and Mamba exhibit fundamentally different information-loss mechanisms.Source: §A.1 · p. 12
A.1.1 CRITICAL-INFORMATION LOSS IN TRANSFORMERS
The passage describes irreversible information loss when Transformer compression excludes critical features from the approximation's projection subspace, resulting in unbounded reconstruction error.
A.1.1 CRITICAL-INFORMATION LOSS IN TRANSFORMERS
- Transformer compression may use low-rank approximations such as Nyström or sparse attention.Source: §A.1.1 · p. 12
- If critical features fall outside the approximation's projection subspace, their loss is irreversible.Source: §A.1.1 · p. 12
- When critical features are excluded from that subspace, the passage says reconstruction error is unbounded.Here, x is the original critical signal, x̂ its reconstruction, and the projection subspace is the low-rank subspace used by the approximation.Source: §A.1.1 · p. 12
A.1.2 CRITICAL-INFORMATION LOSS IN MAMBA
Mamba's state compression makes long-range critical information vulnerable to forgetting within and across layers.
A.1.2 CRITICAL-INFORMATION LOSS IN MAMBA
- Mamba's state compression makes it easy to forget long-range critical information both within and across layers.Source: §A.1.2 · p. 12
1. Intra-layer dependence. Each layer’s state h(l)
The analysis characterizes temporal and cross-layer information loss in Mamba, then derives MemMamba’s memory, efficiency, and robustness properties. Its proposed summaries and attention preserve long-range information while retaining linear complexity.
A.1.3 THEORY OF CROSS-LAYER TRANSMISSION OF CRITICAL INFORMATION; A.2.1 EXPECTED TOKEN MEMORY FIDELITY (ETMF); A.2.2 EXPECTED CROSS-LAYER MEMORY FIDELITY (ECLMF)
- When temporal transition magnitudes are below one, early information decays exponentially across time and depth; ETMF and ECLMF measure token-level and cross-layer fidelity, respectively.The framework reports that state summarization and cross-layer attention improve both fidelity measures.Source: §A.1.3 · p. 13, §A.2.1 · p. 13, §A.2.2 · p. 14
A.3.1 ERROR BOUND OF POOLING APPROXIMATION
- Max pooling preserves maxima by design, while its reconstruction error is bounded and controllable when local fluctuation Δ is small.The bound scales with √n and Δ is described as typically small for text.Source: §A.3.1 · p. 14
A.3.2 EQUAL-BUDGET COMPARISON: MEMMAMBA VS. TRANSFORMER/MAMBA
- Under equal memory, MemMamba’s long-range recall is at least 1 − δ (δ < 0.05), and it achieves higher accuracy than Mamba.At equal compute, the passage gives an example of roughly 100× longer contexts for MemMamba.Source: §A.3.2 · p. 15
A.3.3 BIBO STABILITY OF MEMMAMBA; A.3.4 CONVERGENCE OF GRADIENT PROPAGATION
- The analysis establishes BIBO stability for bounded inputs and a nonzero gradient path that avoids vanilla Mamba’s |A|^L vanishing-gradient issue.The gradient argument assumes a bounded Softmax derivative.Source: §A.3.3 · p. 15, §A.3.4 · p. 15
A.3.5 LONG-SEQUENCE RECALL: VANILLA MAMBA VS. MEMMAMBA
- For k=100, the analysis estimates Recall_Mamba < 0.01, versus Recall_MemMamba ≥ 0.9 under the stated α, Δ, and θ assumptions.The MemMamba example uses α=0.8, Δ=0.1γ, and θ=0.7γ.Source: §A.3.5 · p. 16
A.4.1 LINEAR TIME COMPLEXITY
- With constant layer count and width, state summarization and cross-layer attention yield O(n) time; the idealized low-rank approximation has zero error and costs O(nkd).The passage contrasts O(nkd) with Transformer’s O(n^2d).Source: §A.4.1 · p. 16
A.4.2 LINEAR SPACE COMPLEXITY
- MemMamba’s state pool keeps active memory linear, O(nC), with small constants, and the analysis gives total space O(nd).The pool size is treated as constant.Source: §A.4.2 · p. 16
A.5 SENSITIVITY ANALYSIS; A.6 IMPLEMENTATION DETAILS
- Across broad ranges, window and pool sizes have negligible performance effects, max pooling performs best among tested pooling choices, and residual or weighted fusion has lower PPL at most long context lengths.The reported model uses 24 layers, a state-pool size of 50, and 8k-token training sequences.Source: §A.5 · p. 17, §A.6 · p. 18