用专家路由提升状态空间模型扩展性,1.3B激活参数达2.3倍密集模型效果
Routing Mamba: Scaling State Space Models with Mixture-of-Experts Projection
- 通过共享路由决策的线性专家投影实现稀疏扩展
- 1.3B激活参数下达到2.3倍参数密集模型性能
- 适合需要高效长序列建模的系统级部署
线性状态空间模型(SSMs)在高效序列建模中表现优异,具有恒定的推理时间和内存复杂度。近期进展如Mamba通过输入相关门控和硬件感知实现进一步优化,使其成为长序列建模对Transformer的有力替代。然而,如何高效扩展SSMs的表达能力,特别是结合混合专家(MoE)仍具挑战,直接集成常导致性能下降。本文提出路由马尔可夫(Routing Mamba, RoM),采用稀疏线性投影专家混合来扩展SSM参数。通过在投影层与马尔可夫内部轻量子模块间共享路由决策,RoM利用线性投影专家间的协同效应,实现有效且高效的稀疏扩展。在1.3B激活参数(总参数10B)、16K训练序列长度下,RoM的语言建模范式性能相当于需超过2.3倍激活参数的密集型马尔可夫模型,并在不同上下文长度下保持一致困惑度。实验表明,该方法能有效扩展混合语言模型,在相近性能下相较密集型马尔可夫扩展节省23% FLOPS。
原文摘要 · Abstract (English)
Linear State Space Models (SSMs) offer remarkable performance gains in efficient sequence modeling, with constant inference-time computation and memory complexity. Recent advances, such as Mamba, further enhance SSMs with input-dependent gating and hardware-aware implementations, positioning them as strong alternatives to Transformers for long sequence modeling. However, efficiently scaling the expressive power of SSMs, particularly with Mixture of Experts (MoE), remains challenging, as naive integration attempts often falter or degrade performance. In this work, we introduce Routing Mamba (RoM), a novel approach that scales SSM parameters using sparse mixtures of linear projection experts. By sharing routing decisions between projection layers and lightweight sub-modules within Mamba across experts, RoM leverages synergies among linear projection experts for effective and efficient sparse scaling of Mamba layers. At a scale of 1.3B active parameters (10B total) and 16K training sequence length, RoM achieves language modeling performance equivalent to a dense Mamba model requiring over 2.3x more active parameters, and demonstrates consistent perplexity across context lengths. Experimental results further show RoM effectively scales hybrid language models, yielding a 23% FLOPS saving compared to dense Mamba scaling for similar performance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。