改进Mamba模型的注意力分配，提升长文本理解与检索能力
Differential Mamba
- 为Mamba设计新型差分机制，解决无关上下文过度关注问题
- 在语言建模任务上实现比原版Mamba更好的检索性能
- 适合关注高效大模型与长序列建模的研究者与工程师
序列模型如Transformer和RNN常对无关上下文过度分配注意力，导致中间表示噪声大，削弱了大语言模型的幻觉抑制、长距离依赖建模与检索能力。已有研究表明，差分设计可缓解Transformer中的该问题。本文探究此类方法能否应用于Mamba——一种基于选择性状态空间层的新架构，在效率上媲美Transformer。我们发现，直接套用差分设计不足以改善Mamba，需进行精细结构改造。为此，提出一种新型差分机制，经语言建模基准验证，显著提升检索能力并优于原始Mamba。通过广泛的消融实验与实证分析，证明该方法有效缓解了Mamba模型中的注意力过度分配问题。代码已开源：https://github.com/NadavSc/Diff-Mamba
原文摘要 · Abstract (English)
Sequence models like Transformers and RNNs often overallocate attention to irrelevant context, leading to noisy intermediate representations. This degrades LLM capabilities by promoting hallucinations, weakening long-range and retrieval abilities, and reducing robustness. Recent work has shown that differential design can mitigate this issue in Transformers, improving their effectiveness across various applications. In this paper, we explore whether these techniques, originally developed for Transformers, can be applied to Mamba, a recent architecture based on selective state-space layers that achieves Transformer-level performance with greater efficiency. We show that a naive adaptation of differential design to Mamba is insufficient and requires careful architectural modifications. To address this, we introduce a novel differential mechanism for Mamba, empirically validated on language modeling benchmarks, demonstrating improved retrieval capabilities and superior performance over vanilla Mamba. Finally, we conduct extensive ablation studies and empirical analyses to justify our design choices and provide evidence that our approach effectively mitigates the overallocation problem in Mamba-based models. Our code is publicly available: https://github.com/NadavSc/Diff-Mamba
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。