MemMamba: Repensando Padrões de Memória em Modelos de Espaço de Estados

MemMamba: Rethinking Memory Patterns in State Space Model

September 28, 2025

Autores: Youjin Wang, Yangjingyi Chen, Jiahao Yan, Jiaxuan Lu, Xiao Sun

cs.AI

Resumo

Com o crescimento explosivo de dados, a modelagem de sequências longas tornou-se cada vez mais importante em tarefas como processamento de linguagem natural e bioinformática. No entanto, os métodos existentes enfrentam trade-offs inerentes entre eficiência e memória. Redes neurais recorrentes sofrem com o problema de desaparecimento e explosão de gradientes, dificultando sua escalabilidade. Transformers conseguem modelar dependências globais, mas são limitados pela complexidade quadrática. Recentemente, modelos de espaço de estados seletivos, como o Mamba, demonstraram alta eficiência com tempo O(n) e inferência recorrente O(1), porém sua memória de longo alcance decai exponencialmente. Neste trabalho, realizamos derivações matemáticas e análises teóricas da informação para desvendar sistematicamente o mecanismo de decaimento de memória do Mamba, respondendo a uma questão fundamental: qual é a natureza da memória de longo alcance do Mamba e como ela retém informações? Para quantificar a perda de informações-chave, introduzimos métricas de fidelidade de memória horizontal-vertical que capturam a degradação tanto dentro quanto entre camadas. Inspirados pela forma como os seres humanos destilam e retêm informações relevantes ao ler documentos longos, propomos o MemMamba, uma nova arquitetura que integra um mecanismo de sumarização de estados com atenção entre camadas e entre tokens, o que atenua o esquecimento de longo alcance enquanto preserva a complexidade linear. O MemMamba alcança melhorias significativas em relação às variantes existentes do Mamba e aos Transformers em benchmarks de sequências longas, como PG19 e Passkey Retrieval, além de oferecer um ganho de 48% na eficiência de inferência. Tanto a análise teórica quanto os resultados empíricos demonstram que o MemMamba representa um avanço no trade-off entre complexidade e memória, oferecendo um novo paradigma para a modelagem de sequências ultra-longas.

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.