MemMamba: Replanteando los patrones de memoria en los modelos de espacio 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

Resumen

Con el crecimiento explosivo de los datos, el modelado de secuencias largas se ha vuelto cada vez más importante en tareas como el procesamiento del lenguaje natural y la bioinformática. Sin embargo, los métodos existentes enfrentan compensaciones inherentes entre eficiencia y memoria. Las redes neuronales recurrentes sufren de desvanecimiento y explosión del gradiente, lo que dificulta su escalabilidad. Los Transformers pueden modelar dependencias globales pero están limitados por su complejidad cuadrática. Recientemente, los modelos de espacio de estados selectivos como Mamba han demostrado alta eficiencia con complejidad temporal O(n) e inferencia recurrente O(1), aunque su memoria de largo alcance decae exponencialmente. En este trabajo, realizamos derivaciones matemáticas y análisis teórico de la información para descubrir sistemáticamente el mecanismo de decaimiento de memoria de Mamba, respondiendo a una pregunta fundamental: ¿cuál es la naturaleza de la memoria de largo alcance de Mamba y cómo retiene la información? Para cuantificar la pérdida de información clave, introducimos métricas de fidelidad de memoria horizontal-vertical que capturan la degradación tanto dentro como entre capas. Inspirados por cómo los humanos destilan y retienen información relevante al leer documentos extensos, proponemos MemMamba, un marco arquitectónico novedoso que integra un mecanismo de resumen de estados junto con atención entre capas y entre tokens, lo que mitiga el olvido de largo alcance mientras preserva la complejidad lineal. MemMamba logra mejoras significativas sobre las variantes existentes de Mamba y Transformers en benchmarks de secuencias largas como PG19 y Passkey Retrieval, además de ofrecer un aumento del 48% en la eficiencia de inferencia. Tanto el análisis teórico como los resultados empíricos demuestran que MemMamba alcanza un avance en la compensación complejidad-memoria, ofreciendo un nuevo paradigma para el modelado de secuencias ultralargas.

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.