Routing Mamba: 基于稀疏专家投影扩展状态空间模型
机器学习
2025-06-24 v1 人工智能
摘要
线性状态空间模型（SSMs）在高效序列建模方面表现卓越，推理时间计算和内存复杂度为常数。近期进展如Mamba进一步通过输入依赖门控和硬件感知实现提升了SSMs，使其成为长序列建模中强大的Transformer替代方案。然而，对SSMs表达力度的高效扩展尤其是采用混合专家（MoE）方式，仍面临挑战，因直接集成往往适得其反或导致性能下降。本文引入Routing Mamba（RoM）一种新方法，通过稀疏线性投影专家的混合来扩展SSM参数。RoM在投影层与Mamba中轻量子模块之间共享路由决策，利用线性投影专家之间的协同作用，实现对Mamba层的有效且高效的稀疏扩展。在13亿活跃参数（总计100亿）和16K训练序列长度规模下，RoM实现的语言建模性能相当于需活跃参数超过2.3倍的稠密Mamba模型，并展现出跨上下文长度的稳定困惑度。实验结果进一步表明，RoM能够有效扩展混合语言模型，相较于稠密Mamba扩展，实现23%的FLOPS节省。
引用
@article{arxiv.2506.18145,
title = {Routing Mamba: Scaling State Space Models with Mixture-of-Experts Projection},
author = {Zheng Zhan and Liliang Ren and Shuohang Wang and Liyuan Liu and Yang Liu and Yeyun Gong and Yanzhi Wang and Yelong Shen},
journal= {arXiv preprint arXiv:2506.18145},
year = {2025}
}