Mamba模型依赖局部模式捷径，新模块让其处理分散信息能力提升
Revealing and Mitigating the Local Pattern Shortcuts of Mamba
- 在Mamba中加入全局选择模块，突破局部记忆限制
- 仅增400万参数，130万参数模型在分布式任务上性能从0提至80.54
- 适合需长距离信息建模的场景，如复杂文本推理
大型语言模型因注意力机制取得显著进展，但其二次复杂度和线性内存需求限制了长上下文任务表现。近期提出的Mamba基于状态空间模型（SSMs），具备线性复杂度和恒定内存。尽管报道显示其性能可媲美或超越注意力模型，我们的分析揭示性能差距：Mamba擅长包含局部关键信息的任务，但在需要处理分散关键信息的任务上表现不佳。控制实验表明，这种不一致性源于Mamba对局部模式捷径的依赖，使其能在有限内存中记忆局部关键信息，却难以保留更分散的信息。为此，我们向Mamba引入全局选择模块以解决该问题。在现有及新设计的合成任务、真实世界任务上的实验验证了方法有效性。值得注意的是，仅增加400万额外参数，我们的方法使1300万参数的Mamba模型在分布式信息任务上性能从0提升至80.54分。
原文摘要 · Abstract (English)
Large language models (LLMs) have advanced significantly due to the attention mechanism, but their quadratic complexity and linear memory demands limit their performance on long-context tasks. Recently, researchers introduced Mamba, an advanced model built upon State Space Models(SSMs) that offers linear complexity and constant memory. Although Mamba is reported to match or surpass the performance of attention-based models, our analysis reveals a performance gap: Mamba excels in tasks that involve localized key information but faces challenges with tasks that require handling distributed key information. Our controlled experiments suggest that this inconsistency arises from Mamba's reliance on local pattern shortcuts, which enable the model to remember local key information within its limited memory but hinder its ability to retain more dispersed information. Therefore, we introduce a global selection module into the Mamba model to address this issue. Experiments on both existing and proposed synthetic tasks, as well as real-world tasks, demonstrate the effectiveness of our method. Notably, with the introduction of only 4M extra parameters, our approach enables the Mamba model(130M) to achieve a significant improvement on tasks with distributed information, increasing its performance from 0 to 80.54 points.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。