用分层记忆增强Mamba,让长序列建模更准更快。
Mamba with Hierarchical Memory: Solving Representation Bottleneck in Long Sequence Modeling

- 模仿人类记忆结构,引入工作记忆与长期记忆分层机制。
- 在长文本任务中,检索成功率提升34.3%~37.1%,推理准确率提升1.6%~14.2%。
- 仅增加2%参数,适合需要高效长序列处理的场景。
递归线性注意力模型(如Mamba)以线性时间复杂度提供高效序列建模,但其固定容量的递归状态限制了长序列建模能力。受人类记忆层级结构启发,本文提出分层记忆Mamba(HMM),在预训练Mamba主干基础上,集成轻量级工作记忆,从主干隐藏状态嵌入的快速感觉记忆中提取慢速段落级语义(PLS),并将该信息压缩为任务相关的持久长期记忆以支持检索。这种分层语义处理克服了传统RLA的表征瓶颈,通过参数化学习实现跨任务泛化能力,这是其他增强型Mamba变体所不具备的。在Passkey Retrieval和LongBench-E任务上的评估表明,相比强基线模型,HMM在检索成功率上提升34.3%~37.1%,推理准确率提升1.6%~14.2%,同时仅增加2%额外参数,训练开销极小。
原文摘要 · Abstract (English)
Recurrent linear attention models (RLAs) such as Mamba offer efficient linear-time sequence modeling as an alternative to Transformers, yet their fixed-capacity recurrent states limit long-sequence modeling. Drawing inspiration from hierarchical human memory, we propose Hierarchical Memory Mamba (HMM) to address this limitation. Building upon a pre-trained Mamba backbone, HMM integrates a lightweight working memory that extracts slow paragraph-level semantics (PLS) from the fast sensory memory embedded in the backbone's hidden states. The PLS is subsequently compressed into persistent long-term memory for task-relevant retrieval. The hierarchical processing of semantic information overcomes the representation bottleneck of RLAs and endows HMM cross-task generalization through parametric learning, which is not observed in other long-context enhanced Mamba variants. Evaluations on Passkey Retrieval and LongBench-E tasks demonstrate that HMM improves retrieval success by 34.3--37.1% and reasoning accuracy by 1.6--14.2% over strong Mamba-based models, while adding only 2% extra parameters and with minimal training overhead.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。