MemMamba: Heroverweging van geheugenpatronen in state space-modellen
MemMamba: Rethinking Memory Patterns in State Space Model
September 28, 2025
Auteurs: Youjin Wang, Yangjingyi Chen, Jiahao Yan, Jiaxuan Lu, Xiao Sun
cs.AI
Samenvatting
Met de explosieve groei van data is modellering van lange sequenties steeds belangrijker geworden in taken zoals natuurlijke taalverwerking en bio-informatica. Bestaande methoden kampen echter met inherente afwegingen tussen efficiëntie en geheugengebruik. Recursieve neurale netwerken lijden onder het verdwijnen en exploderen van gradiënten, wat het opschalen bemoeilijkt. Transformers kunnen globale afhankelijkheden modelleren, maar worden beperkt door kwadratische complexiteit. Recentelijk hebben selectieve toestandsruimtemodellen zoals Mamba een hoge efficiëntie getoond met O(n)-tijd en O(1) recursieve inferentie, maar hun langeafstandsgeheugen vervalt exponentieel. In dit werk voeren we wiskundige afleidingen en informatie-theoretische analyses uit om systematisch het geheugenvervalmechanisme van Mamba te ontrafelen, waarbij we een fundamentele vraag beantwoorden: wat is de aard van Mamba's langeafstandsgeheugen en hoe behoudt het informatie? Om sleutelinformatieverlies te kwantificeren, introduceren we horizontaal-verticale geheugengetrouwheidsmetrieken die degradatie zowel binnen als tussen lagen vastleggen. Geïnspireerd door hoe mensen essentiële informatie destilleren en behouden bij het lezen van lange documenten, stellen we MemMamba voor, een nieuw architecturaal framework dat een toestandssamenvattingsmechanisme integreert samen met aandacht tussen lagen en tokens, wat langetermijnvergeten vermindert terwijl lineaire complexiteit behouden blijft. MemMamba behaalt significante verbeteringen ten opzichte van bestaande Mamba-varianten en Transformers op lange-sequentiebenchmarks zoals PG19 en Passkey Retrieval, terwijl het een snelheidswinst van 48% biedt in inferentie-efficiëntie. Zowel theoretische analyse als empirische resultaten tonen aan dat MemMamba een doorbraak bereikt in de complexiteit-geheugenafweging, en biedt een nieuw paradigma voor ultra-lange sequentiemodellering.
English
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.