让Mamba模型更可控、可观测、稳定,提升效率与性能。
Sparse Mamba: Introducing Controllability, Observability, And Stability To Structural State Space Models
- 引入可控性与可观测性,优化状态空间矩阵计算
- 参数减少,困惑度降5%,训练时间快3%
- 稀疏稳定A矩阵,适合构建下一代Mamba模型
近期研究中的结构化状态空间模型(SSMs),如Mamba和Mamba2,已在中小规模下超越Transformer,在计算效率上实现突破。本文提出Sparse-Mamba(S-Mamba),首次将可控性与可观测性引入原始Mamba架构,用于自然语言处理任务。同时,强化Mamba2中$nxn$ $A$矩阵的稳定性。现有Mamba模型在每步计算$A$、$B$、$C$、$D$矩阵时缺乏可控性,导致复杂度与开销上升;Mamba2的$A$矩阵亦不总稳定。本工作展示参数量低于首个Mamba与Mamba2,困惑度降低5%,训练时间缩短3%。进一步对Mamba2的$A$矩阵施加稳定性约束,有效降低损失与困惑度。所提稀疏$A$矩阵仅含$n$个自由参数,兼具可控、可观测与稳定特性,为构建Mamba3奠定关键基础。
原文摘要 · Abstract (English)
Structured state space models' (SSMs) development in recent studies, such as Mamba and Mamba2, outperformed and solved the computational inefficiency of transformers and large language models at small to medium scale. In this work, we introduce the concept of controllability and observability to the original Mamba SSM's architecture in our Sparse-Mamba (S-Mamba) for natural language processing (NLP) applications. Moreover, we reinforce stability on the $nxn$ $A$ matrix on Mmaba2. The Mamba SSMs architecture drops the need for attention layers or multilayer perception blocks in transformers. However, current Mamba models lack reinforcement of controllability in state-space equations for computing the $A$, $B$, $C$, and $D$ matrices at each time step, leading to increased complexity and computational costs. Furthermore, the $A$ matrix in Mamba2 is not always stable. We demonstrate a reduction of parameters compared to the first published Mamba and Mamba2. We showcase an improvement in perplexity by 5\% and a decrease in training time by 3\% after reinforcing controllability and observability on the original Mamba architecture in our proposed S-Mamba. We further enforce stability on the $A$ matrix in Mamba2 to improve the loss and perplexity of the model. The controllable and stable $n \times n$ state matrix $A$ is sparse, and it has only $n$ free parameters. Our novel approach will ensure controllable/observable and stable SSMs, which will be the gate key for Mamba3.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。