MemMamba : Repenser les motifs de mémoire dans les modèles d'espace d'état

MemMamba: Rethinking Memory Patterns in State Space Model

September 28, 2025

papers.authors: Youjin Wang, Yangjingyi Chen, Jiahao Yan, Jiaxuan Lu, Xiao Sun

cs.AI

papers.abstract

Avec l'explosion des données, la modélisation de séquences longues est devenue de plus en plus cruciale dans des tâches telles que le traitement du langage naturel et la bioinformatique. Cependant, les méthodes existantes sont confrontées à des compromis inhérents entre efficacité et mémoire. Les réseaux de neurones récurrents souffrent de problèmes de disparition et d'explosion des gradients, ce qui les rend difficiles à mettre à l'échelle. Les Transformers peuvent modéliser des dépendances globales, mais sont limités par une complexité quadratique. Récemment, des modèles à espace d'états sélectifs tels que Mamba ont démontré une grande efficacité avec un temps O(n) et une inférence récurrente O(1), mais leur mémoire à long terme décroît de manière exponentielle. Dans ce travail, nous menons des dérivations mathématiques et une analyse informationnelle pour systématiquement révéler le mécanisme de décroissance de la mémoire de Mamba, répondant à une question fondamentale : quelle est la nature de la mémoire à long terme de Mamba et comment conserve-t-elle l'information ? Pour quantifier la perte d'information clé, nous introduisons des métriques de fidélité mémoire horizontale-verticale qui capturent la dégradation à la fois au sein et entre les couches. Inspirés par la manière dont les humains distillent et retiennent les informations saillantes lors de la lecture de documents longs, nous proposons MemMamba, un nouveau cadre architectural qui intègre un mécanisme de synthèse d'état ainsi qu'une attention inter-couches et inter-tokens, atténuant ainsi l'oubli à long terme tout en préservant une complexité linéaire. MemMamba obtient des améliorations significatives par rapport aux variantes existantes de Mamba et aux Transformers sur des benchmarks de séquences longues tels que PG19 et Passkey Retrieval, tout en offrant une accélération de 48 % en efficacité d'inférence. L'analyse théorique et les résultats empiriques démontrent que MemMamba réalise une percée dans le compromis complexité-mémoire, offrant un nouveau paradigme pour la modélisation de séquences ultra-longues.

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.