Youjin Wang Yangjingyi Chen Jiahao Yan Jiaxuan Lu Xiao Sun

摘要

随着数据的爆炸式增长，长序列建模在自然语言处理、生物信息学等任务中变得日益重要。然而，现有方法在效率与内存消耗之间面临固有的权衡：循环神经网络（RNN）易受梯度消失与爆炸问题困扰，难以实现有效扩展；Transformer 虽能捕捉全局依赖关系，却受限于二次方复杂度。近期，选择性状态空间模型（如 Mamba）展现出 O(n) 时间复杂度与 O(1) 的递归推理能力，具有高效性，但其长程记忆能力呈指数衰减。在本工作中，我们通过数学推导与信息论分析，系统性地揭示了 Mamba 模型中记忆衰减的机制，回答了一个根本性问题：Mamba 的长程记忆本质是什么？它是如何实现信息保留的？为量化关键信息的丢失程度，我们进一步提出“横向-纵向记忆保真度”度量指标，用以捕捉层内与层间的信息退化情况。受人类阅读长文档时提炼并保留关键信息机制的启发，我们提出 MemMamba——一种新型架构框架，融合状态摘要机制与跨层、跨标记注意力机制，在保持线性复杂度的同时显著缓解了长程遗忘问题。在 PG19 和 Passkey Retrieval 等长序列基准测试中，MemMamba 相较于现有 Mamba 变体与 Transformer 均取得了显著性能提升，同时推理效率提升达 48%。理论分析与实证结果共同表明，MemMamba 在复杂度与记忆能力的权衡上实现了突破性进展，为超长序列建模提供了一种全新的范式。