路由Mamba:利用混合专家投影扩展状态空间模型
Routing Mamba: Scaling State Space Models with Mixture-of-Experts Projection
AI总结:
本文提出路由Mamba(RoM),通过在线性投影专家间共享路由决策,以稀疏混合专家方式扩展状态空间模型,在1.3B激活参数下达到密集Mamba超2.3倍参数的性能,并节省23%计算量。
摘要:
线性状态空间模型(SSMs)在高效序列建模中提供了显著的性能提升,其推理时的计算和内存复杂度为常数。最近的进展,如Mamba,通过输入相关的门控和硬件感知的实现进一步增强了SSMs,使其成为Transformer在长序列建模中的有力替代方案。然而,高效扩展SSMs的表达能力,特别是使用混合专家(MoE),仍然具有挑战性,因为简单的集成尝试常常失败或降低性能。在这项工作中,我们引入了路由Mamba(RoM),一种使用稀疏的线性投影专家混合来扩展SSM参数的新方法。通过在Mamba中的投影层和轻量子模块之间共享路由决策,RoM利用线性投影专家之间的协同作用,实现Mamba层有效且高效的稀疏扩展。在1.3B激活参数(总计10B)和16K训练序列长度的规模下,RoM实现了与需要超过2.3倍激活参数的密集Mamba模型相当的语言建模性能,并在不同上下文长度下表现出一致的困惑度。实验结果进一步表明,RoM有效扩展了混合语言模型,与密集Mamba扩展相比,在相似性能下节省了23%的FLOPS。
英文摘要:
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.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。