⏶68
MemMamba:重新思考状态空间模型中的记忆模式
发表
由
Jiaxuan Lu 提交
作者: Youjin Wang, Yangjingyi Chen, Jiahao Yan, Jiaxuan Lu, Xiao Sun
摘要
AI 生成总结
MemMamba,一种整合状态摘要和交叉注意力的创新架构,与Mamba和Transformers相比,在序列建模方面提高了远距离记忆和效率。随着数据量的爆炸式增长,长序列建模在自然语言处理和生物信息学等任务中变得越来越重要。然而,现有方法在效率和内存之间面临固有的权衡。循环神经网络遭受梯度消失和爆炸的困扰,难以扩展。Transformer 可以模拟全局依赖关系,但受限于二次复杂度。最近,Mamba 等选择性状态空间模型已展示出高效率,具有 O(n) 的时间和 O(1) 的递归推理,但其长程记忆呈指数衰减。在这项工作中,我们进行了数学推导和信息论分析,系统地揭示了 Mamba 的记忆衰减机制,回答了一个基本问题:Mamba 的长程记忆的本质是什么?它如何保留信息?为了量化关键信息损失,我们进一步引入了水平-垂直记忆保真度指标,该指标可以捕获层内和层间的退化。受人类在阅读长文档时如何提炼和保留重要信息的启发,我们提出了 MemMamba,一种将状态摘要机制与跨层和跨标记注意力相结合的新型架构框架,它减轻了长程遗忘,同时保持了线性复杂度。MemMamba 在 PG19 和 Passkey Retrieval 等长序列基准测试上取得了比现有的 Mamba 变体和 Transformers 显著的改进,同时将推理效率提高了 48%。理论分析和经验结果都表明,MemMamba 在复杂性-内存权衡方面取得了突破,为超长序列建模提供了一种新范式。
评论

> 基线评估方法
表 1 报告了 Megalodon 模型的结果,该模型 未公开发布检查点。作者能否确认这些基线是否从头开始训练的?如果可以,能否提供训练细节(例如,是否使用了未修改的 XuezheMax/megalodon
,超参数和收敛曲线如何)?
报告的 Megalodon 困惑度 (PPL) 分数约为 64-66,在基础情况下似乎相当……糟糕,并且与其关于强大长上下文性能的公开声明不符。除了高数字之外,分数还随着序列长度略有……提高?
MemMamba:重新思考状态空间模型中的记忆模式