LongMamba：通过训练自由地扩大Mamba的感受野以增强其长上下文能力
计算与语言
2025-04-23 v1 人工智能
摘要
状态空间模型（SSMs）已成为Transformer模型的高效替代方案，能够在上下文长度增加时保持线性计算复杂度和恒定内存使用。然而，尽管SSMs在处理长上下文方面效率高，但近期研究表明，诸如Mamba模型之类的SSMs在长上下文理解任务中通常表现不如Transformer。为解决这一显著不足以实现高效且准确的长上下文理解，我们提出了LongMamba，这是一种无需训练即可显著增强Mamba模型长上下文能力的技术。LongMamba建立在我们发现的隐藏通道可根据其感受野长度分为局部和全局通道的观察基础上，全局通道主要负责长上下文能力。当输入上下文长度大大超过训练序列长度时，全局通道成为关键瓶颈，表现出在适应性地扩展其感受野方面存在限制，导致Mamba的长上下文性能差。LongMamba的关键思想是通过防止其记忆中累积无关标记来缓解全局通道中的隐藏状态记忆衰减。这通过首先识别全局通道中的关键标记，然后对其应用标记过滤以仅累积这些关键标记来实现。通过在合成和真实世界的长上下文场景中进行大规模基准测试，LongMamba为Mamba的长上下文性能树立了新标准，在无需额外训练的情况下显著扩展了其操作范围。我们的代码可在 https://github.com/GATECH-EIC/LongMamba 上获取。
引用
@article{arxiv.2504.16053,
title = {LongMamba: Enhancing Mamba's Long Context Capabilities via Training-Free Receptive Field Enlargement},
author = {Zhifan Ye and Kejing Xia and Yonggan Fu and Xin Dong and Jihoon Hong and Xiangchi Yuan and Shizhe Diao and Jan Kautz and Pavlo Molchanov and Yingyan Celine Lin},
journal= {arXiv preprint arXiv:2504.16053},
year = {2025}
}
备注
Accepted by ICLR 2025