MemMamba: Переосмысление паттернов памяти в моделях пространства состояний
MemMamba: Rethinking Memory Patterns in State Space Model
September 28, 2025
Авторы: Youjin Wang, Yangjingyi Chen, Jiahao Yan, Jiaxuan Lu, Xiao Sun
cs.AI
Аннотация
С экспоненциальным ростом данных моделирование длинных последовательностей становится все более важным в задачах, таких как обработка естественного языка и биоинформатика. Однако существующие методы сталкиваются с неизбежным компромиссом между эффективностью и использованием памяти. Рекуррентные нейронные сети страдают от проблем исчезающих и взрывающихся градиентов, что затрудняет их масштабирование. Трансформеры способны моделировать глобальные зависимости, но ограничены квадратичной сложностью. Недавно селективные модели пространства состояний, такие как Mamba, продемонстрировали высокую эффективность с линейной временной сложностью O(n) и постоянной сложностью O(1) при рекуррентном выводе, однако их долгосрочная память экспоненциально затухает. В данной работе мы проводим математические выкладки и информационно-теоретический анализ, чтобы систематически раскрыть механизм затухания памяти в Mamba, отвечая на фундаментальный вопрос: какова природа долгосрочной памяти Mamba и как она сохраняет информацию? Для количественной оценки ключевой потери информации мы вводим метрики горизонтально-вертикальной верности памяти, которые фиксируют деградацию как внутри слоев, так и между ними. Вдохновленные тем, как люди выделяют и сохраняют важную информацию при чтении длинных документов, мы предлагаем MemMamba — новую архитектурную структуру, которая интегрирует механизм суммирования состояний вместе с вниманием между слоями и токенами, что смягчает проблему забывания в длинных последовательностях, сохраняя при этом линейную сложность. MemMamba демонстрирует значительные улучшения по сравнению с существующими вариантами Mamba и Трансформерами на бенчмарках длинных последовательностей, таких как PG19 и Passkey Retrieval, обеспечивая при этом ускорение вывода на 48%. Как теоретический анализ, так и эмпирические результаты показывают, что MemMamba достигает прорыва в компромиссе между сложностью и памятью, предлагая новую парадигму для моделирования сверхдлинных последовательностей.
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.