arXiv:2510.03279cs.LGcs.AI2025-10被引 2

解决Mamba模型长序列记忆衰减问题,提升信息保留能力。

MemMamba: Rethinking Memory Patterns in State Space Model

  • 引入状态摘要与跨层注意力机制,缓解长期遗忘
  • 在PG19和Passkey任务上超越现有Mamba与Transformer
  • 保持线性复杂度,推理速度提升48%

随着数据爆炸式增长,长序列建模在自然语言处理和生物信息学中愈发重要。现有方法在效率与记忆之间存在固有权衡:循环神经网络易出现梯度消失/爆炸,难以扩展;Transformer虽能建模全局依赖,但计算复杂度为二次方。近期的选通状态空间模型如Mamba实现了O(n)时间复杂度和O(1)递归推理,但其长程记忆呈指数衰减。本文通过数学推导与信息论分析,系统揭示了Mamba的记忆衰减机制,回答了其长程记忆的本质及信息保留方式。为量化关键信息损失,提出水平-垂直记忆保真度指标,捕捉层内与跨层退化。受人类阅读长文档时提取关键信息的启发,提出MemMamba架构,融合状态摘要机制与跨层、跨标记注意力,缓解长期遗忘,同时保持线性复杂度。在PG19和Passkey Retrieval等长序列基准上显著优于现有Mamba变体与Transformer,并实现48%的推理加速。理论与实证均表明,MemMamba在复杂度-记忆权衡上取得突破,为超长序列建模提供新范式。

原文摘要 · Abstract (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.

状态空间模型长序列建模记忆机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。