arXiv:2504.16053cs.CLcs.AI2025-04ICLR被引 38

不训练提升Mamba长文本理解能力,解决其记忆衰减问题。

LongMamba: Enhancing Mamba's Long Context Capabilities via Training-Free Receptive Field Enlargement

论文配图:LongMamba: Enhancing Mamba's Long Context Capabilities via Training-Free Receptive Field Enlargement
图 1 · 摘自论文原文
  • 通过筛选关键令牌,防止无关信息堆积导致全局通道记忆衰减。
  • 在合成与真实场景中显著提升Mamba长上下文表现,扩展可用上下文长度。
  • 无需重新训练,适配现有Mamba模型,适合追求高效长文本处理的开发者。

状态空间模型(SSMs)作为语言建模的高效替代方案,具备线性计算复杂度和恒定内存消耗,但相较于Transformer,在长上下文理解任务中仍表现欠佳。本文提出LongMamba,一种无需训练的改进方法,显著增强Mamba模型的长上下文能力。研究发现Mamba的隐藏通道可分为局部与全局通道,其中全局通道主导长上下文理解,但在输入长度远超训练序列长度时,其接收域无法自适应扩展,导致性能下降。LongMamba通过识别全局通道中的关键令牌,并仅累积这些关键信息,有效缓解隐藏状态的记忆衰减。在多种合成与真实长上下文场景中进行广泛评测,LongMamba显著提升Mamba性能,拓展其适用范围。代码已开源:https://github.com/GATECH-EIC/LongMamba。

原文摘要 · Abstract (English)

State space models (SSMs) have emerged as an efficient alternative to Transformer models for language modeling, offering linear computational complexity and constant memory usage as context length increases. However, despite their efficiency in handling long contexts, recent studies have shown that SSMs, such as Mamba models, generally underperform compared to Transformers in long-context understanding tasks. To address this significant shortfall and achieve both efficient and accurate long-context understanding, we propose LongMamba, a training-free technique that significantly enhances the long-context capabilities of Mamba models. LongMamba builds on our discovery that the hidden channels in Mamba can be categorized into local and global channels based on their receptive field lengths, with global channels primarily responsible for long-context capability. These global channels can become the key bottleneck as the input context lengthens. Specifically, when input lengths largely exceed the training sequence length, global channels exhibit limitations in adaptively extend their receptive fields, leading to Mamba's poor long-context performance. The key idea of LongMamba is to mitigate the hidden state memory decay in these global channels by preventing the accumulation of unimportant tokens in their memory. This is achieved by first identifying critical tokens in the global channels and then applying token filtering to accumulate only those critical tokens. Through extensive benchmarking across synthetic and real-world long-context scenarios, LongMamba sets a new standard for Mamba's long-context performance, significantly extending its operational range without requiring additional training. Our code is available at https://github.com/GATECH-EIC/LongMamba.

长文本Mamba状态空间无训练

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